Golang中的简单工作者池

时间:2019-12-27 04:12:26

标签: go goroutine worker-pool

我必须完成一些与kafka并行的任务。所以我想在go中实现一个简单的工作组模式。

下面的代码是否足够体面,可以投入生产?


var workerCount = runtime.NumCPU()*7 + 1

func WorkerPattern() {
    taskWg := &sync.WaitGroup{}
    taskWg.Add(2)

    autoCancelChan := make(chan string, workerCount*3) // *3, just to make enough room. workers will be slower anyways
    go produceTaskOfType1ToChan(taskWg, autoCancelChan)
    go produceTaskOfType2ToChan(taskWg, autoCancelChan)

    // start workers to push autoCancel to kafka
    workerWg := &sync.WaitGroup{}
    go kafkaProducerWorkers(autoCancelChan, workerWg)

    // wait to close autoCancelChan channel till all the task is written
    taskWg.Wait()
    close(autoCancelChan)

    // wait till all workers finish their task
    workerWg.Wait()

    fmt.Println("Done!!!")
}

func produceTaskOfType1ToChan(wg *sync.WaitGroup, autoCancelChan chan string) {
    defer wg.Done()
    // can produce random number of task on autoCancelChan
    autoCancelChan <- "task of type of 1"
}

func produceTaskOfType2ToChan(wg *sync.WaitGroup, autoCancelChan chan string) {
    defer wg.Done()
    // can produce random number of task on autoCancelChan
    autoCancelChan <- "task of type of 2"
}

func kafkaProducerWorkers(autoCancelChan chan string, workerWg *sync.WaitGroup) {
    workerWg.Add(workerCount)
    for i := 0; i < workerCount; i++ {
        go produceToKafka(autoCancelChan, workerWg)
    }
}

func produceToKafka(autoCancelChan chan string, workerWg *sync.WaitGroup) {
    defer workerWg.Done()

    // for loop will terminate once autoCancelChan is closed
    for autoCancel := range autoCancelChan {
        KafkaClient.PublishToKafkaTopic(autoCancel)
    }
}

您可以对这段代码提出任何建议吗?

0 个答案:

没有答案