Add option to force drain rtcService/agentService connections. (#4618)

When force: true, drain as fast as possible.

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)

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()
+		}
 	}
 }
 

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()
+		}
 	}
 }