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