parent
e805b244d7
commit
604bd05e3b
@ -0,0 +1,58 @@
|
||||
package server
|
||||
|
||||
import (
|
||||
"context"
|
||||
"sync"
|
||||
|
||||
"git.noahlan.cn/noahlan/nnet/v2/internal/logger"
|
||||
"git.noahlan.cn/noahlan/nnet/v2/pkg/config"
|
||||
"git.noahlan.cn/noahlan/nnet/v2/pkg/executor"
|
||||
)
|
||||
|
||||
// handlerExecutor owns the optional asynchronous handler executor. An
|
||||
// injected executor is deliberately not owned by the server.
|
||||
type handlerExecutor struct {
|
||||
exec executor.Executor
|
||||
owned bool
|
||||
|
||||
shutdownOnce sync.Once
|
||||
shutdownErr error
|
||||
}
|
||||
|
||||
func newHandlerExecutor(cfg *config.Config, log logger.Logger) *handlerExecutor {
|
||||
if cfg == nil || cfg.HandlerExecutionMode != config.HandlerExecutionPerConnection {
|
||||
return nil
|
||||
}
|
||||
|
||||
if cfg.HandlerExecutor != nil {
|
||||
return &handlerExecutor{exec: cfg.HandlerExecutor}
|
||||
}
|
||||
|
||||
return &handlerExecutor{
|
||||
exec: executor.NewPerKey(executor.Config{
|
||||
Workers: cfg.HandlerWorkers,
|
||||
QueueSize: cfg.HandlerQueueSize,
|
||||
OnPanic: func(value interface{}) {
|
||||
log.Error("Handler panic recovered: %v", value)
|
||||
},
|
||||
}),
|
||||
owned: true,
|
||||
}
|
||||
}
|
||||
|
||||
func (h *handlerExecutor) submit(key string, task executor.Task) error {
|
||||
if h == nil || h.exec == nil {
|
||||
return nil
|
||||
}
|
||||
return h.exec.Submit(key, task)
|
||||
}
|
||||
|
||||
func (h *handlerExecutor) shutdown(ctx context.Context) error {
|
||||
if h == nil || !h.owned || h.exec == nil {
|
||||
return nil
|
||||
}
|
||||
h.shutdownOnce.Do(func() {
|
||||
h.shutdownErr = h.exec.Shutdown(ctx)
|
||||
})
|
||||
return h.shutdownErr
|
||||
}
|
||||
@ -0,0 +1,192 @@
|
||||
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()
|
||||
}
|
||||
@ -0,0 +1,187 @@
|
||||
package executor
|
||||
|
||||
import (
|
||||
"context"
|
||||
"errors"
|
||||
"sync"
|
||||
"testing"
|
||||
"time"
|
||||
)
|
||||
|
||||
func TestPerKeyPreservesOrderAndRunsDifferentKeysConcurrently(t *testing.T) {
|
||||
e := NewPerKey(Config{Workers: 2, QueueSize: 8})
|
||||
defer func() {
|
||||
ctx, cancel := context.WithTimeout(context.Background(), time.Second)
|
||||
defer cancel()
|
||||
if err := e.Shutdown(ctx); err != nil {
|
||||
t.Fatalf("shutdown error: %v", err)
|
||||
}
|
||||
}()
|
||||
|
||||
var mu sync.Mutex
|
||||
var order []int
|
||||
firstStarted := make(chan struct{})
|
||||
releaseFirst := make(chan struct{})
|
||||
secondDone := make(chan struct{})
|
||||
|
||||
if err := e.Submit("a", func() {
|
||||
close(firstStarted)
|
||||
<-releaseFirst
|
||||
mu.Lock()
|
||||
order = append(order, 1)
|
||||
mu.Unlock()
|
||||
}); err != nil {
|
||||
t.Fatalf("submit first task: %v", err)
|
||||
}
|
||||
if err := e.Submit("a", func() {
|
||||
mu.Lock()
|
||||
order = append(order, 2)
|
||||
mu.Unlock()
|
||||
}); err != nil {
|
||||
t.Fatalf("submit second task: %v", err)
|
||||
}
|
||||
if err := e.Submit("b", func() {
|
||||
close(secondDone)
|
||||
}); err != nil {
|
||||
t.Fatalf("submit different-key task: %v", err)
|
||||
}
|
||||
|
||||
select {
|
||||
case <-firstStarted:
|
||||
case <-time.After(time.Second):
|
||||
t.Fatal("first task did not start")
|
||||
}
|
||||
select {
|
||||
case <-secondDone:
|
||||
case <-time.After(time.Second):
|
||||
t.Fatal("different-key task did not run concurrently")
|
||||
}
|
||||
close(releaseFirst)
|
||||
|
||||
deadline := time.After(time.Second)
|
||||
for {
|
||||
mu.Lock()
|
||||
done := len(order) == 2
|
||||
mu.Unlock()
|
||||
if done {
|
||||
break
|
||||
}
|
||||
select {
|
||||
case <-deadline:
|
||||
t.Fatal("same-key tasks did not drain")
|
||||
default:
|
||||
time.Sleep(time.Millisecond)
|
||||
}
|
||||
}
|
||||
|
||||
mu.Lock()
|
||||
defer mu.Unlock()
|
||||
if order[0] != 1 || order[1] != 2 {
|
||||
t.Fatalf("order = %v, want [1 2]", order)
|
||||
}
|
||||
}
|
||||
|
||||
func TestPerKeyRejectsFullQueueWithoutBlocking(t *testing.T) {
|
||||
e := NewPerKey(Config{Workers: 1, QueueSize: 1})
|
||||
defer func() {
|
||||
closeTask := make(chan struct{})
|
||||
_ = e.Submit("a", func() { close(closeTask) })
|
||||
ctx, cancel := context.WithTimeout(context.Background(), time.Second)
|
||||
defer cancel()
|
||||
_ = e.Shutdown(ctx)
|
||||
}()
|
||||
|
||||
started := make(chan struct{})
|
||||
release := make(chan struct{})
|
||||
if err := e.Submit("a", func() {
|
||||
close(started)
|
||||
<-release
|
||||
}); err != nil {
|
||||
t.Fatalf("submit running task: %v", err)
|
||||
}
|
||||
select {
|
||||
case <-started:
|
||||
case <-time.After(time.Second):
|
||||
t.Fatal("running task did not start")
|
||||
}
|
||||
|
||||
if err := e.Submit("a", func() {}); err != nil {
|
||||
t.Fatalf("submit queued task: %v", err)
|
||||
}
|
||||
|
||||
done := make(chan error, 1)
|
||||
go func() {
|
||||
done <- e.Submit("a", func() {})
|
||||
}()
|
||||
select {
|
||||
case err := <-done:
|
||||
if !errors.Is(err, ErrQueueFull) {
|
||||
t.Fatalf("submit error = %v, want ErrQueueFull", err)
|
||||
}
|
||||
case <-time.After(100 * time.Millisecond):
|
||||
t.Fatal("queue-full submit blocked")
|
||||
}
|
||||
close(release)
|
||||
}
|
||||
|
||||
func TestPerKeyShutdownDrainsAcceptedTasks(t *testing.T) {
|
||||
e := NewPerKey(Config{Workers: 1, QueueSize: 4})
|
||||
var countMu sync.Mutex
|
||||
count := 0
|
||||
for i := 0; i < 3; i++ {
|
||||
if err := e.Submit("a", func() {
|
||||
countMu.Lock()
|
||||
count++
|
||||
countMu.Unlock()
|
||||
}); err != nil {
|
||||
t.Fatalf("submit task %d: %v", i, err)
|
||||
}
|
||||
}
|
||||
|
||||
ctx, cancel := context.WithTimeout(context.Background(), time.Second)
|
||||
defer cancel()
|
||||
if err := e.Shutdown(ctx); err != nil {
|
||||
t.Fatalf("shutdown error: %v", err)
|
||||
}
|
||||
if err := e.Submit("a", func() {}); !errors.Is(err, ErrClosed) {
|
||||
t.Fatalf("submit after shutdown error = %v, want ErrClosed", err)
|
||||
}
|
||||
|
||||
countMu.Lock()
|
||||
defer countMu.Unlock()
|
||||
if count != 3 {
|
||||
t.Fatalf("completed tasks = %d, want 3", count)
|
||||
}
|
||||
}
|
||||
|
||||
func TestPerKeyRecoversTaskPanic(t *testing.T) {
|
||||
panicSeen := make(chan interface{}, 1)
|
||||
e := NewPerKey(Config{
|
||||
Workers: 1,
|
||||
QueueSize: 2,
|
||||
OnPanic: func(value interface{}) {
|
||||
panicSeen <- value
|
||||
},
|
||||
})
|
||||
if err := e.Submit("a", func() { panic("boom") }); err != nil {
|
||||
t.Fatalf("submit panic task: %v", err)
|
||||
}
|
||||
if err := e.Submit("a", func() {}); err != nil {
|
||||
t.Fatalf("submit follow-up task: %v", err)
|
||||
}
|
||||
|
||||
select {
|
||||
case value := <-panicSeen:
|
||||
if value != "boom" {
|
||||
t.Fatalf("panic value = %v, want boom", value)
|
||||
}
|
||||
case <-time.After(time.Second):
|
||||
t.Fatal("panic callback was not called")
|
||||
}
|
||||
|
||||
ctx, cancel := context.WithTimeout(context.Background(), time.Second)
|
||||
defer cancel()
|
||||
if err := e.Shutdown(ctx); err != nil {
|
||||
t.Fatalf("shutdown error: %v", err)
|
||||
}
|
||||
}
|
||||
Loading…
Reference in New Issue