address pr comments

somdoron · somdoron · commit 2ff19d1254b3 · 2025-10-26T10:24:59.000+02:00
diff --git a/client-server-protocol/src/main/scala/zio/raft/protocol/package.scala b/client-server-protocol/src/main/scala/zio/raft/protocol/package.scala
@@ -86,6 +86,8 @@ package object protocol {
       def isLowerThan(other: RequestId): Boolean = value < other.value
 
       def isGreaterThan(other: RequestId): Boolean = value > other.value
+
+      def increaseBy(x: Int) = RequestId(value + x)
     }
   }
 
diff --git a/session-state-machine/src/main/scala/zio/raft/sessionstatemachine/SessionCommand.scala b/session-state-machine/src/main/scala/zio/raft/sessionstatemachine/SessionCommand.scala
@@ -21,8 +21,8 @@ object SessionCommand:
   /** A client request containing a user command to execute.
     *
     * This is the primary command type. The SessionStateMachine base class:
-    *   1. Updates highestLowestRequestIdSeen if lowestRequestId is higher 2. Checks if requestId <
-    *      highestLowestRequestIdSeen AND not in cache → return Left(SessionError.ResponseEvicted) 3. Checks if
+    *   1. Updates highestLowestRequestIdSeen if lowestRequestId is higher 2. Checks if requestId <=
+    *      highestLowestRequestIdSeen AND not in cache → return Left(RequestError.ResponseEvicted) 3. Checks if
     *      (sessionId, requestId) is already in the cache (idempotency) 4. If cached, returns the cached response 5. If
     *      not cached, narrows state to UserSchema and calls user's applyCommand method 6. Caches the response and
     *      returns it
@@ -33,13 +33,13 @@ object SessionCommand:
     *   The request ID (for idempotency checking)
     * @param lowestRequestId
     *   The lowest request ID for which client hasn't received response (for cache cleanup). Client is saying "I have
-    *   received all responses for requestIds < this value"
+    *   received all responses for requestIds <= this value (inclusive)"
     * @param command
     *   The user's command to execute
     *
     * @note
-    *   Response type is Either[RequestError, (command.Response, List[ServerRequestWithContext[SR]])] Left: error (e.g.,
-    *   response evicted) Right: (user command response, server requests with context)
+    *   Response type is Either[RequestError, (command.Response, List[ServerRequestEnvelope[SR]])] Left: error (e.g.,
+    *   response evicted) Right: (user command response, server request envelopes)
     * @note
     *   lowestRequestId enables the "Lowest Sequence Number Protocol" from Raft dissertation Ch. 6.3
     * @note
@@ -54,8 +54,8 @@ object SessionCommand:
     lowestRequestId: RequestId,
     command: UC
   ) extends SessionCommand[UC, SR]:
-    // Response type can be an error or the user command's response with server requests
-    type Response = Either[RequestError, (command.Response, List[ServerRequestWithContext[SR]])]
+    // Response type can be an error or the user command's response with server request envelopes
+    type Response = Either[RequestError, (command.Response, List[ServerRequestEnvelope[SR]])]
 
   /** Acknowledgment from a client for a server-initiated request.
     *
@@ -90,7 +90,7 @@ object SessionCommand:
     sessionId: SessionId,
     capabilities: Map[String, String]
   ) extends SessionCommand[Nothing, SR]:
-    type Response = List[ServerRequestWithContext[SR]] // server requests with context
+    type Response = List[ServerRequestEnvelope[SR]] // server request envelopes
 
   /** Notification that a session has expired.
     *
@@ -105,7 +105,7 @@ object SessionCommand:
     createdAt: Instant,
     sessionId: SessionId
   ) extends SessionCommand[Nothing, SR]:
-    type Response = List[ServerRequestWithContext[SR]] // server requests with context
+    type Response = List[ServerRequestEnvelope[SR]] // server request envelopes
 
   /** Command to atomically retrieve requests needing retry and update lastSentAt.
     *
diff --git a/session-state-machine/src/main/scala/zio/raft/sessionstatemachine/SessionStateMachine.scala b/session-state-machine/src/main/scala/zio/raft/sessionstatemachine/SessionStateMachine.scala
@@ -4,6 +4,7 @@ import zio.{UIO, Chunk}
 import zio.prelude.State
 import zio.raft.{Command, HMap, StateMachine, Index}
 import zio.raft.protocol.{SessionId, RequestId}
+import zio.raft.protocol.RequestId.RequestIdSyntax
 import zio.stream.Stream
 import java.time.Instant
 
@@ -267,14 +268,15 @@ trait SessionStateMachine[UC <: Command, R, SR, UserSchema <: Tuple]
     *
     * Implementation of Raft dissertation Chapter 6.3 session management protocol:
     *   1. Check cache for (sessionId, requestId) 2. If cache hit, return cached response 3. If cache miss, check if
-    *      requestId < highestLowestRequestIdSeen → response was evicted, return error 4. If cache miss + requestId >=
+    *      requestId <= highestLowestRequestIdSeen → response was evicted, return error 4. If cache miss + requestId >
     *      highestLowestRequestIdSeen, execute command and update highestLowestRequestIdSeen
     *
     * This correctly handles out-of-order requests. The lowestRequestId from the client tells us which responses have
-    * been acknowledged and can be evicted. We only update highestLowestRequestIdSeen for requests we actually process.
+    * been acknowledged (inclusive) and can be evicted. We only update highestLowestRequestIdSeen for requests we
+    * actually process.
     */
   private def handleClientRequest(cmd: SessionCommand.ClientRequest[UC, SR])
-    : State[HMap[Schema], Either[RequestError, (cmd.command.Response, List[ServerRequestWithContext[SR]])]] =
+    : State[HMap[Schema], Either[RequestError, (cmd.command.Response, List[ServerRequestEnvelope[SR]])]] =
     for
       highestLowestSeen <- getHighestLowestRequestIdSeen(cmd.sessionId)
       cachedOpt <- getCachedResponse((cmd.sessionId, cmd.requestId))
@@ -285,9 +287,9 @@ trait SessionStateMachine[UC <: Command, R, SR, UserSchema <: Tuple]
 
         case None =>
           // Cache miss - check if response was evicted
-          // If requestId < highestLowestRequestIdSeen, client has acknowledged receiving this response
-          if cmd.requestId.isLowerThan(highestLowestSeen) then
-            // Client said "I have responses for all requestIds < highestLowest", so this was evicted
+          // If requestId <= highestLowestRequestIdSeen, client has acknowledged receiving this response
+          if cmd.requestId.isLowerOrEqual(highestLowestSeen) then
+            // Client said "I have responses for all requestIds <= highestLowest", so this was evicted
             State.succeed(Left(RequestError.ResponseEvicted(cmd.sessionId, cmd.requestId)))
           else
             // requestId >= highestLowestRequestIdSeen
@@ -327,8 +329,8 @@ trait SessionStateMachine[UC <: Command, R, SR, UserSchema <: Tuple]
 
   /** Update the highest lowestRequestId seen from the client (only if lowestRequestId > current highest).
     *
-    * The lowestRequestId from the client indicates "I have received all responses for requestIds < this value". We
-    * track the highest such value to detect evicted responses.
+    * The lowestRequestId from the client indicates "I have received all responses for requestIds <= this value
+    * (inclusive)". We track the highest such value to detect evicted responses.
     */
   private def updateHighestLowestRequestIdSeen(
     sessionId: SessionId,
@@ -354,7 +356,7 @@ trait SessionStateMachine[UC <: Command, R, SR, UserSchema <: Tuple]
       val keysToRemove = state.range["serverRequests"](
         (cmd.sessionId, RequestId.zero),
         (cmd.sessionId, upperBoundExclusive)
-      ).map(_._1)
+      ).map((key, _) => key)
 
       // Remove all acknowledged requests in one efficient operation
       state.removedAll["serverRequests"](keysToRemove)
@@ -363,7 +365,7 @@ trait SessionStateMachine[UC <: Command, R, SR, UserSchema <: Tuple]
   /** Handle CreateSession command.
     */
   private def handleCreateSession(cmd: SessionCommand.CreateSession[SR])
-    : State[HMap[Schema], List[ServerRequestWithContext[SR]]] =
+    : State[HMap[Schema], List[ServerRequestEnvelope[SR]]] =
     for
       _ <- createSessionMetadata(cmd.sessionId, cmd.capabilities, cmd.createdAt)
       (serverRequestsLog, _) <- handleSessionCreated(cmd.sessionId, cmd.capabilities, cmd.createdAt).withLog
@@ -382,7 +384,7 @@ trait SessionStateMachine[UC <: Command, R, SR, UserSchema <: Tuple]
   /** Handle SessionExpired command (internal name to avoid conflict with abstract method).
     */
   private def handleSessionExpired_internal(cmd: SessionCommand.SessionExpired[SR])
-    : State[HMap[Schema], List[ServerRequestWithContext[SR]]] =
+    : State[HMap[Schema], List[ServerRequestEnvelope[SR]]] =
     for
       capabilities <- getSessionCapabilities(cmd.sessionId)
       (serverRequestsLog, _) <- handleSessionExpired(cmd.sessionId, capabilities, cmd.createdAt).withLog
@@ -408,9 +410,7 @@ trait SessionStateMachine[UC <: Command, R, SR, UserSchema <: Tuple]
       // Use foldRight to collect updated requests and build new state
       // This avoids vars and the need to reverse at the end
       state.iterator["serverRequests"].foldRight((List.empty[PendingServerRequest[SR]], state)) {
-        case (((sessionId, requestId), pendingAny), (accumulated, currentState)) =>
-          val pending = pendingAny.asInstanceOf[PendingServerRequest[SR]]
-
+        case (((sessionId, requestId), pending), (accumulated, currentState)) =>
           // Check if this request needs retry (lastSentAt before threshold)
           if pending.lastSentAt.isBefore(cmd.lastSentBefore) then
             // Update lastSentAt and add to accumulated list
@@ -439,7 +439,7 @@ trait SessionStateMachine[UC <: Command, R, SR, UserSchema <: Tuple]
   private def addServerRequests(
     createdAt: Instant,
     serverRequests: Chunk[ServerRequestForSession[SR]]
-  ): State[HMap[Schema], List[ServerRequestWithContext[SR]]] =
+  ): State[HMap[Schema], List[ServerRequestEnvelope[SR]]] =
     if serverRequests.isEmpty then
       State.succeed(Nil)
     else
@@ -448,52 +448,45 @@ trait SessionStateMachine[UC <: Command, R, SR, UserSchema <: Tuple]
         val groupedBySession = serverRequests.groupBy(_.sessionId)
 
         // Process each session's requests
-        val (finalState, allRequestsWithContext) =
-          groupedBySession.foldLeft((state, List.empty[ServerRequestWithContext[SR]])) {
+        val (finalState, allEnvelopes) =
+          groupedBySession.foldLeft((state, List.empty[ServerRequestEnvelope[SR]])) {
             case ((currentState, accumulated), (sessionId, sessionRequests)) =>
               // Get last assigned ID for this session
               val lastId = currentState.get["lastServerRequestId"](sessionId)
                 .getOrElse(RequestId.zero)
 
-              // Assign IDs starting from lastId + 1
-              val requestsWithIds = sessionRequests.zipWithIndex.map { case (reqWithSession, index) =>
-                val newId = RequestId(RequestId.unwrap(lastId) + index + 1)
-                (
-                  sessionId,
-                  newId,
-                  PendingServerRequest(
-                    payload = reqWithSession.payload,
-                    lastSentAt = createdAt
-                  )
-                )
-              }
-
-              // Update lastServerRequestId for this session
-              val newLastId = requestsWithIds.last._2
+              // Calculate new last ID: lastId + number of requests
+              val newLastId = lastId.increaseBy(sessionRequests.length)
               val stateWithNewId = currentState.updated["lastServerRequestId"](
                 sessionId,
                 newLastId
               )
 
-              // Add all requests with composite keys
-              val stateWithRequests = requestsWithIds.foldLeft(stateWithNewId) {
-                case (s, (sid, requestId, pending)) =>
-                  s.updated["serverRequests"]((sid, requestId), pending)
-              }
-
-              // Create requests with context
-              val requestsWithContext = requestsWithIds.map { case (sid, requestId, pending) =>
-                ServerRequestWithContext(
-                  sessionId = sid,
-                  requestId = requestId,
-                  payload = pending.payload
-                )
-              }.toList
-
-              (stateWithRequests, accumulated ++ requestsWithContext)
+              // Add all requests with composite keys and create envelopes in one pass
+              val (stateWithRequests, envelopes) =
+                sessionRequests.zipWithIndex.foldLeft((stateWithNewId, List.empty[ServerRequestEnvelope[SR]])) {
+                  case ((s, envs), (reqWithSession, index)) =>
+                    val requestId = lastId.increaseBy(index + 1)
+                    val pending = PendingServerRequest(
+                      payload = reqWithSession.payload,
+                      lastSentAt = createdAt
+                    )
+                    val updatedState = s.updated["serverRequests"]((sessionId, requestId), pending)
+                    val envelope = ServerRequestEnvelope(
+                      sessionId = sessionId,
+                      requestId = requestId,
+                      payload = reqWithSession.payload
+                    )
+                    (updatedState, envelope :: envs)
+                }
+
+              // Reverse to maintain order (since we prepended)
+              val orderedEnvelopes = envelopes.reverse
+
+              (stateWithRequests, accumulated ++ orderedEnvelopes)
           }
 
-        (allRequestsWithContext, finalState)
+        (allEnvelopes, finalState)
       }
 
   /** Remove all session data when session expires.
@@ -504,13 +497,13 @@ trait SessionStateMachine[UC <: Command, R, SR, UserSchema <: Tuple]
       val cacheKeysToRemove = state.range["cache"](
         (sessionId, RequestId.zero),
         (sessionId, RequestId.max)
-      ).map(_._1)
+      ).map((key, _) => key)
 
       // Remove all server requests for this session using range query
       val serverRequestKeysToRemove = state.range["serverRequests"](
         (sessionId, RequestId.zero),
         (sessionId, RequestId.max)
-      ).map(_._1)
+      ).map((key, _) => key)
 
       // Remove all session data in batch
       state
@@ -527,28 +520,30 @@ trait SessionStateMachine[UC <: Command, R, SR, UserSchema <: Tuple]
 
   /** Clean up cached responses based on lowestRequestId (Lowest Sequence Number Protocol).
     *
-    * Removes all cached responses for the session with requestId < lowestRequestId. This allows the client to control
+    * Removes all cached responses for the session with requestId <= lowestRequestId. This allows the client to control
     * cache cleanup by telling the server which responses it no longer needs (Chapter 6.3 of Raft dissertation).
     *
+    * The client sends lowestRequestId to indicate "I have received all responses up to and including this ID".
+    *
     * Uses range queries to efficiently find and remove old cache entries.
     */
   private def cleanupCache(
     sessionId: SessionId,
     lowestRequestId: RequestId
   ): State[HMap[Schema], Unit] =
     State.update { state =>
-      // Use range to find all cache entries for this session with requestId < lowestRequestId
-      // Range is [from, until), so we use (sessionId, RequestId.zero) to (sessionId, lowestRequestId)
+      // Use range to find all cache entries for this session with requestId <= lowestRequestId
+      // Range is [from, until), so to include lowestRequestId, we use lowestRequestId.next as upper bound
       val keysToRemove = state.range["cache"](
         (sessionId, RequestId.zero),
-        (sessionId, lowestRequestId)
-      ).map(_._1)
+        (sessionId, lowestRequestId.next)
+      ).map((key, _) => key)
 
       // Remove all old cache entries in one efficient operation
       state.removedAll["cache"](keysToRemove)
     }
 
-  /** Dirty read helper (FR-027) - check if ANY session has pending requests needing retry.
+  /** Dirty read helper - check if ANY session has pending requests needing retry.
     *
     * This method can be called directly (outside Raft consensus) to optimize the retry process. The retry process
     * performs a dirty read, applies policy locally, and only sends GetRequestsForRetry command if retries are needed.
diff --git a/session-state-machine/src/main/scala/zio/raft/sessionstatemachine/package.scala b/session-state-machine/src/main/scala/zio/raft/sessionstatemachine/package.scala
@@ -55,7 +55,7 @@ package object sessionstatemachine:
 
   /** Server request with assigned ID after being added to state.
     */
-  case class ServerRequestWithContext[SR](
+  case class ServerRequestEnvelope[SR](
     sessionId: SessionId,
     requestId: RequestId,
     payload: SR
@@ -140,8 +140,8 @@ package object sessionstatemachine:
     /** Response was cached but has been evicted. Client must create a new session.
       *
       * This error occurs when:
-      *   1. Client retries a request with requestId < highestLowestRequestIdSeen for the session 2. The response is not
-      *      in the cache (was evicted via cleanupCache)
+      *   1. Client retries a request with requestId <= highestLowestRequestIdSeen for the session 2. The response is
+      *      not in the cache (was evicted via cleanupCache)
       *
       * Per Raft dissertation Chapter 6.3, the client should create a new session and retry the operation.
       */
@@ -168,9 +168,9 @@ package object sessionstatemachine:
     *   - No data duplication: sessionId and requestId are not stored in the value, only in the key
     *
     * The highestLowestRequestIdSeen prefix enables detection of evicted responses:
-    *   - Client sends lowestRequestId indicating "I have received all responses < this ID"
+    *   - Client sends lowestRequestId indicating "I have received all responses <= this ID (inclusive)"
     *   - We track the highest lowestRequestId value we've seen from the client
-    *   - When a ClientRequest arrives, we check if requestId < highestLowestRequestIdSeen
+    *   - When a ClientRequest arrives, we check if requestId <= highestLowestRequestIdSeen
     *   - If yes and response is not in cache, we know it was evicted (client already acknowledged it)
     *   - This correctly handles out-of-order requests while preventing re-execution of acknowledged commands
     */
diff --git a/session-state-machine/src/test/scala/zio/raft/sessionstatemachine/IdempotencySpec.scala b/session-state-machine/src/test/scala/zio/raft/sessionstatemachine/IdempotencySpec.scala

Original file line number	Diff line number	Diff line change
`@@ -86,6 +86,8 @@ package object protocol {`
`86`	`86`	`def isLowerThan(other: RequestId): Boolean = value < other.value`
`87`	`87`
`88`	`88`	`def isGreaterThan(other: RequestId): Boolean = value > other.value`
	`89`	`+`
	`90`	`+ def increaseBy(x: Int) = RequestId(value + x)`
`89`	`91`	`}`
`90`	`92`	`}`
`91`	`93`