nnet/pkg/executor/executor.go

package executor

import (
	"context"
	"errors"
	"runtime"
	"sync"
)

var (
	ErrClosed    = errors.New("executor is closed")
	ErrQueueFull = errors.New("executor queue is full")
	ErrNilTask   = errors.New("executor task is nil")
)

type Task func()

type Executor interface {
	Submit(key string, task Task) error
	Shutdown(ctx context.Context) error
}

type Config struct {
	Workers   int
	QueueSize int
	OnPanic   func(value interface{})
}

type PerKey struct {
	mu        sync.Mutex
	cond      *sync.Cond
	queues    map[string]*taskQueue
	ready     []*taskQueue
	workers   int
	queueSize int
	onPanic   func(value interface{})
	pending   int
	closed    bool
	done      chan struct{}
	wg        sync.WaitGroup
}

type taskQueue struct {
	key     string
	tasks   []Task
	pending int
	running bool
	ready   bool
}

func NewPerKey(cfg Config) *PerKey {
	workers := cfg.Workers
	if workers <= 0 {
		workers = runtime.GOMAXPROCS(0)
	}
	if workers <= 0 {
		workers = 1
	}
	queueSize := cfg.QueueSize
	if queueSize <= 0 {
		queueSize = 1024
	}

	e := &PerKey{
		queues:    make(map[string]*taskQueue),
		workers:   workers,
		queueSize: queueSize,
		onPanic:   cfg.OnPanic,
		done:      make(chan struct{}),
	}
	e.cond = sync.NewCond(&e.mu)
	e.wg.Add(workers)
	for i := 0; i < workers; i++ {
		go e.worker()
	}
	go func() {
		e.wg.Wait()
		close(e.done)
	}()
	return e
}

func (e *PerKey) Submit(key string, task Task) error {
	if task == nil {
		return ErrNilTask
	}

	e.mu.Lock()
	defer e.mu.Unlock()

	if e.closed {
		return ErrClosed
	}

	queue := e.queues[key]
	if queue == nil {
		queue = &taskQueue{key: key}
		e.queues[key] = queue
	}
	if len(queue.tasks) >= e.queueSize {
		return ErrQueueFull
	}

	queue.tasks = append(queue.tasks, task)
	queue.pending++
	e.pending++
	if !queue.running && !queue.ready {
		queue.ready = true
		e.ready = append(e.ready, queue)
		e.cond.Signal()
	}
	return nil
}

func (e *PerKey) Shutdown(ctx context.Context) error {
	if ctx == nil {
		ctx = context.Background()
	}

	e.mu.Lock()
	if !e.closed {
		e.closed = true
		e.cond.Broadcast()
	}
	e.mu.Unlock()

	select {
	case <-e.done:
		return nil
	case <-ctx.Done():
		return ctx.Err()
	}
}

func (e *PerKey) worker() {
	defer e.wg.Done()

	for {
		task, queue, ok := e.nextTask()
		if !ok {
			return
		}

		e.run(task)

		e.mu.Lock()
		queue.running = false
		queue.pending--
		e.pending--
		if len(queue.tasks) > 0 {
			queue.ready = true
			e.ready = append(e.ready, queue)
			e.cond.Signal()
		} else {
			delete(e.queues, queue.key)
		}
		e.cond.Broadcast()
		e.mu.Unlock()
	}
}

func (e *PerKey) nextTask() (Task, *taskQueue, bool) {
	e.mu.Lock()
	defer e.mu.Unlock()

	for len(e.ready) == 0 {
		if e.closed && e.pending == 0 {
			return nil, nil, false
		}
		e.cond.Wait()
	}

	queue := e.ready[0]
	e.ready[0] = nil
	e.ready = e.ready[1:]
	queue.ready = false

	task := queue.tasks[0]
	queue.tasks[0] = nil
	queue.tasks = queue.tasks[1:]
	queue.running = true
	return task, queue, true
}

func (e *PerKey) run(task Task) {
	defer func() {
		if value := recover(); value != nil && e.onPanic != nil {
			e.onPanic(value)
		}
	}()
	task()
}