address PR comments

Author: somdoron

client-server-client/src/main/scala/zio/raft/client/ClientTransport.scala

@@ -0,0 +1,69 @@
+package zio.raft.client
+
+import zio.*
+import zio.stream.*
+import zio.raft.protocol.*
+import scodec.bits.BitVector
+import zio.zmq.{ZContext, ZSocket}
+import zio.raft.protocol.Codecs.{clientMessageCodec, serverMessageCodec}
+
+/** Transport abstraction for client-server communication.
+  */
+trait ClientTransport {
+  def connect(address: String): ZIO[Any, Throwable, Unit]
+  def disconnect(): ZIO[Any, Throwable, Unit]
+  def sendMessage(message: ClientMessage): ZIO[Any, Throwable, Unit]
+  def incomingMessages: ZStream[Any, Throwable, ServerMessage]
+}
+
+object ClientTransport {
+
+  /** Create a ZeroMQ-based client transport.
+    */
+  def make(config: ClientConfig): ZIO[ZContext & Scope, Throwable, ClientTransport] =
+    for {
+      socket <- ZSocket.client
+      _ <- socket.options.setLinger(0)
+      _ <- socket.options.setHeartbeat(1.seconds, 10.second, 30.second)
+      timeoutConnectionClosed = serverMessageCodec
+        .encode(SessionClosed(SessionCloseReason.ConnectionClosed, None))
+        .require
+        .toByteArray
+      _ <- socket.options.setHiccupMessage(timeoutConnectionClosed)
+      _ <- socket.options.setHighWatermark(200000, 200000)
+      lastAddressRef <- Ref.make(Option.empty[String])
+      transport = new Zmq(socket, lastAddressRef)
+    } yield transport
+
+  /** ZeroMQ-based implementation of ClientTransport.
+    */
+  private[client] class Zmq(socket: ZSocket, lastAddressRef: Ref[Option[String]])
+      extends ClientTransport {
+    override def connect(address: String): ZIO[Any, Throwable, Unit] =
+      for {
+        _ <- lastAddressRef.set(Some(address))
+        _ <- socket.connect(address)
+      } yield ()
+
+    override def disconnect(): ZIO[Any, Throwable, Unit] =
+      for {
+        lastAddress <- lastAddressRef.get
+        _ <- lastAddress match {
+          case Some(address) => socket.disconnect(address)
+          case None          => ZIO.unit
+        }
+        _ <- lastAddressRef.set(None)
+      } yield ()
+
+    override def sendMessage(message: ClientMessage): ZIO[Any, Throwable, Unit] =
+      for {
+        bytes <- ZIO.attempt(clientMessageCodec.encode(message).require.toByteArray)
+        _ <- socket.sendImmediately(bytes)
+      } yield ()
+
+    override def incomingMessages: ZStream[Any, Throwable, ServerMessage] =
+      socket.stream.mapZIO { msg =>
+        ZIO.attempt(serverMessageCodec.decode(BitVector(msg.data())).require.value)
+      }
+  }
+}

client-server-client/src/main/scala/zio/raft/client/PendingRequests.scala

@@ -0,0 +1,72 @@
+package zio.raft.client
+
+import zio.*
+import zio.raft.protocol.*
+import scodec.bits.ByteVector
+import java.time.Instant
+
+/** Manages pending requests with lastSentAt timestamps for retry.
+  */
+case class PendingRequests(
+  requests: Map[RequestId, PendingRequests.PendingRequestData]
+) {
+  def add(
+    requestId: RequestId,
+    payload: ByteVector,
+    promise: Promise[Throwable, ByteVector],
+    sentAt: Instant
+  ): PendingRequests =
+    copy(requests = requests.updated(requestId, PendingRequests.PendingRequestData(payload, promise, sentAt, sentAt)))
+
+  def complete(requestId: RequestId, result: ByteVector): ZIO[Any, Nothing, PendingRequests] =
+    requests.get(requestId) match {
+      case Some(data) =>
+        data.promise.succeed(result).as(copy(requests = requests.removed(requestId)))
+      case None =>
+        ZIO.succeed(this)
+    }
+
+  /** Resend all pending requests (used after successful connection). Returns updated PendingRequests with new
+    * lastSentAt timestamps.
+    */
+  def resendAll(transport: ClientTransport): UIO[PendingRequests] =
+    ZIO.foldLeft(requests.toList)(this) { case (pending, (requestId, data)) =>
+      for {
+        now <- Clock.instant
+        request = ClientRequest(requestId, data.payload, now)
+        _ <- transport.sendMessage(request).orDie
+        _ <- ZIO.logDebug(s"Resending pending request: $requestId")
+        updatedData = data.copy(lastSentAt = now)
+      } yield PendingRequests(pending.requests.updated(requestId, updatedData))
+    }
+
+  /** Resend expired requests and update lastSentAt.
+    */
+  def resendExpired(transport: ClientTransport, currentTime: Instant, timeout: Duration): UIO[PendingRequests] = {
+    val timeoutSeconds = timeout.toSeconds
+    ZIO.foldLeft(requests.toList)(this) { case (pending, (requestId, data)) =>
+      val elapsed = Duration.fromInterval(data.lastSentAt, currentTime)
+      if (elapsed > timeout) {
+        val request = ClientRequest(requestId, data.payload, currentTime)
+        for {
+          _ <- transport.sendMessage(request).orDie
+          _ <- ZIO.logDebug(s"Resending timed out request: $requestId")
+          updatedData = data.copy(lastSentAt = currentTime)
+        } yield PendingRequests(pending.requests.updated(requestId, updatedData))
+      } else {
+        ZIO.succeed(pending)
+      }
+    }
+  }
+}
+
+object PendingRequests {
+  def empty: PendingRequests = PendingRequests(Map.empty)
+
+  case class PendingRequestData(
+    payload: ByteVector,
+    promise: Promise[Throwable, ByteVector],
+    createdAt: Instant,
+    lastSentAt: Instant
+  )
+}