Merge pull request #15 from bluesky-social/jc/zset

Subset of Commands for Ordered Sets

Author: jcalabro

.gitignore

@@ -30,3 +30,6 @@ go.work.sum
 # Editor/IDE
 # .idea/
 # .vscode/
+
+# Delve build artifacts
+**__debug_bin**

README.md

@@ -68,3 +68,4 @@ FoundationDB provides strictly serializable transactions atop an ordered key-val
 |`redis_v0/<user_id>/obj/<obj_id>`|Storage of data objects|
 |`redis_v0/<user_id>/list/<list_id>`|Metadata protobuf for information on each list|
 |`redis_v0/<user_id>/list/<list_id>/<item_id>`|Metadata protobuf for information on each item in a list|
+|`redis_v0/<user_id>/zset/<set_id>`|Stores an ordered list of item scores for ordered sets|

internal/server/redis/auth.go

@@ -117,13 +117,19 @@ func (s *session) setSessionUser(user *types.User) error {
 		return fmt.Errorf("failed to initialize user list directory: %w", err)
 	}
 
+	zsetDir, err := userDir.CreateOrOpen(s.fdb, []string{"zset"}, nil)
+	if err != nil {
+		return fmt.Errorf("failed to initialize user zset directory: %w", err)
+	}
+
 	s.userMu.Lock()
-	s.user = &sessionUser{
+	s.user = &userSession{
 		objDir:        objDir,
 		metaDir:       metaDir,
 		uidDir:        uidDir,
 		reverseUIDDir: reverseUIDDir,
 		listDir:       listDir,
+		zsetDir:       zsetDir,
 		user:          user,
 	}
 	s.userMu.Unlock()

internal/server/redis/list.go

@@ -97,7 +97,7 @@ func (s *session) handlePush(ctx context.Context, args []resp.Value, left bool)
 	}
 
 	_, err = s.fdb.Transact(func(tx fdb.Transaction) (any, error) {
-		listMetaKey, objMeta, err := s.getObjectMeta(ctx, tx, key)
+		listMetaKey, objMeta, err := s.getMeta(ctx, tx, key)
 		if err != nil {
 			return nil, fmt.Errorf("failed to get list meta: %w", err)
 		}

internal/server/redis/object.go

@@ -23,6 +23,7 @@ type objectKind int64
 const (
 	objectKindBasic objectKind = 1 << iota
 	objectKindSet
+	objectKindSortedSet
 	objectKindList
 	objectKindListItem
 )
@@ -44,6 +45,9 @@ func getNumChunks(meta *types.ObjectMeta, kind gt.Option[objectKind]) (uint32, e
 	case *types.ObjectMeta_Set:
 		objKind = objectKindSet
 		numChunks = typ.Set.NumChunks
+	case *types.ObjectMeta_SortedSet:
+		objKind = objectKindSortedSet
+		numChunks = typ.SortedSet.NumChunks
 	case *types.ObjectMeta_ListItem:
 		objKind = objectKindListItem
 		numChunks = typ.ListItem.NumChunks
@@ -64,7 +68,7 @@ func (s *session) getObject(ctx context.Context, tx fdb.ReadTransaction, kind ob
 	ctx, span := s.tracer.Start(ctx, "getObject")
 	defer span.End()
 
-	_, meta, err := s.getObjectMeta(ctx, tx, id)
+	_, meta, err := s.getMeta(ctx, tx, id)
 	if err != nil {
 		span.RecordError(err)
 		return nil, nil, err
@@ -116,7 +120,7 @@ func (s *session) writeObject(ctx context.Context, tx fdb.Transaction, id string
 	now := timestamppb.Now()
 
 	// check if the object already exists and should be overwritten
-	metaKey, meta, err := s.getObjectMeta(ctx, tx, id)
+	metaKey, meta, err := s.getMeta(ctx, tx, id)
 	if err != nil {
 		span.RecordError(err)
 		return err
@@ -137,6 +141,8 @@ func (s *session) writeObject(ctx context.Context, tx fdb.Transaction, id string
 			meta.Type = &types.ObjectMeta_Basic{Basic: &types.BasicObjectMeta{}}
 		case objectKindSet:
 			meta.Type = &types.ObjectMeta_Set{Set: &types.SetMeta{}}
+		case objectKindSortedSet:
+			meta.Type = &types.ObjectMeta_SortedSet{SortedSet: &types.SortedSetMeta{}}
 		case objectKindListItem:
 			meta.Type = &types.ObjectMeta_ListItem{ListItem: &types.ListItemMeta{}}
 		}
@@ -157,6 +163,8 @@ func (s *session) writeObject(ctx context.Context, tx fdb.Transaction, id string
 		typ.Basic.NumChunks = numChunks
 	case *types.ObjectMeta_Set:
 		typ.Set.NumChunks = numChunks
+	case *types.ObjectMeta_SortedSet:
+		typ.SortedSet.NumChunks = numChunks
 	case *types.ObjectMeta_ListItem:
 		typ.ListItem.NumChunks = numChunks
 	}
@@ -217,10 +225,7 @@ func (s *session) deleteObject(ctx context.Context, tx fdb.Transaction, id strin
 		span.RecordError(err)
 		return fmt.Errorf("failed to get object start key: %w", err)
 	}
-	tx.ClearRange(fdb.KeyRange{
-		Begin: begin,
-		End:   end,
-	})
+	tx.ClearRange(fdb.KeyRange{Begin: begin, End: end})
 
 	metrics.SpanOK(span)
 	return nil
@@ -275,7 +280,7 @@ func (s *session) handleExists(ctx context.Context, args []resp.Value) (string,
 	}
 
 	existsAny, err := s.fdb.ReadTransact(func(tx fdb.ReadTransaction) (any, error) {
-		_, meta, err := s.getObjectMeta(ctx, tx, key)
+		_, meta, err := s.getMeta(ctx, tx, key)
 		if err != nil {
 			return nil, err
 		}
@@ -337,7 +342,7 @@ func (s *session) handleDelete(ctx context.Context, args []resp.Value) (string,
 	}
 
 	existsAny, err := s.fdb.Transact(func(tx fdb.Transaction) (any, error) {
-		_, meta, err := s.getObjectMeta(ctx, tx, key)
+		_, meta, err := s.getMeta(ctx, tx, key)
 		if err != nil {
 			return false, err
 		}
@@ -350,6 +355,19 @@ func (s *session) handleDelete(ctx context.Context, args []resp.Value) (string,
 			return false, err
 		}
 
+		// if we're deleting a sorted set, also clear its score directory
+		if _, ok := meta.Type.(*types.ObjectMeta_SortedSet); ok {
+			scoreDir, err := s.sortedSetScoreDir(key)
+			if err != nil {
+				return false, fmt.Errorf("failed to get sorted set score directory: %w", err)
+			}
+
+			_, err = scoreDir.Remove(tx, []string{})
+			if err != nil {
+				return false, fmt.Errorf("failed to remove score directory: %w", err)
+			}
+		}
+
 		return true, nil
 	})
 	if err != nil {
@@ -578,7 +596,7 @@ func (s *session) handleExpire(ctx context.Context, args []resp.Value) (string,
 	delta := time.Duration(secs) * time.Second
 
 	resAny, err := s.fdb.Transact(func(tx fdb.Transaction) (any, error) {
-		metaKey, meta, err := s.getObjectMeta(ctx, tx, key)
+		metaKey, meta, err := s.getMeta(ctx, tx, key)
 		if err != nil {
 			return 0, fmt.Errorf("failed to get object meta: %w", err)
 		}

internal/server/redis/redis.go

@@ -151,16 +151,17 @@ type session struct {
 	dirs *Directories
 
 	// Only set once the user has authenticated
-	user   *sessionUser
+	user   *userSession
 	userMu *sync.RWMutex
 }
 
-type sessionUser struct {
+type userSession struct {
 	objDir        directory.DirectorySubspace
 	metaDir       directory.DirectorySubspace
 	uidDir        directory.DirectorySubspace
 	reverseUIDDir directory.DirectorySubspace
 	listDir       directory.DirectorySubspace
+	zsetDir       directory.DirectorySubspace
 
 	user *types.User
 }
@@ -247,6 +248,18 @@ func (s *session) uidKey(id string) (fdb.Key, error) {
 	return s.user.uidDir.Pack(tuple.Tuple{id}), nil
 }
 
+// Returns the FDB directory
+func (s *session) sortedSetScoreDir(key string) (directory.DirectorySubspace, error) {
+	s.userMu.RLock()
+	defer s.userMu.RUnlock()
+
+	if s.user == nil {
+		return nil, fmt.Errorf("authentication is required")
+	}
+
+	return s.user.zsetDir.CreateOrOpen(s.fdb, []string{key}, nil)
+}
+
 // Returns the FDB key of an object in the per-user uid directory
 func (s *session) reverseUIDKey(id string) (fdb.Key, error) {
 	s.userMu.RLock()
@@ -382,6 +395,14 @@ func (s *session) handleCommand(ctx context.Context, cmd *resp.Command) string {
 		res, err = s.handleSetUnion(ctx, cmd.Args)
 	case "sdiff":
 		res, err = s.handleSetDiff(ctx, cmd.Args)
+	case "zadd":
+		res, err = s.handleZAdd(ctx, cmd.Args)
+	case "zcount":
+		res, err = s.handleZCount(ctx, cmd.Args)
+	case "zremrangebyscore":
+		res, err = s.handleZRemRangeByScore(ctx, cmd.Args)
+	case "zcard":
+		res, err = s.handleSetCard(ctx, cmd.Args)
 	case "llen":
 		res, err = s.handleLLen(ctx, cmd.Args)
 	case "lpush":
@@ -552,61 +573,43 @@ func userIsAdmin(user *types.User) bool {
 	return false
 }
 
-func (s *session) getMeta(ctx context.Context, tx fdb.ReadTransaction, id string, meta proto.Message) (bool, fdb.Key, error) {
+func (s *session) getMeta(ctx context.Context, tx fdb.ReadTransaction, id string) (fdb.Key, *types.ObjectMeta, error) {
 	ctx, span := s.tracer.Start(ctx, "getMeta") // nolint
 	defer span.End()
 
 	metaKey, err := s.metaKey(id)
 	if err != nil {
 		span.RecordError(err)
-		return false, nil, fmt.Errorf("failed to get meta key: %w", err)
+		return nil, nil, fmt.Errorf("failed to get meta key: %w", err)
 	}
 
 	metaBuf, err := tx.Get(metaKey).Get()
 	if err != nil {
 		span.RecordError(err)
-		return false, metaKey, fmt.Errorf("failed to get object meta: %w", err)
+		return nil, nil, fmt.Errorf("failed to get object meta: %w", err)
 	}
 
 	if len(metaBuf) == 0 {
 		metrics.SpanOK(span)
-		return false, metaKey, nil
-	}
-
-	if err = proto.Unmarshal(metaBuf, meta); err != nil {
-		span.RecordError(err)
-		return false, metaKey, fmt.Errorf("failed to proto unmarshal object meta: %w", err)
+		return metaKey, nil, nil
 	}
 
-	metrics.SpanOK(span)
-	return true, metaKey, nil
-}
-
-// Returns the associated metadata for an object of any type. Returns nil if it not exist.
-func (s *session) getObjectMeta(ctx context.Context, tx fdb.ReadTransaction, id string) (fdb.Key, *types.ObjectMeta, error) {
-	ctx, span := s.tracer.Start(ctx, "getObjectMeta")
-	defer span.End()
-
 	meta := &types.ObjectMeta{}
-	exists, key, err := s.getMeta(ctx, tx, id, meta)
-	if err != nil {
+	if err = proto.Unmarshal(metaBuf, meta); err != nil {
 		span.RecordError(err)
-		return nil, nil, err
-	}
-	if !exists {
-		meta = nil
+		return nil, nil, fmt.Errorf("failed to proto unmarshal object meta: %w", err)
 	}
 
 	metrics.SpanOK(span)
-	return key, meta, nil
+	return metaKey, meta, nil
 }
 
 // Returns the associated metadata for an object. Returns nil if the object does not exist.
 func (s *session) getListMeta(ctx context.Context, tx fdb.ReadTransaction, id string) (fdb.Key, *types.ListMeta, error) {
 	ctx, span := s.tracer.Start(ctx, "getListMeta")
 	defer span.End()
 
-	key, objMeta, err := s.getObjectMeta(ctx, tx, id)
+	key, objMeta, err := s.getMeta(ctx, tx, id)
 	if err != nil {
 		span.RecordError(err)
 		return nil, nil, fmt.Errorf("failed to get object meta: %w", err)
@@ -698,12 +701,13 @@ func (s *session) allocateNewUID(ctx context.Context, tx fdb.Transaction) (uint6
 	return newUID, nil
 }
 
-// Returns the UID for the given member string, creating a new one if it does not exist
-func (s *session) getOrAllocateUID(ctx context.Context, tx fdb.Transaction, member string) (uint64, error) {
+// Returns the UID for the given member string, creating a new one if it does not exist. If peek is true, it
+// will refuse to create a new UID if one does not already exist.
+func (s *session) getOrAllocateUID(ctx context.Context, tx fdb.Transaction, member *types.SetMember) (uint64, error) {
 	ctx, span := s.tracer.Start(ctx, "getOrAllocateUID")
 	defer span.End()
 
-	memberToUIDKey, err := s.reverseUIDKey(member)
+	memberToUIDKey, err := s.reverseUIDKey(member.Member)
 	if err != nil {
 		span.RecordError(err)
 		return 0, fmt.Errorf("failed to get uid key: %w", err)
@@ -718,13 +722,13 @@ func (s *session) getOrAllocateUID(ctx context.Context, tx fdb.Transaction, memb
 
 	if len(val) == 0 {
 		// allocate a new UID for this member string
-		uid, err := s.allocateNewUID(ctx, tx)
+		member.Uid, err = s.allocateNewUID(ctx, tx)
 		if err != nil {
 			span.RecordError(err)
 			return 0, fmt.Errorf("failed to allocate new uid: %w", err)
 		}
 
-		uidStr := strconv.FormatUint(uid, 10)
+		uidStr := strconv.FormatUint(member.Uid, 10)
 		uidToMemberKey, err := s.uidKey(uidStr)
 		if err != nil {
 			span.RecordError(err)
@@ -733,7 +737,9 @@ func (s *session) getOrAllocateUID(ctx context.Context, tx fdb.Transaction, memb
 
 		// store the bi-directional mapping
 		tx.Set(memberToUIDKey, []byte(uidStr))
-		tx.Set(uidToMemberKey, []byte(member))
+		if err := setProtoItem(tx, uidToMemberKey, member); err != nil {
+			return 0, fmt.Errorf("failed to set uid to member: %w", err)
+		}
 
 		val = []byte(uidStr)
 	}
@@ -780,30 +786,29 @@ func (s *session) peekUID(ctx context.Context, tx fdb.ReadTransaction, member st
 	return uid, nil
 }
 
-func (s *session) memberFromUID(ctx context.Context, tx fdb.ReadTransaction, uid uint64) (string, error) {
+func (s *session) memberFromUID(ctx context.Context, tx fdb.ReadTransaction, uid uint64) (*types.SetMember, error) {
 	ctx, span := s.tracer.Start(ctx, "memberFromUID") // nolint
 	defer span.End()
 
 	key, err := s.uidKey(strconv.FormatUint(uid, 10))
 	if err != nil {
 		span.RecordError(err)
-		return "", fmt.Errorf("failed to get uid key: %w", err)
+		return nil, fmt.Errorf("failed to get uid key: %w", err)
 	}
 
-	member, err := tx.Get(key).Get()
+	member := &types.SetMember{}
+	exists, err := getProtoItem(tx, key, member)
 	if err != nil {
-		span.RecordError(err)
-		return "", fmt.Errorf("failed to get member for UID %d: %w", uid, err)
+		return nil, err
 	}
-
-	if len(member) == 0 {
+	if !exists {
 		err := fmt.Errorf("uid %d not found", uid)
 		span.RecordError(err)
-		return "", err
+		return nil, err
 	}
 
 	metrics.SpanOK(span)
-	return string(member), nil
+	return member, nil
 }
 
 func parseVariadicArguments(args []resp.Value) ([]string, error) {