fix(agent): retire replaced tunnel sessions (#541)

go-gost/x/internal/util/sessionretire/sessionretire.go

@@ -0,0 +1,58 @@
+package sessionretire
+
+import "time"
+
+const (
+	defaultIdleGrace  = time.Second
+	defaultPollPeriod = 100 * time.Millisecond
+)
+
+// Session is the lifecycle surface shared by the multiplexed dialers.
+type Session interface {
+	Close() error
+	IsClosed() bool
+	NumStreams() int
+}
+
+// Gracefully closes a retired session after all existing streams have drained.
+// A short idle grace covers the Dial/Handshake hand-off used by several dialers.
+func Gracefully(session Session) {
+	if session == nil {
+		return
+	}
+	go waitUntilIdle(session, defaultIdleGrace, defaultPollPeriod)
+}
+
+func waitUntilIdle(session Session, idleGrace, pollPeriod time.Duration) {
+	if session == nil {
+		return
+	}
+	if idleGrace <= 0 {
+		idleGrace = defaultIdleGrace
+	}
+	if pollPeriod <= 0 {
+		pollPeriod = defaultPollPeriod
+	}
+
+	ticker := time.NewTicker(pollPeriod)
+	defer ticker.Stop()
+
+	var idleSince time.Time
+	for {
+		if session.IsClosed() {
+			_ = session.Close()
+			return
+		}
+		if session.NumStreams() == 0 {
+			if idleSince.IsZero() {
+				idleSince = time.Now()
+			} else if time.Since(idleSince) >= idleGrace {
+				_ = session.Close()
+				return
+			}
+		} else {
+			idleSince = time.Time{}
+		}
+		<-ticker.C
+	}
+}

go-gost/x/internal/util/sessionretire/sessionretire_test.go

@@ -0,0 +1,59 @@
+package sessionretire
+
+import (
+	"sync"
+	"testing"
+	"time"
+)
+
+type testSession struct {
+	mu      sync.Mutex
+	streams int
+	closed  bool
+}
+
+func (s *testSession) Close() error {
+	s.mu.Lock()
+	s.closed = true
+	s.mu.Unlock()
+	return nil
+}
+
+func (s *testSession) IsClosed() bool {
+	s.mu.Lock()
+	defer s.mu.Unlock()
+	return s.closed
+}
+
+func (s *testSession) NumStreams() int {
+	s.mu.Lock()
+	defer s.mu.Unlock()
+	return s.streams
+}
+
+func TestWaitUntilIdlePreservesActiveStreams(t *testing.T) {
+	session := &testSession{streams: 1}
+	done := make(chan struct{})
+	go func() {
+		waitUntilIdle(session, 20*time.Millisecond, time.Millisecond)
+		close(done)
+	}()
+
+	time.Sleep(30 * time.Millisecond)
+	if session.IsClosed() {
+		t.Fatal("active session was closed")
+	}
+
+	session.mu.Lock()
+	session.streams = 0
+	session.mu.Unlock()
+
+	select {
+	case <-done:
+	case <-time.After(250 * time.Millisecond):
+		t.Fatal("idle session was not closed")
+	}
+	if !session.IsClosed() {
+		t.Fatal("retired session did not close after becoming idle")
+	}
+}