diff --git a/internal/daemon/worker/session/manager_test.go b/internal/daemon/worker/session/manager_test.go index e042f6bad..cae3e0e51 100644 --- a/internal/daemon/worker/session/manager_test.go +++ b/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 }) diff --git a/internal/daemon/worker/session/session.go b/internal/daemon/worker/session/session.go index 6c7bb7fee..20588eb26 100644 --- a/internal/daemon/worker/session/session.go +++ b/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()) } diff --git a/internal/daemon/worker/session/session_test.go b/internal/daemon/worker/session/session_test.go index baea5f8c7..0968846ff 100644 --- a/internal/daemon/worker/session/session_test.go +++ b/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)} diff --git a/internal/daemon/worker/testing.go b/internal/daemon/worker/testing.go index 4ccc34eb9..77f115fd7 100644 --- a/internal/daemon/worker/testing.go +++ b/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, } } diff --git a/internal/daemon/worker/testing_test.go b/internal/daemon/worker/testing_test.go index dd3594177..3992d640c 100644 --- a/internal/daemon/worker/testing_test.go +++ b/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") diff --git a/internal/tests/helper/testing_helper.go b/internal/tests/helper/testing_helper.go index e9a3e2ea6..17afac679 100644 --- a/internal/tests/helper/testing_helper.go +++ b/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 }