golang workerpool

README.md

references

• https://www.freecodecamp.org/news/million-websockets-and-go-cc58418460bb/
• https://github.com/mosn/mosn/blob/master/pkg/sync/workerpool.go

workerpool.go

package pool

import (
	"log"
	"runtime"
	"runtime/debug"
)

type workerPool struct {
	work chan func()
	sem  chan struct{}
}

// New создает экземпляр worker pool с размером семафора равным кол-ва ядер умноженное на 256
func New() *workerPool {
	return NewWithSize(runtime.NumCPU() * 256)
}

func NewWithSize(size int) *workerPool {
	return &workerPool{
		work: make(chan func()),
		sem:  make(chan struct{}, size),
	}
}

func (p *workerPool) Schedule(task func()) {
	select {
	case p.work <- task:
	case p.sem <- struct{}{}:
		go p.spawnWorker(task)
	}
}

func (p *workerPool) spawnWorker(task func()) {
	defer func() {
		if r := recover(); r != nil {
			log.Printf("panic %v\n%s", r, string(debug.Stack()))
		}
		<-p.sem
	}()
	for {
		task()
		task = <-p.work
	}
}

workerpool_test.go

package pool

import (
	"testing"
	"time"
)

func TestSchedule(t *testing.T) {
	size := 5
	p := NewWithSize(size)
	for i := 0; i < 5; i++ {
		p.Schedule(func() {
			time.Sleep(100 * time.Millisecond)
		})
	}
	now := time.Now()
	p.Schedule(func() {})
	if time.Now().Before(now.Add(50 * time.Millisecond)) {
		t.Errorf("Test Schedule() error, should wait for 20 millisecond")
	}
	if len(p.sem) != size {
		t.Errorf("Test Schedule() error, should be 5")
	}
}