diff --git a/internal/proxy/udprelay/run.go b/internal/proxy/udprelay/run.go index 1d80a8e..a7395a5 100644 --- a/internal/proxy/udprelay/run.go +++ b/internal/proxy/udprelay/run.go @@ -51,6 +51,7 @@ type Deps struct { ConnectedStreams *atomic.Int32 OnTURNServer func(ip net.IP) fatalCh chan error + allocPace *allocPacer } func (d *Deps) log() logx.Logger { @@ -91,6 +92,7 @@ func Run(ctx context.Context, dtlsDialer *dtlsdial.Dialer, auth AuthHandler, log ConnectedStreams: connectedStreams, OnTURNServer: onTURNServer, fatalCh: fatalCh, + allocPace: newAllocPacer(allocPaceInterval), } runCtx, runCancel := context.WithCancel(ctx) @@ -105,6 +107,18 @@ func Run(ctx context.Context, dtlsDialer *dtlsdial.Dialer, auth AuthHandler, log } }() + var fatalErr atomic.Pointer[error] + watcherDone := make(chan struct{}) + go func() { + defer close(watcherDone) + select { + case err := <-fatalCh: + fatalErr.Store(&err) + runCancel() + case <-runCtx.Done(): + } + }() + inboundChan := make(chan *Packet, inboundQueueCap) wg := sync.WaitGroup{} wg.Go(deps.guard(func() { @@ -140,18 +154,6 @@ func Run(ctx context.Context, dtlsDialer *dtlsdial.Dialer, auth AuthHandler, log })) } - var fatalErr atomic.Pointer[error] - watcherDone := make(chan struct{}) - go func() { - defer close(watcherDone) - select { - case err := <-fatalCh: - fatalErr.Store(&err) - runCancel() - case <-runCtx.Done(): - } - }() - wg.Wait() runCancel() <-deadlineSet diff --git a/internal/proxy/udprelay/run_test.go b/internal/proxy/udprelay/run_test.go index e116be5..909f721 100644 --- a/internal/proxy/udprelay/run_test.go +++ b/internal/proxy/udprelay/run_test.go @@ -119,3 +119,27 @@ func TestRunReturnsOnStreamPanic(t *testing.T) { t.Fatal("Run hung on panic in stream goroutine") } } + +func TestRunFatalDoesNotWaitWarmupBarrier(t *testing.T) { + t.Parallel() + dialer, params, peer, local := runFatalDeps(t) + + var connected atomic.Int32 + done := make(chan error, 1) + start := time.Now() + go func() { + done <- Run(context.Background(), dialer, stubAuth{}, logx.Nop(), &connected, nil, params, peer, local, 4) + }() + + select { + case err := <-done: + if !errors.Is(err, ErrFatal) { + t.Fatalf("err = %v, want ErrFatal", err) + } + if elapsed := time.Since(start); elapsed >= streamStartBarrier { + t.Fatalf("Run took %v, want less than warm-up barrier %v", elapsed, streamStartBarrier) + } + case <-time.After(2 * streamStartBarrier): + t.Fatal("Run hung on fatal provider error") + } +}