初步增加协程池

Author: FishGoddess

pool.go

@@ -0,0 +1,74 @@
+// Copyright 2025 FishGoddess. All rights reserved.
+// Use of this source code is governed by a MIT style
+// license that can be found in the LICENSE file.
+
+package goes
+
+import (
+	"sync"
+
+	"github.com/FishGoddess/goes/pkg/spinlock"
+)
+
+type Pool struct {
+	workers []*worker
+	index   int
+	closed  bool
+
+	wg   sync.WaitGroup
+	lock sync.Locker
+}
+
+func NewPool(size int, workerLimit int) *Pool {
+	if size <= 0 {
+		panic("goes: pool size <= 0")
+	}
+
+	pool := &Pool{
+		workers: make([]*worker, 0, size),
+		index:   0,
+		closed:  false,
+		lock:    spinlock.New(),
+	}
+
+	for range size {
+		worker := newWorker(pool, workerLimit)
+		pool.workers = append(pool.workers, worker)
+	}
+
+	return pool
+}
+
+func (p *Pool) Go(f func()) {
+	p.lock.Lock()
+	defer p.lock.Unlock()
+
+	if p.closed {
+		return
+	}
+
+	worker := p.workers[p.index]
+	worker.Accept(f)
+
+	p.index++
+
+	if p.index >= len(p.workers) {
+		p.index = 0
+	}
+}
+
+func (p *Pool) Wait() {
+	p.wg.Wait()
+}
+
+func (p *Pool) Close() {
+	p.lock.Lock()
+	defer p.lock.Unlock()
+
+	for _, worker := range p.workers {
+		worker.Done()
+	}
+
+	p.wg.Wait()
+	p.closed = true
+}

pool_test.go

@@ -0,0 +1,50 @@
+// Copyright 2025 FishGoddess. All rights reserved.
+// Use of this source code is governed by a MIT style
+// license that can be found in the LICENSE file.
+
+package goes
+
+import (
+	"sync"
+	"testing"
+	"time"
+)
+
+// go test -v -cover -run=^TestPool$
+func TestPool(t *testing.T) {
+	size := 16
+	workerLimit := 1024
+	pool := NewPool(size, workerLimit)
+
+	var countMap = make(map[int64]int, 16)
+	var lock sync.Mutex
+
+	totalCount := 10 * size
+	for i := 0; i < totalCount; i++ {
+		pool.Go(func() {
+			now := time.Now().UnixMilli() / 10
+
+			lock.Lock()
+			countMap[now] = countMap[now] + 1
+			lock.Unlock()
+
+			time.Sleep(10 * time.Millisecond)
+		})
+	}
+
+	pool.Close()
+	pool.Wait()
+
+	gotTotalCount := 0
+	for now, count := range countMap {
+		gotTotalCount = gotTotalCount + count
+
+		if count != size {
+			t.Fatalf("now %d: count %d != size %d", now, count, size)
+		}
+	}
+
+	if gotTotalCount != totalCount {
+		t.Fatalf("gotTotalCount %d != totalCount %d", gotTotalCount, totalCount)
+	}
+}

worker.go

@@ -0,0 +1,67 @@
+// Copyright 2025 FishGoddess. All rights reserved.
+// Use of this source code is governed by a MIT style
+// license that can be found in the LICENSE file.
+
+package goes
+
+type worker struct {
+	pool  *Pool
+	tasks chan func()
+}
+
+func newWorker(pool *Pool, limit int) *worker {
+	if limit <= 0 {
+		panic("goes: worker limit <= 0")
+	}
+
+	w := &worker{
+		pool:  pool,
+		tasks: make(chan func(), limit),
+	}
+
+	w.work()
+	return w
+}
+
+func (w *worker) handle(task func()) {
+	if task == nil {
+		return
+	}
+
+	defer func() {
+		if r := recover(); r != nil {
+			panic(r) // TODO recover from panic
+		}
+	}()
+
+	task()
+}
+
+func (w *worker) work() {
+	w.pool.wg.Add(1)
+	go func() {
+		defer w.pool.wg.Done()
+
+		for task := range w.tasks {
+			if task == nil {
+				break
+			}
+
+			w.handle(task)
+		}
+
+		close(w.tasks)
+	}()
+}
+
+func (w *worker) Accept(task func()) {
+	if task == nil {
+		return
+	}
+
+	w.tasks <- task
+}
+
+func (w *worker) Done() {
+	w.tasks <- nil
+}

worker_test.go

@@ -0,0 +1,5 @@
+// Copyright 2025 FishGoddess. All rights reserved.
+// Use of this source code is governed by a MIT style
+// license that can be found in the LICENSE file.
+
+package goes