From 04469d97a513b6f946e8488dd20a511374c7eca6 Mon Sep 17 00:00:00 2001 From: YuminosukeSato Date: Fri, 6 Feb 2026 23:03:34 +0900 Subject: [PATCH 1/2] feat: serialize dispatch per worker --- pkg/pyproc/config.go | 12 +-- pkg/pyproc/pool.go | 163 +++++++++++++++++++++++++---------- pkg/pyproc/pool_transport.go | 5 +- 3 files changed, 128 insertions(+), 52 deletions(-) diff --git a/pkg/pyproc/config.go b/pkg/pyproc/config.go index 8fbf61c..07a58b5 100644 --- a/pkg/pyproc/config.go +++ b/pkg/pyproc/config.go @@ -21,11 +21,12 @@ type Config struct { // PoolConfig defines worker pool settings type PoolConfig struct { - Workers int `mapstructure:"workers"` - MaxInFlight int `mapstructure:"max_in_flight"` - StartTimeout time.Duration `mapstructure:"start_timeout"` - HealthInterval time.Duration `mapstructure:"health_interval"` - Restart RestartConfig `mapstructure:"restart"` + Workers int `mapstructure:"workers"` + MaxInFlight int `mapstructure:"max_in_flight"` + MaxInFlightPerWorker int `mapstructure:"max_in_flight_per_worker"` + StartTimeout time.Duration `mapstructure:"start_timeout"` + HealthInterval time.Duration `mapstructure:"health_interval"` + Restart RestartConfig `mapstructure:"restart"` } // RestartConfig defines restart policy @@ -122,6 +123,7 @@ func setDefaults(v *viper.Viper) { // Pool defaults v.SetDefault("pool.workers", 4) v.SetDefault("pool.max_in_flight", 10) + v.SetDefault("pool.max_in_flight_per_worker", 1) v.SetDefault("pool.start_timeout", 30) v.SetDefault("pool.health_interval", 30) v.SetDefault("pool.restart.max_attempts", 5) diff --git a/pkg/pyproc/pool.go b/pkg/pyproc/pool.go index 2f4d543..fe5a876 100644 --- a/pkg/pyproc/pool.go +++ b/pkg/pyproc/pool.go @@ -35,16 +35,20 @@ type workerHandle interface { // Pool manages multiple Python workers with load balancing type Pool struct { - opts PoolOptions - logger *Logger - workers []*poolWorker - nextIdx atomic.Uint64 - shutdown atomic.Bool - started atomic.Bool - wg sync.WaitGroup + opts PoolOptions + logger *Logger + workers []*poolWorker + nextIdx atomic.Uint64 + shutdown atomic.Bool + started atomic.Bool + activeCallsWG sync.WaitGroup + callsMu sync.Mutex + wg sync.WaitGroup // Backpressure control - semaphore chan struct{} + semaphore chan struct{} + workerAvailable chan struct{} + shutdownCh chan struct{} // Health monitoring healthMu sync.RWMutex @@ -67,10 +71,11 @@ type activeRequest struct { // poolWorker wraps a Worker with connection pooling type poolWorker struct { - worker workerHandle - connPool chan net.Conn - requestID atomic.Uint64 - healthy atomic.Bool + worker workerHandle + connPool chan net.Conn + inflightGate chan struct{} + requestID atomic.Uint64 + healthy atomic.Bool } // HealthStatus represents the health of the pool @@ -94,6 +99,9 @@ func NewPool(opts PoolOptions, logger *Logger) (*Pool, error) { if opts.Config.MaxInFlight <= 0 { opts.Config.MaxInFlight = 10 } + if opts.Config.MaxInFlightPerWorker <= 0 { + opts.Config.MaxInFlightPerWorker = 1 + } if opts.Config.HealthInterval <= 0 { opts.Config.HealthInterval = 30 * time.Second } @@ -103,11 +111,13 @@ func NewPool(opts PoolOptions, logger *Logger) (*Pool, error) { } pool := &Pool{ - opts: opts, - logger: logger, - workers: make([]*poolWorker, opts.Config.Workers), - semaphore: make(chan struct{}, opts.Config.Workers*opts.Config.MaxInFlight), - activeRequests: make(map[uint64]*activeRequest), + opts: opts, + logger: logger, + workers: make([]*poolWorker, opts.Config.Workers), + semaphore: make(chan struct{}, opts.Config.MaxInFlight), + workerAvailable: make(chan struct{}, opts.Config.Workers*opts.Config.MaxInFlightPerWorker), + shutdownCh: make(chan struct{}), + activeRequests: make(map[uint64]*activeRequest), } // Create workers @@ -121,8 +131,9 @@ func NewPool(opts PoolOptions, logger *Logger) (*Pool, error) { worker := NewWorker(workerCfg, logger) pool.workers[i] = &poolWorker{ - worker: worker, - connPool: make(chan net.Conn, opts.Config.MaxInFlight), + worker: worker, + connPool: make(chan net.Conn, opts.Config.MaxInFlightPerWorker), + inflightGate: make(chan struct{}, opts.Config.MaxInFlightPerWorker), } } @@ -141,6 +152,9 @@ func newExternalPool(opts PoolOptions, logger *Logger) (*Pool, error) { if opts.Config.MaxInFlight <= 0 { opts.Config.MaxInFlight = 10 } + if opts.Config.MaxInFlightPerWorker <= 0 { + opts.Config.MaxInFlightPerWorker = 1 + } if opts.Config.HealthInterval <= 0 { opts.Config.HealthInterval = 30 * time.Second } @@ -150,18 +164,21 @@ func newExternalPool(opts PoolOptions, logger *Logger) (*Pool, error) { } pool := &Pool{ - opts: opts, - logger: logger, - workers: make([]*poolWorker, numWorkers), - semaphore: make(chan struct{}, numWorkers*opts.Config.MaxInFlight), - activeRequests: make(map[uint64]*activeRequest), + opts: opts, + logger: logger, + workers: make([]*poolWorker, numWorkers), + semaphore: make(chan struct{}, opts.Config.MaxInFlight), + workerAvailable: make(chan struct{}, numWorkers*opts.Config.MaxInFlightPerWorker), + shutdownCh: make(chan struct{}), + activeRequests: make(map[uint64]*activeRequest), } for i, sockPath := range opts.ExternalSocketPaths { worker := NewExternalWorker(sockPath, opts.WorkerConfig.StartTimeout) pool.workers[i] = &poolWorker{ - worker: worker, - connPool: make(chan net.Conn, opts.Config.MaxInFlight), + worker: worker, + connPool: make(chan net.Conn, opts.Config.MaxInFlightPerWorker), + inflightGate: make(chan struct{}, opts.Config.MaxInFlightPerWorker), } } @@ -198,7 +215,7 @@ func (p *Pool) Start(ctx context.Context) error { select { case pw.connPool <- conn: default: - // Pool is full (shouldn't happen with MaxInFlight=1), close connection + // Pool is full (shouldn't happen with MaxInFlightPerWorker=1), close connection if err := conn.Close(); err != nil { p.logger.Error("failed to close connection", "error", err) } @@ -224,9 +241,14 @@ func (p *Pool) Start(ctx context.Context) error { // Call invokes a method on one of the workers using round-robin func (p *Pool) Call(ctx context.Context, method string, input interface{}, output interface{}) error { + p.callsMu.Lock() if p.shutdown.Load() { + p.callsMu.Unlock() return errors.New("pool is shut down") } + p.activeCallsWG.Add(1) + p.callsMu.Unlock() + defer p.activeCallsWG.Done() // Acquire semaphore for backpressure select { @@ -234,31 +256,27 @@ func (p *Pool) Call(ctx context.Context, method string, input interface{}, outpu defer func() { <-p.semaphore }() case <-ctx.Done(): return ctx.Err() + case <-p.shutdownCh: + return errors.New("pool is shut down") } - // Select worker using round-robin - idx := p.nextIdx.Add(1) - 1 - workerIdx := int(idx % uint64(len(p.workers))) - pw := p.workers[workerIdx] - - if !pw.healthy.Load() { - // Try to find a healthy worker - for i, w := range p.workers { - if w.healthy.Load() { - pw = w - workerIdx = i - break - } - } - if !pw.healthy.Load() { - return errors.New("no healthy workers available") - } + pw, workerIdx, err := p.acquireWorker(ctx) + if err != nil { + return err } + defer func() { + <-pw.inflightGate + p.signalWorkerAvailable() + }() // Get connection from pool var conn net.Conn select { - case conn = <-pw.connPool: + case pooledConn, ok := <-pw.connPool: + if !ok { + return errors.New("connection pool closed") + } + conn = pooledConn default: // Create new connection if pool is empty var err error @@ -296,12 +314,14 @@ func (p *Pool) Call(ctx context.Context, method string, input interface{}, outpu p.activeRequestsMu.Unlock() // Return connection to pool only if not closed - if !connClosed { + if !connClosed && !p.shutdown.Load() { select { case pw.connPool <- conn: default: _ = conn.Close() } + } else if !connClosed { + _ = conn.Close() } }() @@ -370,6 +390,7 @@ func (p *Pool) Shutdown(_ context.Context) error { if !p.shutdown.CompareAndSwap(false, true) { return nil // Already shutting down } + close(p.shutdownCh) p.logger.Info("shutting down worker pool") @@ -378,6 +399,11 @@ func (p *Pool) Shutdown(_ context.Context) error { p.healthCancel() } + // Wait for in-flight calls to complete before closing pools + p.callsMu.Lock() + p.callsMu.Unlock() + p.activeCallsWG.Wait() + // Close all connection pools for _, pw := range p.workers { close(pw.connPool) @@ -405,6 +431,51 @@ func (p *Pool) Shutdown(_ context.Context) error { return nil } +func (p *Pool) signalWorkerAvailable() { + select { + case p.workerAvailable <- struct{}{}: + default: + } +} + +func (p *Pool) acquireWorker(ctx context.Context) (*poolWorker, int, error) { + workers := p.workers + if len(workers) == 0 { + return nil, -1, errors.New("no workers available") + } + + for { + startIdx := int(p.nextIdx.Add(1) - 1) + healthyFound := false + + for i := 0; i < len(workers); i++ { + idx := (startIdx + i) % len(workers) + pw := workers[idx] + if !pw.healthy.Load() { + continue + } + healthyFound = true + select { + case pw.inflightGate <- struct{}{}: + return pw, idx, nil + default: + } + } + + if !healthyFound { + return nil, -1, errors.New("no healthy workers available") + } + + select { + case <-ctx.Done(): + return nil, -1, ctx.Err() + case <-p.shutdownCh: + return nil, -1, errors.New("pool is shut down") + case <-p.workerAvailable: + } + } +} + // Health returns the current health status of the pool func (p *Pool) Health() HealthStatus { p.healthMu.RLock() diff --git a/pkg/pyproc/pool_transport.go b/pkg/pyproc/pool_transport.go index 57a695a..38536f9 100644 --- a/pkg/pyproc/pool_transport.go +++ b/pkg/pyproc/pool_transport.go @@ -134,6 +134,9 @@ func NewPoolWithTransport(opts PoolOptions, logger *Logger) (*PoolWithTransport, if opts.Config.MaxInFlight <= 0 { opts.Config.MaxInFlight = 10 } + if opts.Config.MaxInFlightPerWorker <= 0 { + opts.Config.MaxInFlightPerWorker = 1 + } if opts.Config.HealthInterval <= 0 { opts.Config.HealthInterval = 30 * time.Second } @@ -146,7 +149,7 @@ func NewPoolWithTransport(opts PoolOptions, logger *Logger) (*PoolWithTransport, opts: opts, logger: logger, workers: make([]*Worker, opts.Config.Workers), - semaphore: make(chan struct{}, opts.Config.Workers*opts.Config.MaxInFlight), + semaphore: make(chan struct{}, opts.Config.MaxInFlight), } // Create workers (they still manage the Python processes) From b4cb3609b17d669f158529bb397afa50f08e562b Mon Sep 17 00:00:00 2001 From: YuminosukeSato Date: Sat, 7 Feb 2026 00:08:47 +0900 Subject: [PATCH 2/2] fix(pool): move WaitGroup.Wait inside lock to fix SA2001 Move activeCallsWG.Wait() inside the callsMu critical section to satisfy staticcheck SA2001 (empty critical section). The lock still serves as a barrier ensuring all in-progress Call() goroutines have completed their Add(1) before Wait() is called. Co-Authored-By: Claude Opus 4.6 --- pkg/pyproc/pool.go | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/pkg/pyproc/pool.go b/pkg/pyproc/pool.go index fe5a876..5b7a073 100644 --- a/pkg/pyproc/pool.go +++ b/pkg/pyproc/pool.go @@ -399,10 +399,12 @@ func (p *Pool) Shutdown(_ context.Context) error { p.healthCancel() } - // Wait for in-flight calls to complete before closing pools + // Wait for in-flight calls to complete before closing pools. + // The lock ensures all in-progress Call() goroutines have finished + // their activeCallsWG.Add(1) before we start waiting. p.callsMu.Lock() - p.callsMu.Unlock() p.activeCallsWG.Wait() + p.callsMu.Unlock() // Close all connection pools for _, pw := range p.workers {