Move sorting to the store

xmtp · Oct 24, 2023 · 6b57691 · 6b57691
1 parent bd821d9
commit 6b57691
Show file tree

Hide file tree

Showing 2 changed files with 36 additions and 30 deletions.
diff --git a/pkg/api/message/v3/service.go b/pkg/api/message/v3/service.go
@@ -2,7 +2,6 @@ package api
 
 import (
 	"context"
-	"sort"
 
 	wakunode "github.com/waku-org/go-waku/waku/v2/node"
 	wakupb "github.com/waku-org/go-waku/waku/v2/protocol/pb"
@@ -212,7 +211,7 @@ func (s *Service) GetIdentityUpdates(ctx context.Context, req *proto.GetIdentity
 	resUpdates := make([]*proto.GetIdentityUpdatesResponse_WalletUpdates, len(walletAddresses))
 	for i, walletAddress := range walletAddresses {
 		walletUpdates := updates[walletAddress]
-		sort.Sort(walletUpdates)
+
 		resUpdates[i] = &proto.GetIdentityUpdatesResponse_WalletUpdates{
 			Updates: []*proto.GetIdentityUpdatesResponse_Update{},
 		}

diff --git a/pkg/mlsstore/store.go b/pkg/mlsstore/store.go
@@ -6,6 +6,7 @@ import (
 	"database/sql"
 	"encoding/hex"
 	"errors"
+	"sort"
 	"time"
 
 	"github.com/uptrace/bun"
@@ -120,36 +121,9 @@ func (s *Store) ConsumeKeyPackages(ctx context.Context, installationIds []string
 	return keyPackages, nil
 }
 
-type IdentityUpdateKind int
-
-const (
-	Create IdentityUpdateKind = iota
-	Revoke
-)
-
-type IdentityUpdate struct {
-	Kind           IdentityUpdateKind
-	InstallationId string
-	TimestampNs    uint64
-}
-
-// Add the required methods to make a valid sort.Sort interface
-type IdentityUpdateList []IdentityUpdate
-
-func (a IdentityUpdateList) Len() int {
-	return len(a)
-}
-
-func (a IdentityUpdateList) Swap(i, j int) {
-	a[i], a[j] = a[j], a[i]
-}
-
-func (a IdentityUpdateList) Less(i, j int) bool {
-	return a[i].TimestampNs < a[j].TimestampNs
-}
-
 func (s *Store) GetIdentityUpdates(ctx context.Context, walletAddresses []string, startTimeNs int64) (map[string]IdentityUpdateList, error) {
 	updated := make([]*Installation, 0)
+	// Find all installations that were changed since the startTimeNs
 	err := s.db.NewSelect().
 		Model(&updated).
 		Where("wallet_address IN (?)", bun.In(walletAddresses)).
@@ -163,6 +137,7 @@ func (s *Store) GetIdentityUpdates(ctx context.Context, walletAddresses []string
 		return nil, err
 	}
 
+	// The returned list is only partially sorted
 	out := make(map[string]IdentityUpdateList)
 	for _, installation := range updated {
 		if installation.CreatedAt > startTimeNs {
@@ -180,6 +155,10 @@ func (s *Store) GetIdentityUpdates(ctx context.Context, walletAddresses []string
 			})
 		}
 	}
+	// Sort the updates by timestamp now that the full list is assembled
+	for _, updates := range out {
+		sort.Sort(updates)
+	}
 
 	return out, nil
 }
@@ -241,3 +220,31 @@ func buildKeyPackageId(keyPackageData []byte) string {
 	digest := sha256.Sum256(keyPackageData)
 	return hex.EncodeToString(digest[:])
 }
+
+type IdentityUpdateKind int
+
+const (
+	Create IdentityUpdateKind = iota
+	Revoke
+)
+
+type IdentityUpdate struct {
+	Kind           IdentityUpdateKind
+	InstallationId string
+	TimestampNs    uint64
+}
+
+// Add the required methods to make a valid sort.Sort interface
+type IdentityUpdateList []IdentityUpdate
+
+func (a IdentityUpdateList) Len() int {
+	return len(a)
+}
+
+func (a IdentityUpdateList) Swap(i, j int) {
+	a[i], a[j] = a[j], a[i]
+}
+
+func (a IdentityUpdateList) Less(i, j int) bool {
+	return a[i].TimestampNs < a[j].TimestampNs
+}