feat: add atomic room allocation API with TTL-based reservations (#684)

Author: fbgava-wildlife

cmd/roomsapi/rooms_api.go

@@ -161,6 +161,8 @@ func buildMuxWithMetricsMdlw(mdlw middleware.Middleware, mux *runtime.ServeMux)
 			handler = std.Handler("/scheduler/:schedulerName/rooms/:roomID/status", mdlw, mux)
 		case commom.MatchPath(path, fmt.Sprintf("^/scheduler/%s/rooms/%s/address", anyWordRegex, anyWordRegex)):
 			handler = std.Handler("/scheduler/:schedulerName/rooms/:roomID/address", mdlw, mux)
+		case commom.MatchPath(path, fmt.Sprintf("^/scheduler/%s/rooms/allocate$", anyWordRegex)):
+			handler = std.Handler("/scheduler/:schedulerName/rooms/allocate", mdlw, mux)
 		default:
 			handler = std.Handler("", mdlw, mux)
 		}
@@ -190,6 +192,8 @@ func buildMuxWithTracing(hndl http.Handler) http.Handler {
 			handler = otelhttp.NewHandler(hndl, "/scheduler/:schedulerName/rooms/:roomID/status")
 		case commom.MatchPath(path, fmt.Sprintf("^/scheduler/%s/rooms/%s/address", anyWordRegex, anyWordRegex)):
 			handler = otelhttp.NewHandler(hndl, "/scheduler/:schedulerName/rooms/:roomID/address")
+		case commom.MatchPath(path, fmt.Sprintf("^/scheduler/%s/rooms/allocate$", anyWordRegex)):
+			handler = otelhttp.NewHandler(hndl, "/scheduler/:schedulerName/rooms/allocate")
 		default:
 			handler = otelhttp.NewHandler(hndl, "")
 		}

config/config.yaml

@@ -91,6 +91,7 @@ services:
     roomPingTimeoutMillis: 240000
     roomInitializationTimeoutMillis: 120000
     roomDeletionTimeoutMillis: 120000
+    roomAllocationTTLMillis: 300000
     roomValidationAttempts: 3
   operationManager:
     operationLeaseTTLMillis: 5000

internal/adapters/storage/redis/room/redis.go

@@ -50,8 +50,38 @@ const (
 	versionKey          = "version"
 	isValidationRoomKey = "is_validation_room"
 	createdAtKey        = "created_at"
+	allocatedAtKey      = "allocated_at"
 )
 
+// allocateRoomScript atomically selects a ready room and marks it as occupied
+var allocateRoomScript = redis.NewScript(`
+local status_key = KEYS[1]
+local room_key_prefix = KEYS[2]
+local ready_status = 2  -- GameStatusReady
+local allocated_status = 8  -- GameStatusAllocated
+local current_time = ARGV[1]
+
+-- Get one ready room (lowest score first, which is actually by room ID since all ready rooms have same score)
+local ready_rooms = redis.call('ZRANGEBYSCORE', status_key, ready_status, ready_status, 'LIMIT', 0, 1)
+if #ready_rooms == 0 then
+    return "NO_ROOMS_AVAILABLE"
+end
+
+local room_id = ready_rooms[1]
+
+-- Atomically change the room status from ready to allocated using ZAddXXCh
+-- This only updates if the member already exists and returns 1 if score was changed
+local changed = redis.call('ZADD', status_key, 'XX', 'CH', allocated_status, room_id)
+if changed == 1 then
+    -- Set the AllocatedAt timestamp
+    local room_key = room_key_prefix .. room_id
+    redis.call('HSET', room_key, 'allocated_at', current_time)
+    return room_id
+else
+    return "ALLOCATION_FAILED"
+end
+`)
+
 type redisStateStorage struct {
 	client *redis.Client
 }
@@ -112,6 +142,17 @@ func (r redisStateStorage) GetRoom(ctx context.Context, scheduler, roomID string
 	}
 
 	room.CreatedAt = time.Unix(createdAt, 0)
+
+	// Parse AllocatedAt if it exists
+	if allocatedAtStr := roomHashCmd.Val()[allocatedAtKey]; allocatedAtStr != "" {
+		allocatedAtUnix, err := strconv.ParseInt(allocatedAtStr, 10, 64)
+		if err != nil {
+			return nil, errors.NewErrEncoding("error parsing room %s allocatedAtKey", roomID).WithError(err)
+		}
+		allocatedAt := time.Unix(allocatedAtUnix, 0)
+		room.AllocatedAt = &allocatedAt
+	}
+
 	return room, nil
 }
 
@@ -440,3 +481,35 @@ func getRoomStatusUpdateChannel(scheduler, roomID string) string {
 func getRoomOccupancyRedisKey(scheduler string) string {
 	return fmt.Sprintf("scheduler:%s:occupancy", scheduler)
 }
+
+func (r *redisStateStorage) AllocateRoom(ctx context.Context, schedulerName string) (string, error) {
+	statusKey := getRoomStatusSetRedisKey(schedulerName)
+	roomKeyPrefix := fmt.Sprintf("scheduler:%s:rooms:", schedulerName)
+	currentTime := time.Now().Unix()
+
+	var result interface{}
+	var err error
+	metrics.RunWithMetrics(roomStorageMetricLabel, func() error {
+		result, err = allocateRoomScript.Run(ctx, r.client, []string{statusKey, roomKeyPrefix}, currentTime).Result()
+		return err
+	})
+
+	if err != nil {
+		return "", errors.NewErrUnexpected("error executing room allocation script for scheduler %s", schedulerName).WithError(err)
+	}
+
+	roomID, ok := result.(string)
+	if !ok || roomID == "" {
+		return "", errors.NewErrUnexpected("unexpected result %v from room allocation script for scheduler %s", result, schedulerName)
+	}
+
+	// Return errors for descriptive error strings from Lua script
+	if roomID == "NO_ROOMS_AVAILABLE" {
+		return "", errors.NewErrUnexpected("no ready rooms available for scheduler %s", schedulerName)
+	}
+	if roomID == "ALLOCATION_FAILED" {
+		return "", errors.NewErrUnexpected("room allocation failed for scheduler %s", schedulerName)
+	}
+
+	return roomID, nil
+}

internal/adapters/storage/redis/room/redis_test.go

@@ -782,3 +782,151 @@ func TestRedisStateStorage_GetRunningMatchesCount(t *testing.T) {
 		require.Equal(t, 6, count)
 	})
 }
+
+func TestRedisStateStorage_AllocateRoom(t *testing.T) {
+	ctx := context.Background()
+
+	t.Run("successfully allocates a ready room", func(t *testing.T) {
+		client := test.GetRedisConnection(t, redisAddress)
+		storage := NewRedisStateStorage(client)
+
+		room := &game_room.GameRoom{
+			ID:          "room-1",
+			SchedulerID: "game",
+			Version:     "1.0",
+			Status:      game_room.GameStatusReady,
+			LastPingAt:  time.Unix(1, 0),
+			Metadata:    map[string]interface{}{"host": "localhost"},
+		}
+
+		require.NoError(t, storage.CreateRoom(ctx, room))
+
+		allocatedRoomID, err := storage.AllocateRoom(ctx, "game")
+		require.NoError(t, err)
+		require.Equal(t, "room-1", allocatedRoomID)
+
+		// Verify the room status was changed to allocated
+		updatedRoom, err := storage.GetRoom(ctx, "game", "room-1")
+		require.NoError(t, err)
+		require.Equal(t, game_room.GameStatusAllocated, updatedRoom.Status)
+
+		// Verify the AllocatedAt timestamp was set
+		require.NotNil(t, updatedRoom.AllocatedAt, "AllocatedAt should be set")
+		require.True(t, time.Since(*updatedRoom.AllocatedAt) < 5*time.Second, "AllocatedAt should be recent")
+	})
+
+	t.Run("returns error when no ready rooms available", func(t *testing.T) {
+		client := test.GetRedisConnection(t, redisAddress)
+		storage := NewRedisStateStorage(client)
+
+		room := &game_room.GameRoom{
+			ID:          "room-1",
+			SchedulerID: "game",
+			Version:     "1.0",
+			Status:      game_room.GameStatusOccupied,
+			LastPingAt:  time.Unix(1, 0),
+			Metadata:    map[string]interface{}{"host": "localhost"},
+		}
+
+		require.NoError(t, storage.CreateRoom(ctx, room))
+
+		allocatedRoomID, err := storage.AllocateRoom(ctx, "game")
+		require.Error(t, err)
+		require.Equal(t, "", allocatedRoomID)
+		require.Contains(t, err.Error(), "no ready rooms available")
+	})
+
+	t.Run("returns error when scheduler has no rooms", func(t *testing.T) {
+		client := test.GetRedisConnection(t, redisAddress)
+		storage := NewRedisStateStorage(client)
+
+		allocatedRoomID, err := storage.AllocateRoom(ctx, "nonexistent")
+		require.Error(t, err)
+		require.Equal(t, "", allocatedRoomID)
+		require.Contains(t, err.Error(), "no ready rooms available")
+	})
+
+	t.Run("allocates only one room when multiple ready rooms exist", func(t *testing.T) {
+		client := test.GetRedisConnection(t, redisAddress)
+		storage := NewRedisStateStorage(client)
+
+		rooms := []*game_room.GameRoom{
+			{
+				ID:          "room-1",
+				SchedulerID: "game",
+				Version:     "1.0",
+				Status:      game_room.GameStatusReady,
+				LastPingAt:  time.Unix(1, 0),
+				Metadata:    map[string]interface{}{"host": "localhost"},
+			},
+			{
+				ID:          "room-2",
+				SchedulerID: "game",
+				Version:     "1.0",
+				Status:      game_room.GameStatusReady,
+				LastPingAt:  time.Unix(2, 0),
+				Metadata:    map[string]interface{}{"host": "localhost"},
+			},
+		}
+
+		for _, room := range rooms {
+			require.NoError(t, storage.CreateRoom(ctx, room))
+		}
+
+		allocatedRoomID, err := storage.AllocateRoom(ctx, "game")
+		require.NoError(t, err)
+		require.NotEmpty(t, allocatedRoomID)
+		require.Contains(t, []string{"room-1", "room-2"}, allocatedRoomID)
+
+		// Verify only the allocated room was changed to allocated
+		readyRooms, err := storage.GetRoomIDsByStatus(ctx, "game", game_room.GameStatusReady)
+		require.NoError(t, err)
+		require.Len(t, readyRooms, 1)
+
+		allocatedRooms, err := storage.GetRoomIDsByStatus(ctx, "game", game_room.GameStatusAllocated)
+		require.NoError(t, err)
+		require.Len(t, allocatedRooms, 1)
+		require.Equal(t, allocatedRoomID, allocatedRooms[0])
+	})
+
+	t.Run("ignores non-ready rooms", func(t *testing.T) {
+		client := test.GetRedisConnection(t, redisAddress)
+		storage := NewRedisStateStorage(client)
+
+		rooms := []*game_room.GameRoom{
+			{
+				ID:          "room-pending",
+				SchedulerID: "game",
+				Version:     "1.0",
+				Status:      game_room.GameStatusPending,
+				LastPingAt:  time.Unix(1, 0),
+				Metadata:    map[string]interface{}{"host": "localhost"},
+			},
+			{
+				ID:          "room-occupied",
+				SchedulerID: "game",
+				Version:     "1.0",
+				Status:      game_room.GameStatusOccupied,
+				LastPingAt:  time.Unix(2, 0),
+				Metadata:    map[string]interface{}{"host": "localhost"},
+			},
+			{
+				ID:          "room-terminating",
+				SchedulerID: "game",
+				Version:     "1.0",
+				Status:      game_room.GameStatusTerminating,
+				LastPingAt:  time.Unix(3, 0),
+				Metadata:    map[string]interface{}{"host": "localhost"},
+			},
+		}
+
+		for _, room := range rooms {
+			require.NoError(t, storage.CreateRoom(ctx, room))
+		}
+
+		allocatedRoomID, err := storage.AllocateRoom(ctx, "game")
+		require.Error(t, err)
+		require.Equal(t, "", allocatedRoomID)
+		require.Contains(t, err.Error(), "no ready rooms available")
+	})
+}

internal/api/handlers/rooms_handler.go

@@ -122,3 +122,33 @@ func (h *RoomsHandler) GetRoomAddress(ctx context.Context, message *api.GetRoomA
 	}
 	return requestadapters.FromInstanceEntityToGameRoomAddressResponse(instance), nil
 }
+
+func (h *RoomsHandler) AllocateRoom(ctx context.Context, message *api.AllocateRoomRequest) (*api.AllocateRoomResponse, error) {
+	handlerLogger := h.logger.With(zap.String(logs.LogFieldSchedulerName, message.SchedulerName))
+	handlerLogger.Debug("handling room allocation request", zap.Any("message", message))
+
+	roomID, err := h.roomManager.AllocateRoom(ctx, message.SchedulerName)
+	if err != nil {
+		handlerLogger.Error("error allocating room", zap.Any("allocation_request", message), zap.Error(err))
+		return nil, status.Error(codes.Unknown, err.Error())
+	}
+
+	// Get room instance to retrieve address information
+	instance, err := h.roomManager.GetRoomInstance(ctx, message.SchedulerName, roomID)
+	if err != nil {
+		handlerLogger.Error("error getting allocated room instance", zap.String(logs.LogFieldRoomID, roomID), zap.Error(err))
+		return nil, status.Error(codes.Unknown, err.Error())
+	}
+
+	// Convert instance to address response format and include in allocation response
+	addressResponse := requestadapters.FromInstanceEntityToGameRoomAddressResponse(instance)
+
+	handlerLogger.Debug("room allocated successfully", zap.String(logs.LogFieldRoomID, roomID))
+	return &api.AllocateRoomResponse{
+		RoomId:  roomID,
+		Success: true,
+		Message: "",
+		Ports:   addressResponse.Ports,
+		Host:    addressResponse.Host,
+	}, nil
+}