Add minor PR changes

internal/db/errors.go

@@ -101,6 +101,7 @@ const (
 	errReplicatorCollections                    string = "failed to get collections for replicator"
 	errReplicatorNotFound                       string = "replicator not found"
 	errCanNotEncryptBuiltinField                string = "can not encrypt build-in field"
+	errFailedToHandleEncKeysReceivedEvent       string = "failed to handle encryption-keys-received event"
 )
 
 var (

internal/db/messages.go

@@ -82,30 +82,25 @@ func (db *db) handleMessages(ctx context.Context, sub *event.Subscription) {
 
 			case encryption.KeyRetrievedEvent:
 				go func() {
-					ctx = encryption.ContextWithStore(ctx, db.Encstore())
-					for encStoreKey, encKey := range evt.Keys {
-						err := encryption.SaveKey(ctx, encStoreKey, encKey)
-
-						if err != nil {
-							log.ErrorContextE(
-								ctx,
-								"Failed to save doc encryption key",
-								err,
-								corelog.Any("Event", evt))
-						}
-					}
-
-					err := db.mergeEncryptedBlocks(ctx, evt)
-
-					if err != nil {
-						log.ErrorContextE(
-							ctx,
-							"Failed to merge encrypted block",
-							err,
-							corelog.Any("Event", evt))
+					if err := db.handleEncryptionKeysRetrievedEvent(ctx, evt); err != nil {
+						log.ErrorContextE(ctx, errFailedToHandleEncKeysReceivedEvent, err, corelog.Any("Event", evt))
 					}
 				}()
 			}
 		}
 	}
 }
+
+// handleEncryptionKeysRetrievedEvent handles the event when requested encryption keys are retrieved from other peers.
+func (db *db) handleEncryptionKeysRetrievedEvent(ctx context.Context, evt encryption.KeyRetrievedEvent) error {
+	ctx = encryption.ContextWithStore(ctx, db.Encstore())
+	for encStoreKey, encKey := range evt.Keys {
+		err := encryption.SaveKey(ctx, encStoreKey, encKey)
+
+		if err != nil {
+			return err
+		}
+	}
+
+	return db.mergeEncryptedBlocks(ctx, evt)
+}

net/server.go

@@ -98,7 +98,9 @@ func newServer(p *Peer, opts ...grpc.DialOption) (*server, error) {
 	return s, nil
 }
 
-func (s *server) extractSessionAndRemoveOldOnes(id string) *session {
+// extractSessionAndRemoveOld extracts a session with the given id from the server's session list
+// and removes any old sessions by comparing their timestamps.
+func (s *server) extractSessionAndRemoveOld(id string) *session {
 	var result *session
 	swapLast := func(i int) {
 		s.sessions[i] = s.sessions[len(s.sessions)-1]

net/server_encryption_key.go

@@ -292,7 +292,7 @@ func (s *server) handleFetchEncryptionKeyResponse(resp rpc.Response, req *pb.Fet
 		return
 	}
 
-	session := s.extractSessionAndRemoveOldOnes(string(keyResp.ReqEphemeralPublicKey))
+	session := s.extractSessionAndRemoveOld(string(keyResp.ReqEphemeralPublicKey))
 	if session == nil {
 		log.ErrorContext(s.peer.ctx, "Failed to find session for ephemeral public key")
 		return

tests/integration/assert_stack.go

@@ -15,6 +15,10 @@ import (
 	"strings"
 )
 
+// assertStack keeps track of the current assertion path.
+// GraphQL response can be traversed by a key of a map and/or an index of an array.
+// So whenever we have a mismatch in a large response, we can use this stack to find the exact path.
+// Example output: "commits[2].links[1].cid"
 type assertStack struct {
 	stack []string
 	isMap []bool

tests/integration/events.go

@@ -351,9 +351,13 @@ func getEventsForCreateDoc(s *state, action CreateDoc) map[string]struct{} {
 	return expect
 }
 
-func waitForKeyRetrievedEvent(s *state, nodeIDs []int) {
+// waitForKeyRetrievedEvent waits for nodes to receive a key retrieved event.
+// If targetNodeIDs is empty, all nodes will be waiting on the event sync.
+// Otherwise, only the nodes with the specified IDs will be checked.
+func waitForKeyRetrievedEvent(s *state, targetNodeIDs []int) {
 	for nodeID := 0; nodeID < len(s.nodes); nodeID++ {
-		if len(nodeIDs) > 0 && !slices.Contains(nodeIDs, nodeID) {
+		// if we have target nodes and the current node is not in the list, skip it
+		if len(targetNodeIDs) > 0 && !slices.Contains(targetNodeIDs, nodeID) {
 			continue
 		}