From 0cf53e2f0dd44616bc29439565f41aefeb85fbe5 Mon Sep 17 00:00:00 2001
From: Raja Subramanian <raja.gobi@tutanota.com>
Date: Tue, 23 Jun 2026 16:10:50 +0530
Subject: [PATCH] Add option to force drain rtcService/agentService
 connections. (#4618)

When force: true, drain as fast as possible.
---
 pkg/agent/testutils/server.go |  4 ++--
 pkg/service/agentservice.go   | 30 ++++++++++++++++++++----------
 pkg/service/rtcservice.go     | 23 +++++++++++++++--------
 3 files changed, 37 insertions(+), 20 deletions(-)

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