spark-server/internal/transport/tcp/taskpool.go

package tcp

import (
	"context"
	"sync"
	"sync/atomic"

	"spark-server/internal/pkg/errs"
)

// TaskPool runs inbound business (MySQL included). EventLoop must only Submit.
type TaskPool interface {
	Submit(ctx context.Context, task func()) error
	Close()
	Closed() bool
}

type pool struct {
	tasks   chan func()
	workers int
	wg      sync.WaitGroup
	once    sync.Once
	closed  atomic.Bool
	onPanic func(any)
}

func NewTaskPool(workers, queue int) TaskPool {
	return NewTaskPoolRecover(workers, queue, nil)
}

func NewTaskPoolRecover(workers, queue int, onPanic func(any)) TaskPool {
	if workers < 1 {
		workers = 1
	}
	if queue < 1 {
		queue = 1
	}
	p := &pool{tasks: make(chan func(), queue), workers: workers, onPanic: onPanic}
	for i := 0; i < workers; i++ {
		p.wg.Add(1)
		go p.loop()
	}
	return p
}

func (p *pool) loop() {
	defer p.wg.Done()
	for task := range p.tasks {
		p.run(task)
	}
}

func (p *pool) run(task func()) {
	if task == nil {
		return
	}
	defer func() {
		if rec := recover(); rec != nil && p.onPanic != nil {
			p.onPanic(rec)
		}
	}()
	task()
}

func (p *pool) Submit(ctx context.Context, task func()) error {
	if task == nil {
		return errs.ErrInvalidArgument
	}
	select {
	case <-ctx.Done():
		return ctx.Err()
	case p.tasks <- task:
		return nil
	default:
		return errs.ErrBusy
	}
}

func (p *pool) Close() {
	p.once.Do(func() {
		p.closed.Store(true)
		close(p.tasks)
		p.wg.Wait()
	})
}

func (p *pool) Closed() bool {
	return p.closed.Load()
}