GO语言通用生产消费模型

pcmode.go

package pcmodel

import "sync"

type callback func(ch chan interface{})

type PCModel struct {
	wg *sync.WaitGroup
	ch chan interface{}
	Producer
	Consumer
}

type Producer struct {
	wg    *sync.WaitGroup
	ch    chan interface{}
	pfunc callback
}

type Consumer struct {
	wg    *sync.WaitGroup
	ch    chan interface{}
	cfunc callback
}

func NewPCModel(pfunc callback, cfunc callback) *PCModel {
	model := PCModel{}
	model.wg = &sync.WaitGroup{}
	model.ch = make(chan interface{})
	model.Producer = *NewProducer(model.wg, model.ch, pfunc)
	model.Consumer = *NewConsumer(model.wg, model.ch, cfunc)
	return &model
}

func NewProducer(wg *sync.WaitGroup, ch chan interface{}, f callback) *Producer {
	p := Producer{}
	p.wg = wg
	p.ch = ch
	p.pfunc = f
	return &p
}

func NewConsumer(wg *sync.WaitGroup, ch chan interface{}, f callback) *Consumer {
	c := Consumer{}
	c.wg = wg
	c.ch = ch
	c.cfunc = f
	return &c
}

func (p *Producer) Produce() {
	p.wg.Add(1)
	go func(ch chan interface{}) {
		defer p.wg.Done()
		p.pfunc(ch)
		close(ch)
	}(p.ch)
}

func (c *Consumer) Consume() {
	c.wg.Add(1)
	go func(ch chan interface{}) {
		defer c.wg.Done()
		c.cfunc(ch)
	}(c.ch)
}

func (pc *PCModel) Wait() {
	pc.wg.Wait()
}

使用方法如下：

package main

import (
	"fmt"
	"pcmodel"
)

func ImportData(ch chan interface{}) {
	for i := 0; i < 10; i++ {
		ch <- i
	}
}

func ExportData(ch chan interface{}) {
	for data := range ch {
		fmt.Println(data)
	}
}

func main() {
	pc := pcmodel.NewPCModel(ImportData, ExportData)
	pc.Produce()
	pc.Consume()
	pc.Wait()
}