Merge pull request #6626 from hashicorp/ICU-18958

Memory leak fix for worker session management by immediately removing closed connections

internal/daemon/worker/session/manager_test.go

@@ -269,7 +269,46 @@ func createTestCert(t *testing.T) []byte {
 	return certBytes
 }
 
-func TestWorkerSetCloseTimeForResponse(t *testing.T) {
+// TestRemoveClosedConnections_CancelFuncInvoked verifies that the cancel func
+// stored on a ConnInfo is called when the controller confirms the connection
+// closed — this is the primary goroutine-leak fix.
+func TestRemoveClosedConnections_CancelFuncInvoked(t *testing.T) {
+	cancelled := false
+	m := new(sync.Map)
+	m.Store("sess1", &sess{
+		resp: &pbs.LookupSessionResponse{Authorization: &targets.SessionAuthorizationData{
+			SessionId: "sess1",
+		}},
+		sessionId: "sess1",
+		connInfoMap: map[string]*ConnInfo{
+			"conn1": {
+				Id: "conn1",
+				connCtxCancelFunc: func() {
+					cancelled = true
+				},
+			},
+		},
+	})
+
+	manager := &manager{sessionMap: m}
+	closedIds, errs := removeClosedConnections(manager, map[string][]*pbs.CloseConnectionResponseData{
+		"sess1": {
+			{ConnectionId: "conn1", Status: pbs.CONNECTIONSTATUS_CONNECTIONSTATUS_CLOSED},
+		},
+	})
+
+	require.Empty(t, errs)
+	require.ElementsMatch(t, []string{"conn1"}, closedIds)
+	assert.True(t, cancelled, "cancel func must be invoked when connection is confirmed closed")
+
+	// Verify the entry is pruned from the map.
+	si := manager.Get("sess1")
+	require.NotNil(t, si)
+	_, exists := si.GetLocalConnections()["conn1"]
+	assert.False(t, exists, "closed connection should be pruned from local state")
+}
+
+func TestWorkerRemoveClosedConnections(t *testing.T) {
 	cases := []struct {
 		name             string
 		sessionCloseInfo map[string][]*pbs.CloseConnectionResponseData
@@ -454,19 +493,21 @@ func TestWorkerSetCloseTimeForResponse(t *testing.T) {
 			require := require.New(t)
 			manager := &manager{sessionMap: new(sync.Map)}
 			manager.sessionMap = tc.sessionInfoMap()
-			actual, actualErr := setCloseTimeForResponse(manager, tc.sessionCloseInfo)
+			actual, actualErr := removeClosedConnections(manager, tc.sessionCloseInfo)
 
-			// Assert all close times were set
+			// Assert closed connections are pruned from local state and
+			// open connections are still present.
 			manager.ForEachLocalSession(func(value Session) bool {
 				t.Helper()
-				for _, ci := range value.GetLocalConnections() {
-					if _, ok := tc.expectedClosed[ci.Id]; ok {
-						require.NotEqual(time.Time{}, ci.CloseTime)
-					} else {
-						require.Equal(time.Time{}, ci.CloseTime)
-					}
+				localConns := value.GetLocalConnections()
+				for connId := range tc.expectedClosed {
+					_, exists := localConns[connId]
+					require.False(exists, "expected closed connection %q to be pruned from local state", connId)
+				}
+				for _, ci := range localConns {
+					_, shouldBeClosed := tc.expectedClosed[ci.Id]
+					require.False(shouldBeClosed, "connection %q should have been pruned from local state", ci.Id)
 				}
-
 				return true
 			})
 

internal/daemon/worker/session/session.go

@@ -36,10 +36,6 @@ type ConnInfo struct {
 	BytesUp func() int64
 	// The number of bytes downloaded to the client.
 	BytesDown func() int64
-
-	// The time the controller has successfully reported that this connection is
-	// closed.
-	CloseTime time.Time
 }
 
 // ConnectionCloseData encapsulates the data we need to send via CloseConnection
@@ -171,7 +167,14 @@ func (s *sess) ApplyLocalConnectionStatus(connId string, status pbs.CONNECTIONST
 
 	connInfo.Status = status
 	if connInfo.Status == pbs.CONNECTIONSTATUS_CONNECTIONSTATUS_CLOSED {
-		connInfo.CloseTime = time.Now()
+		// Cancel the proxy context so the goroutine is stopped in the
+		// controller-initiated close path, then immediately remove the
+		// entry from the map to prevent unbounded growth when a session
+		// carries a large number of short-lived connections.
+		if connInfo.connCtxCancelFunc != nil {
+			connInfo.connCtxCancelFunc()
+		}
+		delete(s.connInfoMap, connId)
 	}
 	return nil
 }
@@ -204,7 +207,6 @@ func (s *sess) GetLocalConnections() map[string]ConnInfo {
 		res[k] = ConnInfo{
 			Id:        v.Id,
 			Status:    v.Status,
-			CloseTime: v.CloseTime,
 			BytesUp:   v.BytesUp,
 			BytesDown: v.BytesDown,
 		}
@@ -356,30 +358,30 @@ func (s *sess) RequestConnectConnection(ctx context.Context, info *pbs.ConnectCo
 	return nil
 }
 
+// cancelLocalConnections iterates connInfoMap, cancels the proxy context for
+// each matching connection, and returns the cancelled connection IDs.
+// If onlyClosed is true, only connections with status CLOSED are cancelled.
+func (s *sess) cancelLocalConnections(onlyClosed bool) []string {
+	s.lock.RLock()
+	defer s.lock.RUnlock()
+	var cancelledIds []string
+	for k, v := range s.connInfoMap {
+		if onlyClosed && v.Status != pbs.CONNECTIONSTATUS_CONNECTIONSTATUS_CLOSED {
+			continue
+		}
+		v.connCtxCancelFunc()
+		cancelledIds = append(cancelledIds, k)
+	}
+	return cancelledIds
+}
+
 // CancelOpenLocalConnections closes the local connections in this session
 // based on the connection's state by calling the connections context cancel
 // function.
 //
 // The returned slice are connection ids that were closed.
 func (s *sess) CancelOpenLocalConnections() []string {
-	s.lock.RLock()
-	defer s.lock.RUnlock()
-	var closedIds []string
-	for k, v := range s.connInfoMap {
-		if v.Status != pbs.CONNECTIONSTATUS_CONNECTIONSTATUS_CLOSED {
-			continue
-		}
-		v.connCtxCancelFunc()
-
-		// CloseTime is set when the controller has reported that the connection
-		// is closed.  The worker should only request a connection be marked
-		// closed after it has already been cancelled.
-		if v.CloseTime.IsZero() {
-			closedIds = append(closedIds, k)
-		}
-	}
-
-	return closedIds
+	return s.cancelLocalConnections(true)
 }
 
 // CancelAllLocalConnections close connections regardless of connection's state
@@ -387,21 +389,7 @@ func (s *sess) CancelOpenLocalConnections() []string {
 //
 // The returned slice is the connection ids which were closed.
 func (s *sess) CancelAllLocalConnections() []string {
-	s.lock.RLock()
-	defer s.lock.RUnlock()
-	var closedIds []string
-	for k, v := range s.connInfoMap {
-		v.connCtxCancelFunc()
-
-		// CloseTime is set when the controller has reported that the connection
-		// is closed.  The worker should only request a connection be marked
-		// closed after it has already been cancelled.
-		if v.CloseTime.IsZero() {
-			closedIds = append(closedIds, k)
-		}
-	}
-
-	return closedIds
+	return s.cancelLocalConnections(false)
 }
 
 // ApplyConnectionCounterCallbacks sets a connection's bytes up and bytes
@@ -479,7 +467,7 @@ func closeConnection(ctx context.Context, sessClient pbs.SessionServiceClient, r
 }
 
 // closeConnections is a helper worker function that sends connection close
-// requests to the controller, and sets close times within the worker. It is
+// requests to the controller and removes closed connections from local state. It is
 // called during the worker session info loop and on connection exit on the proxy.
 //
 // The boolean indicates whether the function was successful, e.g. had any
@@ -527,7 +515,7 @@ func closeConnections(ctx context.Context, sessClient pbs.SessionServiceClient,
 	}
 
 	// Mark connections as closed
-	_, errs := setCloseTimeForResponse(sManager, sessionCloseInfo)
+	_, errs := removeClosedConnections(sManager, sessionCloseInfo)
 	if len(errs) > 0 {
 		for _, err := range errs {
 			event.WriteError(ctx, op, err, event.WithInfoMsg("error marking connection closed in state"))
@@ -565,7 +553,7 @@ func makeCloseConnectionRequest(closeInfo map[string]*ConnectionCloseData) *pbs.
 // our original closeInfo map and makes a map of slices, indexed by
 // session ID, of all of the connection responses. This allows us to
 // easily lock on session once for all connections in
-// setCloseTimeForResponse.
+// removeClosedConnections.
 func makeSessionCloseInfo(
 	closeInfo map[string]*ConnectionCloseData,
 	response *pbs.CloseConnectionResponse,
@@ -583,7 +571,7 @@ func makeSessionCloseInfo(
 	return result, nil
 }
 
-// makeFakeSessionCloseInfo makes a "fake" makeFakeSessionCloseInfo, intended
+// makeFakeSessionCloseInfo makes a "fake" session close info map, intended
 // for use when we can't contact the controller.
 func makeFakeSessionCloseInfo(
 	closeInfo map[string]*ConnectionCloseData,
@@ -604,9 +592,10 @@ func makeFakeSessionCloseInfo(
 	return result, nil
 }
 
-// setCloseTimeForResponse iterates a CloseConnectionResponse and
-// sets the close time for any connection found to be closed to the
-// current time.
+// removeClosedConnections iterates a CloseConnectionResponse and
+// removes any connection confirmed closed by the controller from the
+// session's local connection map. Removing the entry also cancels the
+// connection's proxy context, stopping any in-flight goroutine.
 //
 // sessionCloseInfo can be derived from the closeInfo supplied to
 // makeCloseConnectionRequest through reverseCloseInfo, which creates
@@ -616,7 +605,7 @@ func makeFakeSessionCloseInfo(
 // failed, as some connections may have been marked as closed. The
 // actual list of connection IDs closed is returned as the first
 // return value.
-func setCloseTimeForResponse(sManager Manager, sessionCloseInfo map[string][]*pbs.CloseConnectionResponseData) ([]string, []error) {
+func removeClosedConnections(sManager Manager, sessionCloseInfo map[string][]*pbs.CloseConnectionResponseData) ([]string, []error) {
 	closedIds := make([]string, 0)
 	var result []error
 	for sessionId, responses := range sessionCloseInfo {
@@ -625,13 +614,11 @@ func setCloseTimeForResponse(sManager Manager, sessionCloseInfo map[string][]*pb
 			result = append(result, fmt.Errorf("could not find session ID %q in local state after closing connections", sessionId))
 			continue
 		}
-		connStatus := make(map[string]pbs.CONNECTIONSTATUS, len(responses))
 		for _, response := range responses {
 			if err := si.ApplyLocalConnectionStatus(response.GetConnectionId(), response.GetStatus()); err != nil {
 				result = append(result, err)
 				continue
 			}
-			connStatus[response.GetConnectionId()] = response.GetStatus()
 			if response.GetStatus() == pbs.CONNECTIONSTATUS_CONNECTIONSTATUS_CLOSED {
 				closedIds = append(closedIds, response.GetConnectionId())
 			}

internal/daemon/worker/session/session_test.go

@@ -7,7 +7,6 @@ import (
 	"context"
 	"fmt"
 	"testing"
-	"time"
 
 	pbs "github.com/hashicorp/boundary/internal/gen/controller/servers/services"
 	"github.com/hashicorp/boundary/internal/session"
@@ -49,15 +48,12 @@ func TestSession_CancelAllLocalConnections(t *testing.T) {
 		"3": {
 			Id:                "3",
 			connCtxCancelFunc: cancelFn("3"),
-			CloseTime:         time.Now(),
 		},
 	}
 	sess := &sess{
 		connInfoMap: connInfo,
 	}
-	assert.ElementsMatch(t, sess.CancelAllLocalConnections(), []string{"1", "2"})
-	// We can call the cancel context multiple times, even if it was marked
-	// closed previously.
+	assert.ElementsMatch(t, sess.CancelAllLocalConnections(), []string{"1", "2", "3"})
 	assert.ElementsMatch(t, closedContextCalled, []string{"1", "2", "3"})
 }
 
@@ -74,12 +70,6 @@ func TestSession_CancelOpenLocalConnections(t *testing.T) {
 			connCtxCancelFunc: cancelFn("1"),
 			Status:            pbs.CONNECTIONSTATUS_CONNECTIONSTATUS_CLOSED,
 		},
-		"2": {
-			Id:                "2",
-			connCtxCancelFunc: cancelFn("2"),
-			Status:            pbs.CONNECTIONSTATUS_CONNECTIONSTATUS_CLOSED,
-			CloseTime:         time.Now(),
-		},
 		"3": {
 			Id:                "3",
 			connCtxCancelFunc: cancelFn("3"),
@@ -90,9 +80,7 @@ func TestSession_CancelOpenLocalConnections(t *testing.T) {
 		connInfoMap: connInfo,
 	}
 	assert.ElementsMatch(t, sess.CancelOpenLocalConnections(), []string{"1"})
-	// We call the cancel context multiple times, even if it was marked
-	// closed previously like connection 2 was (by setting CloseTime)
-	assert.ElementsMatch(t, closedContextCalled, []string{"1", "2"})
+	assert.ElementsMatch(t, closedContextCalled, []string{"1"})
 }
 
 func TestSession_RequestActivate(t *testing.T) {
@@ -274,6 +262,126 @@ func TestMakeFakeSessionCloseInfoEmpty(t *testing.T) {
 	)
 }
 
+// TestSession_ApplyLocalConnectionStatus_Closed verifies that marking a
+// connection closed cancels its proxy context and removes it from the map.
+func TestSession_ApplyLocalConnectionStatus_Closed(t *testing.T) {
+	cancelled := false
+	s := &sess{
+		sessionId: "sess1",
+		connInfoMap: map[string]*ConnInfo{
+			"conn1": {
+				Id:     "conn1",
+				Status: pbs.CONNECTIONSTATUS_CONNECTIONSTATUS_AUTHORIZED,
+				connCtxCancelFunc: func() {
+					cancelled = true
+				},
+			},
+		},
+	}
+
+	require.NoError(t, s.ApplyLocalConnectionStatus("conn1", pbs.CONNECTIONSTATUS_CONNECTIONSTATUS_CLOSED))
+	assert.True(t, cancelled, "cancel func should have been called")
+	_, exists := s.connInfoMap["conn1"]
+	assert.False(t, exists, "closed connection should be removed from connInfoMap")
+}
+
+// TestSession_ApplyLocalConnectionStatus_NilCancelFunc verifies that a nil
+// cancel func is handled safely and the entry is still deleted.
+func TestSession_ApplyLocalConnectionStatus_NilCancelFunc(t *testing.T) {
+	s := &sess{
+		sessionId: "sess1",
+		connInfoMap: map[string]*ConnInfo{
+			"conn1": {
+				Id:                "conn1",
+				Status:            pbs.CONNECTIONSTATUS_CONNECTIONSTATUS_AUTHORIZED,
+				connCtxCancelFunc: nil,
+			},
+		},
+	}
+
+	require.NoError(t, s.ApplyLocalConnectionStatus("conn1", pbs.CONNECTIONSTATUS_CONNECTIONSTATUS_CLOSED))
+	_, exists := s.connInfoMap["conn1"]
+	assert.False(t, exists, "closed connection should be removed even when cancel func is nil")
+}
+
+// TestSession_ApplyLocalConnectionStatus_NonClosed verifies that a non-closed
+// status update leaves the entry in the map.
+func TestSession_ApplyLocalConnectionStatus_NonClosed(t *testing.T) {
+	s := &sess{
+		sessionId: "sess1",
+		connInfoMap: map[string]*ConnInfo{
+			"conn1": {
+				Id:     "conn1",
+				Status: pbs.CONNECTIONSTATUS_CONNECTIONSTATUS_AUTHORIZED,
+			},
+		},
+	}
+
+	require.NoError(t, s.ApplyLocalConnectionStatus("conn1", pbs.CONNECTIONSTATUS_CONNECTIONSTATUS_CONNECTED))
+	conn, exists := s.connInfoMap["conn1"]
+	require.True(t, exists, "non-closed connection should remain in connInfoMap")
+	assert.Equal(t, pbs.CONNECTIONSTATUS_CONNECTIONSTATUS_CONNECTED, conn.Status)
+}
+
+// TestSession_ApplyLocalConnectionStatus_UnknownConnection verifies that an
+// error is returned for an unknown connection ID.
+func TestSession_ApplyLocalConnectionStatus_UnknownConnection(t *testing.T) {
+	s := &sess{
+		sessionId:   "sess1",
+		connInfoMap: make(map[string]*ConnInfo),
+	}
+	err := s.ApplyLocalConnectionStatus("nonexistent", pbs.CONNECTIONSTATUS_CONNECTIONSTATUS_CLOSED)
+	require.Error(t, err)
+	assert.Contains(t, err.Error(), "nonexistent")
+	assert.Contains(t, err.Error(), "sess1")
+}
+
+// TestSession_CancelAllLocalConnections_Empty verifies that an empty map
+// returns an empty slice.
+func TestSession_CancelAllLocalConnections_Empty(t *testing.T) {
+	s := &sess{connInfoMap: make(map[string]*ConnInfo)}
+	assert.Empty(t, s.CancelAllLocalConnections())
+}
+
+// TestSession_CancelOpenLocalConnections_Empty verifies that an empty map
+// returns an empty slice.
+func TestSession_CancelOpenLocalConnections_Empty(t *testing.T) {
+	s := &sess{connInfoMap: make(map[string]*ConnInfo)}
+	assert.Empty(t, s.CancelOpenLocalConnections())
+}
+
+// TestSession_CancelOpenLocalConnections_MixedStatuses verifies that only
+// CLOSED connections are cancelled across all possible connection statuses.
+func TestSession_CancelOpenLocalConnections_MixedStatuses(t *testing.T) {
+	var cancelled []string
+	cancelFn := func(id string) context.CancelFunc {
+		return func() { cancelled = append(cancelled, id) }
+	}
+	s := &sess{
+		connInfoMap: map[string]*ConnInfo{
+			"authorized": {
+				Id:                "authorized",
+				connCtxCancelFunc: cancelFn("authorized"),
+				Status:            pbs.CONNECTIONSTATUS_CONNECTIONSTATUS_AUTHORIZED,
+			},
+			"connected": {
+				Id:                "connected",
+				connCtxCancelFunc: cancelFn("connected"),
+				Status:            pbs.CONNECTIONSTATUS_CONNECTIONSTATUS_CONNECTED,
+			},
+			"closed": {
+				Id:                "closed",
+				connCtxCancelFunc: cancelFn("closed"),
+				Status:            pbs.CONNECTIONSTATUS_CONNECTIONSTATUS_CLOSED,
+			},
+		},
+	}
+
+	result := s.CancelOpenLocalConnections()
+	assert.ElementsMatch(t, []string{"closed"}, result)
+	assert.ElementsMatch(t, []string{"closed"}, cancelled)
+}
+
 func TestApplyConnectionCounterCallbacks(t *testing.T) {
 	s := &sess{connInfoMap: make(map[string]*ConnInfo)}
 

internal/daemon/worker/testing.go

@@ -113,9 +113,8 @@ type TestSessionInfo struct {
 // as a part of TestSessionInfo. See that struct for details about
 // the purpose of this data and how it's gathered.
 type TestConnectionInfo struct {
-	Id        string
-	Status    pbs.CONNECTIONSTATUS
-	CloseTime time.Time
+	Id     string
+	Status pbs.CONNECTIONSTATUS
 }
 
 // LookupSession returns session info from the worker's local session
@@ -135,9 +134,8 @@ func (tw *TestWorker) LookupSession(id string) (TestSessionInfo, bool) {
 	conns := make(map[string]TestConnectionInfo)
 	for _, conn := range sess.GetLocalConnections() {
 		conns[conn.Id] = TestConnectionInfo{
-			Id:        conn.Id,
-			Status:    conn.Status,
-			CloseTime: conn.CloseTime,
+			Id:     conn.Id,
+			Status: conn.Status,
 		}
 	}
 

internal/daemon/worker/testing_test.go

@@ -71,18 +71,14 @@ func TestTestWorkerLookupSession(t *testing.T) {
 		sessionManager: manager,
 	}
 
-	closeTime := s.GetLocalConnections()["one"].CloseTime
-	assert.NotZero(t, closeTime)
+	// Closed connections are pruned from the local map immediately; the
+	// session itself must still be visible but with an empty connection set.
+	_, exists := s.GetLocalConnections()["one"]
+	assert.False(t, exists, "closed connection should be pruned from local state")
 	expected := TestSessionInfo{
-		Id:     "foo",
-		Status: pbs.SESSIONSTATUS_SESSIONSTATUS_ACTIVE,
-		Connections: map[string]TestConnectionInfo{
-			"one": {
-				Id:        "one",
-				Status:    pbs.CONNECTIONSTATUS_CONNECTIONSTATUS_CLOSED,
-				CloseTime: closeTime,
-			},
-		},
+		Id:          "foo",
+		Status:      pbs.SESSIONSTATUS_SESSIONSTATUS_ACTIVE,
+		Connections: map[string]TestConnectionInfo{},
 	}
 
 	actual, ok := tw.LookupSession("foo")

internal/tests/helper/testing_helper.go

@@ -178,6 +178,8 @@ func (s *TestSession) ExpectConnectionStateOnController(
 
 // ExpectConnectionStateOnWorker waits until all connections in a
 // session have transitioned to a particular state on the worker.
+// Connections that have been pruned from the worker's local map
+// (i.e. confirmed closed) are treated as StatusClosed.
 func (s *TestSession) ExpectConnectionStateOnWorker(
 	ctx context.Context,
 	t *testing.T,
@@ -216,6 +218,21 @@ func (s *TestSession) ExpectConnectionStateOnWorker(
 			actualStates[conn.Id] = session.ConnectionStatusFromProtoVal(conn.Status)
 		}
 
+		// Closed connections are immediately pruned from the worker's local
+		// map. Treat any connection that has disappeared from the map as
+		// StatusClosed so that callers waiting for StatusClosed converge.
+		si, ok := tw.LookupSession(s.SessionId)
+		if ok {
+			for id, state := range actualStates {
+				if state == sessionStatusUnknown {
+					continue
+				}
+				if _, present := si.Connections[id]; !present {
+					actualStates[id] = session.StatusClosed
+				}
+			}
+		}
+
 		if reflect.DeepEqual(expectStates, actualStates) {
 			break
 		}
@@ -240,10 +257,8 @@ func (s *TestSession) testWorkerConnectionInfo(t *testing.T, tw *worker.TestWork
 	// local state.
 	require.True(ok)
 
-	// Likewise, we require the helper to always be used with
-	// connections.
-	require.Greater(len(si.Connections), 0, "should have at least one connection")
-
+	// Closed connections are immediately pruned from the local map, so
+	// the connections set may legitimately be empty. Do not require > 0.
 	return si.Connections
 }