diff --git a/pkg/agent/testutils/server.go b/pkg/agent/testutils/server.go index 090e37680..3f490918c 100644 --- a/pkg/agent/testutils/server.go +++ b/pkg/agent/testutils/server.go @@ -29,7 +29,7 @@ import ( type AgentService interface { HandleConnection(context.Context, agent.SignalConn, agent.WorkerRegistration) - DrainConnections(time.Duration) + DrainConnections(time.Duration, bool) } type TestServer struct { @@ -140,7 +140,7 @@ func (h *TestServer) SimulateAgentWorker(opts ...SimulatedWorkerOption) *AgentWo } func (h *TestServer) Close() { - h.DrainConnections(1) + h.DrainConnections(1, false) } var _ agent.SignalConn = (*AgentWorker)(nil) diff --git a/pkg/service/agentservice.go b/pkg/service/agentservice.go index 8481a1e3d..5f95f0722 100644 --- a/pkg/service/agentservice.go +++ b/pkg/service/agentservice.go @@ -485,19 +485,29 @@ func (h *AgentHandler) CheckEnabled(ctx context.Context, req *rpc.CheckEnabledRe }, nil } -func (h *AgentHandler) DrainConnections(interval time.Duration) { - // jitter drain start - time.Sleep(time.Duration(rand.Int63n(int64(interval)))) +func (h *AgentHandler) DrainConnections(interval time.Duration, force bool) { + if !force { + // jitter drain start + time.Sleep(time.Duration(rand.Int63n(int64(interval)))) - t := time.NewTicker(interval) - defer t.Stop() + t := time.NewTicker(interval) + defer t.Stop() - h.mu.Lock() - defer h.mu.Unlock() + h.mu.Lock() + defer h.mu.Unlock() - for _, w := range h.workers { - w.Close() - <-t.C + for _, w := range h.workers { + w.Close() + <-t.C + } + } else { + // drain as quickly as possible when forced + h.mu.Lock() + defer h.mu.Unlock() + + for _, w := range h.workers { + w.Close() + } } } diff --git a/pkg/service/rtcservice.go b/pkg/service/rtcservice.go index ff2d87b25..84082f355 100644 --- a/pkg/service/rtcservice.go +++ b/pkg/service/rtcservice.go @@ -628,20 +628,27 @@ func (s *RTCService) serve(w http.ResponseWriter, r *http.Request, needsJoinRequ } } -func (s *RTCService) DrainConnections(interval time.Duration) { +func (s *RTCService) DrainConnections(interval time.Duration, force bool) { s.mu.Lock() conns := maps.Clone(s.connections) s.mu.Unlock() - // jitter drain start - time.Sleep(time.Duration(rand.Int63n(int64(interval)))) + if !force { + // jitter drain start + time.Sleep(time.Duration(rand.Int63n(int64(interval)))) - t := time.NewTicker(interval) - defer t.Stop() + t := time.NewTicker(interval) + defer t.Stop() - for c := range conns { - _ = c.Close() - <-t.C + for c := range conns { + _ = c.Close() + <-t.C + } + } else { + // drain as quickly as possible when forced + for c := range conns { + _ = c.Close() + } } }