session: add context to various session List methods

Author: ellemouton

session/interface.go

@@ -201,14 +201,15 @@ type Store interface {
 	GetSession(ctx context.Context, key *btcec.PublicKey) (*Session, error)
 
 	// ListAllSessions returns all sessions currently known to the store.
-	ListAllSessions() ([]*Session, error)
+	ListAllSessions(ctx context.Context) ([]*Session, error)
 
 	// ListSessionsByType returns all sessions of the given type.
-	ListSessionsByType(t Type) ([]*Session, error)
+	ListSessionsByType(ctx context.Context, t Type) ([]*Session, error)
 
 	// ListSessionsByState returns all sessions currently known to the store
 	// that are in the given states.
-	ListSessionsByState(...State) ([]*Session, error)
+	ListSessionsByState(ctx context.Context, state ...State) ([]*Session,
+		error)
 
 	// UpdateSessionRemotePubKey can be used to add the given remote pub key
 	// to the session with the given local pub key.

session/kvdb_store.go

@@ -352,7 +352,7 @@ func (db *BoltStore) GetSession(_ context.Context, key *btcec.PublicKey) (
 // ListAllSessions returns all sessions currently known to the store.
 //
 // NOTE: this is part of the Store interface.
-func (db *BoltStore) ListAllSessions() ([]*Session, error) {
+func (db *BoltStore) ListAllSessions(_ context.Context) ([]*Session, error) {
 	return db.listSessions(func(s *Session) bool {
 		return true
 	})
@@ -362,7 +362,9 @@ func (db *BoltStore) ListAllSessions() ([]*Session, error) {
 // have the given type.
 //
 // NOTE: this is part of the Store interface.
-func (db *BoltStore) ListSessionsByType(t Type) ([]*Session, error) {
+func (db *BoltStore) ListSessionsByType(_ context.Context, t Type) ([]*Session,
+	error) {
+
 	return db.listSessions(func(s *Session) bool {
 		return s.Type == t
 	})
@@ -372,7 +374,9 @@ func (db *BoltStore) ListSessionsByType(t Type) ([]*Session, error) {
 // are in the given states.
 //
 // NOTE: this is part of the Store interface.
-func (db *BoltStore) ListSessionsByState(states ...State) ([]*Session, error) {
+func (db *BoltStore) ListSessionsByState(_ context.Context, states ...State) (
+	[]*Session, error) {
+
 	return db.listSessions(func(s *Session) bool {
 		for _, state := range states {
 			if s.State == state {

session/store_test.go

@@ -58,18 +58,18 @@ func TestBasicSessionStore(t *testing.T) {
 	s3 := createSession(t, db, "session 3", withType(TypeAutopilot))
 
 	// Test the ListSessionsByType method.
-	sessions, err := db.ListSessionsByType(TypeMacaroonAdmin)
+	sessions, err := db.ListSessionsByType(ctx, TypeMacaroonAdmin)
 	require.NoError(t, err)
 	require.Equal(t, 2, len(sessions))
 	assertEqualSessions(t, s1, sessions[0])
 	assertEqualSessions(t, s2, sessions[1])
 
-	sessions, err = db.ListSessionsByType(TypeAutopilot)
+	sessions, err = db.ListSessionsByType(ctx, TypeAutopilot)
 	require.NoError(t, err)
 	require.Equal(t, 1, len(sessions))
 	assertEqualSessions(t, s3, sessions[0])
 
-	sessions, err = db.ListSessionsByType(TypeMacaroonReadonly)
+	sessions, err = db.ListSessionsByType(ctx, TypeMacaroonReadonly)
 	require.NoError(t, err)
 	require.Empty(t, sessions)
 
@@ -113,37 +113,37 @@ func TestBasicSessionStore(t *testing.T) {
 	require.Equal(t, s1.State, StateRevoked)
 
 	// Test that ListAllSessions works.
-	sessions, err = db.ListAllSessions()
+	sessions, err = db.ListAllSessions(ctx)
 	require.NoError(t, err)
 	require.Equal(t, 3, len(sessions))
 	assertEqualSessions(t, s1, sessions[0])
 	assertEqualSessions(t, s2, sessions[1])
 	assertEqualSessions(t, s3, sessions[2])
 
 	// Test that ListSessionsByState works.
-	sessions, err = db.ListSessionsByState(StateRevoked)
+	sessions, err = db.ListSessionsByState(ctx, StateRevoked)
 	require.NoError(t, err)
 	require.Equal(t, 1, len(sessions))
 	assertEqualSessions(t, s1, sessions[0])
 
-	sessions, err = db.ListSessionsByState(StateCreated)
+	sessions, err = db.ListSessionsByState(ctx, StateCreated)
 	require.NoError(t, err)
 	require.Equal(t, 2, len(sessions))
 	assertEqualSessions(t, s2, sessions[0])
 	assertEqualSessions(t, s3, sessions[1])
 
-	sessions, err = db.ListSessionsByState(StateCreated, StateRevoked)
+	sessions, err = db.ListSessionsByState(ctx, StateCreated, StateRevoked)
 	require.NoError(t, err)
 	require.Equal(t, 3, len(sessions))
 	assertEqualSessions(t, s1, sessions[0])
 	assertEqualSessions(t, s2, sessions[1])
 	assertEqualSessions(t, s3, sessions[2])
 
-	sessions, err = db.ListSessionsByState()
+	sessions, err = db.ListSessionsByState(ctx)
 	require.NoError(t, err)
 	require.Empty(t, sessions)
 
-	sessions, err = db.ListSessionsByState(StateReserved)
+	sessions, err = db.ListSessionsByState(ctx, StateReserved)
 	require.NoError(t, err)
 	require.Empty(t, sessions)
 
@@ -153,7 +153,7 @@ func TestBasicSessionStore(t *testing.T) {
 	// of the sessions are reserved.
 	require.NoError(t, db.DeleteReservedSessions())
 
-	sessions, err = db.ListSessionsByState(StateReserved)
+	sessions, err = db.ListSessionsByState(ctx, StateReserved)
 	require.NoError(t, err)
 	require.Empty(t, sessions)
 
@@ -163,7 +163,7 @@ func TestBasicSessionStore(t *testing.T) {
 	)
 	require.NoError(t, err)
 
-	sessions, err = db.ListSessionsByState(StateReserved)
+	sessions, err = db.ListSessionsByState(ctx, StateReserved)
 	require.NoError(t, err)
 	require.Equal(t, 1, len(sessions))
 	assertEqualSessions(t, s4, sessions[0])
@@ -182,7 +182,7 @@ func TestBasicSessionStore(t *testing.T) {
 	// database and no longer in the group ID/session ID index.
 	require.NoError(t, db.DeleteReservedSessions())
 
-	sessions, err = db.ListSessionsByState(StateReserved)
+	sessions, err = db.ListSessionsByState(ctx, StateReserved)
 	require.NoError(t, err)
 	require.Empty(t, sessions)
 

session_rpcserver.go

@@ -103,8 +103,7 @@ func (s *sessionRpcServer) start(ctx context.Context) error {
 
 	// Start up all previously created sessions.
 	sessions, err := s.cfg.db.ListSessionsByState(
-		session.StateCreated,
-		session.StateInUse,
+		ctx, session.StateCreated, session.StateInUse,
 	)
 	if err != nil {
 		return fmt.Errorf("error listing sessions: %v", err)
@@ -524,10 +523,10 @@ func (s *sessionRpcServer) resumeSession(ctx context.Context,
 }
 
 // ListSessions returns all sessions known to the session store.
-func (s *sessionRpcServer) ListSessions(_ context.Context,
+func (s *sessionRpcServer) ListSessions(ctx context.Context,
 	_ *litrpc.ListSessionsRequest) (*litrpc.ListSessionsResponse, error) {
 
-	sessions, err := s.cfg.db.ListAllSessions()
+	sessions, err := s.cfg.db.ListAllSessions(ctx)
 	if err != nil {
 		return nil, fmt.Errorf("error fetching sessions: %v", err)
 	}
@@ -1243,11 +1242,11 @@ func (s *sessionRpcServer) AddAutopilotSession(ctx context.Context,
 
 // ListAutopilotSessions fetches and returns all the sessions from the DB that
 // are of type TypeAutopilot.
-func (s *sessionRpcServer) ListAutopilotSessions(_ context.Context,
+func (s *sessionRpcServer) ListAutopilotSessions(ctx context.Context,
 	_ *litrpc.ListAutopilotSessionsRequest) (
 	*litrpc.ListAutopilotSessionsResponse, error) {
 
-	sessions, err := s.cfg.db.ListSessionsByType(session.TypeAutopilot)
+	sessions, err := s.cfg.db.ListSessionsByType(ctx, session.TypeAutopilot)
 	if err != nil {
 		return nil, fmt.Errorf("error fetching sessions: %v", err)
 	}