diff --git a/packages/envd/internal/services/process/handler/multiplex.go b/packages/envd/internal/services/process/handler/multiplex.go index 4ba16e9dc0..131f097981 100644 --- a/packages/envd/internal/services/process/handler/multiplex.go +++ b/packages/envd/internal/services/process/handler/multiplex.go @@ -5,40 +5,88 @@ import ( "sync/atomic" ) +// MultiplexedChannel fans out values written to Source to every subscriber +// obtained via Fork. Each subscriber send is guarded by a done channel so +// a cancelled consumer can never wedge the fan-out loop. type MultiplexedChannel[T any] struct { - Source chan T - channels []chan T + Source chan T + mu sync.RWMutex + channels []*subscriber[T] exited atomic.Bool } +type subscriber[T any] struct { + ch chan T + done chan struct{} + once sync.Once +} + +// cancel marks the subscriber as gone. Idempotent and non-blocking. +func (s *subscriber[T]) cancel() { + s.once.Do(func() { + close(s.done) + }) +} + +// isCancelled reports whether cancel has been called. +func (s *subscriber[T]) isCancelled() bool { + select { + case <-s.done: + return true + default: + return false + } +} + func NewMultiplexedChannel[T any](buffer int) *MultiplexedChannel[T] { c := &MultiplexedChannel[T]{ - channels: nil, - Source: make(chan T, buffer), + Source: make(chan T, buffer), } - go func() { - for v := range c.Source { - c.mu.RLock() + go c.run() + + return c +} - for _, cons := range c.channels { - cons <- v +// run is the fan-out loop. It delivers each Source value to every live +// subscriber and closes all consumer channels when Source is closed. +func (m *MultiplexedChannel[T]) run() { + for v := range m.Source { + m.mu.RLock() + subs := m.channels + m.mu.RUnlock() + + for _, s := range subs { + // Skip already-cancelled subscribers. + if s.isCancelled() { + continue } - c.mu.RUnlock() + select { + case s.ch <- v: + case <-s.done: + } } + } - c.exited.Store(true) + m.exited.Store(true) - for _, cons := range c.channels { - close(cons) - } - }() + // Close all remaining consumer channels so `for range` loops exit. + m.mu.Lock() + defer m.mu.Unlock() - return c + for _, s := range m.channels { + s.cancel() + close(s.ch) + } + m.channels = nil } +// Fork registers a new subscriber and returns its channel plus a cancel func. +// If Source is already closed it returns a pre-closed channel and a no-op cancel. +// The channel is bidirectional for backwards compat with start.go which writes +// a bootstrap event into it; new callers should treat it as receive-only. func (m *MultiplexedChannel[T]) Fork() (chan T, func()) { if m.exited.Load() { ch := make(chan T) @@ -50,21 +98,36 @@ func (m *MultiplexedChannel[T]) Fork() (chan T, func()) { m.mu.Lock() defer m.mu.Unlock() - consumer := make(chan T) + // Re-check under lock in case run() finished between the fast path and here. + if m.exited.Load() { + ch := make(chan T) + close(ch) - m.channels = append(m.channels, consumer) + return ch, func() {} + } + + s := &subscriber[T]{ + ch: make(chan T), + done: make(chan struct{}), + } + + m.channels = append(m.channels, s) - return consumer, func() { - m.remove(consumer) + return s.ch, func() { + m.remove(s) } } -func (m *MultiplexedChannel[T]) remove(consumer chan T) { +// remove unsubscribes s. Safe to call multiple times. +func (m *MultiplexedChannel[T]) remove(s *subscriber[T]) { + // Cancel before locking so an in-flight fan-out send can unblock. + s.cancel() + m.mu.Lock() defer m.mu.Unlock() - for i, ch := range m.channels { - if ch == consumer { + for i, sub := range m.channels { + if sub == s { m.channels = append(m.channels[:i], m.channels[i+1:]...) return diff --git a/packages/envd/internal/services/process/handler/multiplex_test.go b/packages/envd/internal/services/process/handler/multiplex_test.go new file mode 100644 index 0000000000..ae79495157 --- /dev/null +++ b/packages/envd/internal/services/process/handler/multiplex_test.go @@ -0,0 +1,314 @@ +package handler + +import ( + "runtime" + "sync" + "testing" + "time" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +// Tests for MultiplexedChannel fan-out, covering the fix for the goroutine +// leak that occurred when a subscriber disconnected mid-send. + +const multiplexTestTimeout = 500 * time.Millisecond + +// recvOrTimeout reads one value from ch or returns ok=false after timeout. +func recvOrTimeout[T any](t *testing.T, ch <-chan T, timeout time.Duration) (T, bool) { + t.Helper() + + select { + case v, ok := <-ch: + return v, ok + case <-time.After(timeout): + var zero T + + return zero, false + } +} + +// sendOrTimeout pushes v into ch or returns false after timeout. +func sendOrTimeout[T any](t *testing.T, ch chan<- T, v T, timeout time.Duration) bool { + t.Helper() + + select { + case ch <- v: + return true + case <-time.After(timeout): + return false + } +} + +func TestMultiplexedChannel_BasicFanOut(t *testing.T) { + t.Parallel() + + m := NewMultiplexedChannel[int](1) + + consA, cancelA := m.Fork() + consB, cancelB := m.Fork() + t.Cleanup(cancelA) + t.Cleanup(cancelB) + + var wg sync.WaitGroup + wg.Add(2) + + gotA := make([]int, 0, 3) + gotB := make([]int, 0, 3) + + go func() { + defer wg.Done() + for v := range consA { + gotA = append(gotA, v) + } + }() + + go func() { + defer wg.Done() + for v := range consB { + gotB = append(gotB, v) + } + }() + + for _, v := range []int{1, 2, 3} { + require.True(t, + sendOrTimeout(t, m.Source, v, multiplexTestTimeout), + "basic fan-out should not block when subscribers drain", + ) + } + + close(m.Source) + wg.Wait() + + assert.Equal(t, []int{1, 2, 3}, gotA) + assert.Equal(t, []int{1, 2, 3}, gotB) +} + +// Regression: an abandoned consumer must not wedge the fan-out loop. +func TestMultiplexedChannel_AbandonedConsumerDoesNotWedgeFanOut(t *testing.T) { + t.Parallel() + + m := NewMultiplexedChannel[int](1) + t.Cleanup(func() { close(m.Source) }) + + abandoned, cancelAbandoned := m.Fork() + + // Consumer reads one value then exits, modeling a disconnected client. + abandonReader := make(chan struct{}) + go func() { + <-abandoned + close(abandonReader) + }() + + require.True(t, + sendOrTimeout(t, m.Source, 1, multiplexTestTimeout), + "first send should be deliverable", + ) + + select { + case <-abandonReader: + case <-time.After(multiplexTestTimeout): + t.Fatal("abandoned consumer should have read its single value") + } + + // Simulate the handler's deferred cancel after return. + cancelDone := make(chan struct{}) + go func() { + cancelAbandoned() + close(cancelDone) + }() + select { + case <-cancelDone: + case <-time.After(multiplexTestTimeout): + t.Fatal("cancel func did not return promptly; fan-out is wedged") + } + + // Producer should still make progress through Source. + for i := 2; i <= 8; i++ { + require.Truef(t, + sendOrTimeout(t, m.Source, i, multiplexTestTimeout), + "send %d should not be back-pressured by an abandoned consumer", i, + ) + } +} + +// Regression: an abandoned consumer must not starve other subscribers. +func TestMultiplexedChannel_AbandonedConsumerDoesNotStarveOthers(t *testing.T) { + t.Parallel() + + m := NewMultiplexedChannel[int](1) + t.Cleanup(func() { close(m.Source) }) + + healthy, cancelHealthy := m.Fork() + t.Cleanup(cancelHealthy) + + healthyReceived := make(chan int, 16) + go func() { + for v := range healthy { + healthyReceived <- v + } + }() + + abandoned, cancelAbandoned := m.Fork() + go func() { + <-abandoned + }() + + require.True(t, + sendOrTimeout(t, m.Source, 1, multiplexTestTimeout), + "first send should pass while both consumers are active", + ) + got, ok := recvOrTimeout(t, healthyReceived, multiplexTestTimeout) + require.True(t, ok, "healthy subscriber should receive value 1") + assert.Equal(t, 1, got) + + // Abandon and cancel the second subscriber. + cancelAbandoned() + + // Healthy subscriber must keep receiving every subsequent value. + for i := 2; i <= 6; i++ { + require.Truef(t, + sendOrTimeout(t, m.Source, i, multiplexTestTimeout), + "send %d should not be back-pressured", i, + ) + got, ok = recvOrTimeout(t, healthyReceived, multiplexTestTimeout) + require.Truef(t, ok, "healthy subscriber should still receive value %d", i) + assert.Equalf(t, i, got, "healthy subscriber received wrong value") + } +} + +// cancel must be idempotent and non-blocking even under producer load. +func TestMultiplexedChannel_CancelIsIdempotentAndPrompt(t *testing.T) { + t.Parallel() + + m := NewMultiplexedChannel[int](0) + + _, cancel := m.Fork() + + // Concurrently push values without anyone draining the consumer chan. + stop := make(chan struct{}) + producerDone := make(chan struct{}) + go func() { + defer close(producerDone) + for { + select { + case <-stop: + return + case m.Source <- 0: + } + } + }() + t.Cleanup(func() { + close(stop) + <-producerDone + close(m.Source) + }) + + // Give the fan-out a chance to enter a per-subscriber select. + time.Sleep(20 * time.Millisecond) + + cancelDone := make(chan struct{}) + go func() { + cancel() + cancel() // idempotent + cancel() + close(cancelDone) + }() + select { + case <-cancelDone: + case <-time.After(multiplexTestTimeout): + t.Fatal("cancel func did not return promptly under producer load") + } +} + +// Closing Source must close all live subscriber channels. +func TestMultiplexedChannel_SourceCloseClosesLiveSubscribers(t *testing.T) { + t.Parallel() + + m := NewMultiplexedChannel[int](1) + + cons, cancel := m.Fork() + t.Cleanup(cancel) + + done := make(chan struct{}) + go func() { + defer close(done) + for range cons { //nolint:revive // drain until closed + } + }() + + require.True(t, + sendOrTimeout(t, m.Source, 1, multiplexTestTimeout), + "send should succeed", + ) + + close(m.Source) + + select { + case <-done: + case <-time.After(multiplexTestTimeout): + t.Fatal("consumer's `for v := range cons` did not terminate after " + + "Source close") + } +} + +// Fork after Source close must return a pre-closed channel. +func TestMultiplexedChannel_ForkAfterSourceCloseReturnsClosedChan(t *testing.T) { + t.Parallel() + + m := NewMultiplexedChannel[int](0) + close(m.Source) + + // Wait for the fan-out goroutine to observe Source close. + deadline := time.Now().Add(multiplexTestTimeout) + for !m.exited.Load() { + if time.Now().After(deadline) { + t.Fatal("fan-out did not mark itself exited after Source close") + } + time.Sleep(time.Millisecond) + } + + cons, cancel := m.Fork() + cancel() // must not panic + + _, ok := recvOrTimeout(t, cons, multiplexTestTimeout) + assert.False(t, ok, "Fork after shutdown must return a pre-closed channel") +} + +// Goroutine count must return to baseline after cancelled subscribers settle. +func TestMultiplexedChannel_NoGoroutineLeakOnAbandon(t *testing.T) { //nolint:paralleltest // relies on a stable goroutine count + const wedges = 16 + + time.Sleep(50 * time.Millisecond) + runtime.GC() //nolint:revive // intentional: settle goroutines before measuring baseline + before := runtime.NumGoroutine() + + for range wedges { + m := NewMultiplexedChannel[int](1) + _, cancel := m.Fork() + // Park one value so fan-out has work, then cancel mid-iteration. + m.Source <- 0 + cancel() + close(m.Source) + } + + // Allow scheduled goroutines to finish. + deadline := time.Now().Add(2 * time.Second) + for time.Now().Before(deadline) { + runtime.GC() //nolint:revive // intentional: help goroutines finalize + runtime.Gosched() + time.Sleep(20 * time.Millisecond) + if runtime.NumGoroutine() <= before+2 { + break + } + } + + after := runtime.NumGoroutine() + leaked := after - before + // Small slack for runtime bookkeeping; the old bug leaked >= wedges. + assert.LessOrEqualf(t, leaked, 2, + "goroutine count grew by %d after %d cancelled wedges; "+ + "before=%d after=%d (expected ~0)", leaked, wedges, before, after, + ) +} diff --git a/packages/envd/pkg/version.go b/packages/envd/pkg/version.go index 9bbb0c85e9..5065015f4b 100644 --- a/packages/envd/pkg/version.go +++ b/packages/envd/pkg/version.go @@ -1,3 +1,3 @@ package pkg -const Version = "0.5.15" +const Version = "0.5.16"