Feature/fetch only needed user presence

ChatInfra/modules/dynamodb/main.tf

@@ -48,6 +48,12 @@ resource "aws_dynamodb_table" "conversation_members" {
     type = "S"
   }
 
+  global_secondary_index {
+    name            = "userId-index"
+    hash_key        = "userId"
+    projection_type = "KEYS_ONLY"
+  }
+
   server_side_encryption {
     enabled = true
   }

ChatInfra/modules/dynamodb/outputs.tf

@@ -3,6 +3,7 @@ output "table_arns" {
     aws_dynamodb_table.server_registry.arn,
     aws_dynamodb_table.user_connections.arn,
     aws_dynamodb_table.conversation_members.arn,
+    "${aws_dynamodb_table.conversation_members.arn}/index/*",
     aws_dynamodb_table.messages.arn,
     aws_dynamodb_table.pending_messages.arn,
   ]

ChatInfra/modules/ecs/main.tf

@@ -186,6 +186,7 @@ resource "aws_iam_role_policy" "chatserver_dynamodb" {
       Action = [
         "dynamodb:DescribeTable",
         "dynamodb:GetItem",
+        "dynamodb:BatchGetItem",
         "dynamodb:PutItem",
         "dynamodb:DeleteItem",
         "dynamodb:UpdateItem",

ChatServer/README.md

@@ -1,20 +1,17 @@
 # ChatServer
 
 ## Overview
-This ChatServer handles WebSocket connections and message delivery for the chat system.
+ChatServer handles WebSocket connections and message delivery for the chat system.
 
-A chat server will register itself and its SQS queue URL with DynamoDB.
+A chat server registers itself and its SQS queue URL with DynamoDB on startup.
 
-When a user joins a given chat server the user is associated with the server in DynamoDB.
-Also, a mapping is kept of each user and their WebSocket session at runtime.
+When a user connects, they are associated with this server in DynamoDB and a runtime mapping of userId to WebSocket session is maintained. A presence snapshot is sent to the connecting user showing which of their groupmates (users who share at least one conversation with them) are currently online. Presence is scoped using a GSI on `ConversationMembers` to avoid a full table scan.
 
-When a user sends a message, the target user's server is found from DynamoDB and the associated SQS queue URL is used to route the message to that server for delivery.
+When a user sends a message, the target user's server is found from DynamoDB and the associated SQS queue URL routes the message to that server for delivery.
 
 If the target user is offline the message is saved to a pending messages table in DynamoDB and delivered when the user reconnects.
 
-Both 1-to-1 and group chat use the same conversation model. 
-A conversation has members stored in DynamoDB and a message sent to a conversation is fanned out to all members. Offline behavior is the same 
-- undelivered messages are saved as pending.
+Both 1-to-1 and group chat use the same conversation model. A conversation has members stored in DynamoDB and a message sent to a conversation is fanned out to all members. Offline behavior is the same -- undelivered messages are saved as pending.
 
 Message history is persisted to DynamoDB on every send and can be fetched via the REST API.
 
@@ -36,10 +33,12 @@ Message history is persisted to DynamoDB on every send and can be fetched via th
 | `serverId` |     | String | Server the user is currently connected to |
 
 ### ConversationMembers
-| Attribute        | Key | Type   | Description                    |
-|------------------|-----|--------|--------------------------------|
-| `conversationId` | PK  | String | Unique conversation identifier |
-| `userId`         | SK  | String | Member of the conversation     |
+| Attribute        | Key         | Type   | Description                    |
+|------------------|-------------|--------|--------------------------------|
+| `conversationId` | PK          | String | Unique conversation identifier |
+| `userId`         | SK / GSI PK | String | Member of the conversation     |
+
+**GSI: `userId-index`** -- partition key `userId`, projection `KEYS_ONLY`. Used to look up all conversations a user belongs to, which enables scoped presence snapshots and targeted presence broadcast delivery.
 
 ### Messages
 | Attribute        | Key | Type   | Description                          |

ChatServer/src/main/kotlin/data/AbstractDynamoCompositeRegistry.kt

@@ -5,6 +5,7 @@ import software.amazon.awssdk.services.dynamodb.DynamoDbClient
 import software.amazon.awssdk.services.dynamodb.model.AttributeDefinition
 import software.amazon.awssdk.services.dynamodb.model.BillingMode
 import software.amazon.awssdk.services.dynamodb.model.CreateTableRequest
+import software.amazon.awssdk.services.dynamodb.model.GlobalSecondaryIndex
 import software.amazon.awssdk.services.dynamodb.model.KeySchemaElement
 import software.amazon.awssdk.services.dynamodb.model.KeyType
 import software.amazon.awssdk.services.dynamodb.model.ResourceInUseException
@@ -16,6 +17,10 @@ abstract class AbstractDynamoCompositeRegistry(protected val dynamoClient: Dynam
     protected abstract val partitionKey: String
     protected abstract val sortKey: String
 
+    protected open fun gsiDefinitions(): List<GlobalSecondaryIndex> = emptyList()
+
+    protected open fun gsiAttributeDefinitions(): List<AttributeDefinition> = emptyList()
+
     private val logger = LoggerFactory.getLogger(AbstractDynamoCompositeRegistry::class.java)
 
     protected fun ensureTableExists() {
@@ -24,20 +29,24 @@ abstract class AbstractDynamoCompositeRegistry(protected val dynamoClient: Dynam
         } catch (e: ResourceNotFoundException) {
             logger.info("$tableName table not found, creating...")
             try {
-                dynamoClient.createTable(
+                val attrDefs =
+                    buildList {
+                        add(AttributeDefinition.builder().attributeName(partitionKey).attributeType(ScalarAttributeType.S).build())
+                        add(AttributeDefinition.builder().attributeName(sortKey).attributeType(ScalarAttributeType.S).build())
+                        addAll(gsiAttributeDefinitions())
+                    }
+                val tableBuilder =
                     CreateTableRequest.builder()
                         .tableName(tableName)
                         .billingMode(BillingMode.PAY_PER_REQUEST)
-                        .attributeDefinitions(
-                            AttributeDefinition.builder().attributeName(partitionKey).attributeType(ScalarAttributeType.S).build(),
-                            AttributeDefinition.builder().attributeName(sortKey).attributeType(ScalarAttributeType.S).build(),
-                        )
+                        .attributeDefinitions(attrDefs)
                         .keySchema(
                             KeySchemaElement.builder().attributeName(partitionKey).keyType(KeyType.HASH).build(),
                             KeySchemaElement.builder().attributeName(sortKey).keyType(KeyType.RANGE).build(),
                         )
-                        .build(),
-                )
+                val gsis = gsiDefinitions()
+                if (gsis.isNotEmpty()) tableBuilder.globalSecondaryIndexes(gsis)
+                dynamoClient.createTable(tableBuilder.build())
                 dynamoClient.waiter().waitUntilTableExists { it.tableName(tableName) }
                 logger.info("$tableName table created")
             } catch (e: ResourceInUseException) {

ChatServer/src/main/kotlin/data/registry/ConversationRegistry.kt

@@ -2,8 +2,14 @@ package org.chatserver.data.registry
 
 import org.chatserver.data.AbstractDynamoCompositeRegistry
 import software.amazon.awssdk.services.dynamodb.DynamoDbClient
+import software.amazon.awssdk.services.dynamodb.model.AttributeDefinition
 import software.amazon.awssdk.services.dynamodb.model.AttributeValue
 import software.amazon.awssdk.services.dynamodb.model.DeleteItemRequest
+import software.amazon.awssdk.services.dynamodb.model.GlobalSecondaryIndex
+import software.amazon.awssdk.services.dynamodb.model.KeySchemaElement
+import software.amazon.awssdk.services.dynamodb.model.KeyType
+import software.amazon.awssdk.services.dynamodb.model.Projection
+import software.amazon.awssdk.services.dynamodb.model.ProjectionType
 import software.amazon.awssdk.services.dynamodb.model.PutItemRequest
 import software.amazon.awssdk.services.dynamodb.model.QueryRequest
 
@@ -12,6 +18,22 @@ class ConversationRegistry(dynamoClient: DynamoDbClient) : AbstractDynamoComposi
     override val partitionKey = "conversationId"
     override val sortKey = "userId"
 
+    companion object {
+        private const val USER_ID_INDEX = "userId-index"
+    }
+
+    // userId is already an attribute (table sort key), so no extra attribute definition needed.
+    override fun gsiDefinitions(): List<GlobalSecondaryIndex> =
+        listOf(
+            GlobalSecondaryIndex.builder()
+                .indexName(USER_ID_INDEX)
+                .keySchema(KeySchemaElement.builder().attributeName(sortKey).keyType(KeyType.HASH).build())
+                .projection(Projection.builder().projectionType(ProjectionType.KEYS_ONLY).build())
+                .build(),
+        )
+
+    override fun gsiAttributeDefinitions(): List<AttributeDefinition> = emptyList()
+
     fun init() {
         ensureTableExists()
     }
@@ -61,4 +83,25 @@ class ConversationRegistry(dynamoClient: DynamoDbClient) : AbstractDynamoComposi
             )
         return response.items().mapNotNull { it[sortKey]?.s() }
     }
+
+    fun getConversationsForUser(userId: String): List<String> {
+        val response =
+            dynamoClient.query(
+                QueryRequest.builder()
+                    .tableName(tableName)
+                    .indexName(USER_ID_INDEX)
+                    .keyConditionExpression("$sortKey = :uid")
+                    .expressionAttributeValues(mapOf(":uid" to AttributeValue.fromS(userId)))
+                    .build(),
+            )
+        return response.items().mapNotNull { it[partitionKey]?.s() }
+    }
+
+    fun getGroupmates(userId: String): Set<String> {
+        val conversationIds = getConversationsForUser(userId)
+        return conversationIds
+            .flatMap { getMembers(it) }
+            .toSet()
+            .minus(userId)
+    }
 }

ChatServer/src/main/kotlin/data/registry/UserRegistry.kt

@@ -5,6 +5,7 @@ import org.slf4j.LoggerFactory
 import software.amazon.awssdk.services.dynamodb.DynamoDbClient
 import software.amazon.awssdk.services.dynamodb.model.AttributeValue
 import software.amazon.awssdk.services.dynamodb.model.DeleteItemRequest
+import software.amazon.awssdk.services.dynamodb.model.KeysAndAttributes
 import software.amazon.awssdk.services.dynamodb.model.PutItemRequest
 import software.amazon.awssdk.services.dynamodb.model.ScanRequest
 
@@ -56,6 +57,22 @@ class UserRegistry(
         return response.items().mapNotNull { it["userId"]?.s() }
     }
 
+    fun getConnectedUsersFrom(userIds: Set<String>): Set<String> {
+        if (userIds.isEmpty()) return emptySet()
+        val connected = mutableSetOf<String>()
+        userIds.chunked(100).forEach { chunk ->
+            val keys = chunk.map { mapOf("userId" to AttributeValue.fromS(it)) }
+            val response =
+                dynamoClient.batchGetItem {
+                    it.requestItems(mapOf(tableName to KeysAndAttributes.builder().keys(keys).build()))
+                }
+            response.responses()[tableName]?.forEach { item ->
+                item["userId"]?.s()?.let { id -> connected.add(id) }
+            }
+        }
+        return connected
+    }
+
     fun deregister(userId: String) {
         dynamoClient.deleteItem(
             DeleteItemRequest.builder()

ChatServer/src/main/kotlin/di/AppModule.kt

@@ -48,7 +48,7 @@ val appModule =
         single { PendingMessageRepository(get()) }
         single { ConversationRegistry(get()) }
         single { MessageRepository(get()) }
-        single { SqsConsumer(get(), get(), get()) }
+        single { SqsConsumer(get(), get(), get(), get()) }
         single { MessageRouter(get(), get(), get(), get(), get(), get()) }
         single { SessionStore() }
         single {