coordinator: user state manager with Async interface

Author: jpnovais

coordinator/app/build.gradle

@@ -22,6 +22,7 @@ dependencies {
   implementation project(':jvm-libs:linea:core:metrics')
   implementation project(':jvm-libs:linea:metrics:micrometer')
   implementation project(':jvm-libs:linea:core:domain-models')
+  implementation project(':jvm-libs:linea:clients:linea-state-manager')
   implementation project(':coordinator:utilities')
   implementation project(':coordinator:core')
   implementation project(':coordinator:clients:shomei-client')

coordinator/app/src/main/kotlin/net/consensys/zkevm/coordinator/app/L1DependentApp.kt

@@ -1,5 +1,7 @@
 package net.consensys.zkevm.coordinator.app
 
+import build.linea.clients.StateManagerClientV1
+import build.linea.clients.StateManagerV1JsonRpcClient
 import io.vertx.core.Vertx
 import kotlinx.datetime.Clock
 import net.consensys.linea.BlockNumberAndHash
@@ -319,6 +321,20 @@ class L1DependentApp(
     )
   }
 
+  private fun createStateManagerClientNew(
+    stateManagerConfig: StateManagerClientConfig,
+    logger: Logger
+  ): StateManagerClientV1 {
+    return StateManagerV1JsonRpcClient(
+      rpcClientFactory = httpJsonRpcClientFactory,
+      endpoints = stateManagerConfig.endpoints.map { it.toURI() },
+      maxInflightRequestsPerClient = stateManagerConfig.requestLimitPerEndpoint,
+      requestRetry = stateManagerConfig.requestRetryConfig,
+      zkStateManagerVersion = stateManagerConfig.version,
+      logger = logger
+    )
+  }
+
   private val lastFinalizedBlock = lastFinalizedBlock().get()
   private val lastProcessedBlockNumber = resumeConflationFrom(
     aggregationsRepository,
@@ -422,8 +438,16 @@ class L1DependentApp(
   private val conflationService: ConflationService =
     ConflationServiceImpl(calculator = conflationCalculator, metricsFacade = metricsFacade)
 
-  private val zkStateClient: Type2StateManagerClient =
-    createStateManagerClient(configs.stateManager, LogManager.getLogger("clients.StateManagerShomeiClient"))
+//  private val zkStateClient: Type2StateManagerClient =
+//    createStateManagerClient(configs.stateManager, LogManager.getLogger("clients.StateManagerShomeiClient"))
+  private val zkStateClient: StateManagerClientV1 = StateManagerV1JsonRpcClient(
+    rpcClientFactory = httpJsonRpcClientFactory,
+    endpoints = configs.stateManager.endpoints.map { it.toURI() },
+    maxInflightRequestsPerClient = configs.stateManager.requestLimitPerEndpoint,
+    requestRetry = configs.stateManager.requestRetryConfig,
+    zkStateManagerVersion = configs.stateManager.version,
+    logger = LogManager.getLogger("clients.StateManagerShomeiClient")
+  )
 
   private val lineaSmartContractClientForDataSubmission: LineaRollupSmartContractClient = run {
     // The below gas provider will act as the primary gas provider if L1

coordinator/app/src/main/kotlin/net/consensys/zkevm/ethereum/coordination/blob/BlobZkStateProviderImpl.kt

@@ -1,34 +1,21 @@
 package net.consensys.zkevm.ethereum.coordination.blob
 
-import com.github.michaelbull.result.Err
-import com.github.michaelbull.result.Ok
-import net.consensys.zkevm.coordinator.clients.GetZkEVMStateMerkleProofResponse
-import net.consensys.zkevm.coordinator.clients.Type2StateManagerClient
+import build.linea.clients.GetStateMerkleProofRequest
+import build.linea.clients.StateManagerClientV1
+import build.linea.domain.BlockInterval
 import tech.pegasys.teku.infrastructure.async.SafeFuture
-import tech.pegasys.teku.infrastructure.unsigned.UInt64
-
-class BlobZkStateProviderImpl(private val zkStateClient: Type2StateManagerClient) : BlobZkStateProvider {
-  private fun rollupGetZkEVMStateMerkleProof(startBlockNumber: ULong, endBlockNumber: ULong):
-    SafeFuture<GetZkEVMStateMerkleProofResponse> {
-    return zkStateClient.rollupGetZkEVMStateMerkleProof(
-      UInt64.valueOf(startBlockNumber.toLong()),
-      UInt64.valueOf(endBlockNumber.toLong())
-    ).thenCompose {
-      when (it) {
-        is Ok -> SafeFuture.completedFuture(it.value)
-        is Err -> {
-          SafeFuture.failedFuture(it.error.asException())
-        }
-      }
-    }
-  }
 
+class BlobZkStateProviderImpl(
+  private val zkStateClient: StateManagerClientV1
+) : BlobZkStateProvider {
   override fun getBlobZKState(blockRange: ULongRange): SafeFuture<BlobZkState> {
-    return rollupGetZkEVMStateMerkleProof(blockRange.first, blockRange.last).thenApply {
-      BlobZkState(
-        parentStateRootHash = it.zkParentStateRootHash.toArray(),
-        finalStateRootHash = it.zkEndStateRootHash.toArray()
-      )
-    }
+    return zkStateClient
+      .makeRequest(GetStateMerkleProofRequest(BlockInterval(blockRange.first, blockRange.last)))
+      .thenApply {
+        BlobZkState(
+          parentStateRootHash = it.zkParentStateRootHash,
+          finalStateRootHash = it.zkEndStateRootHash
+        )
+      }
   }
 }

coordinator/clients/prover-client/file-based-client/src/main/kotlin/net/consensys/zkevm/coordinator/clients/prover/FileBasedExecutionProverClientV2.kt

@@ -61,7 +61,7 @@ internal class ExecutionProofRequestDataDecorator(
         getBlockStateRootHash(request.blocks.first().blockNumber.toULong() - 1UL)
       ) { blocksAndBridgeLogs, previousKeccakStateRootHash ->
         BatchExecutionProofRequestDto(
-          zkParentStateRootHash = request.type2StateData.zkParentStateRootHash.toHexString(),
+          zkParentStateRootHash = request.type2StateData.zkParentStateRootHash.encodeHex(),
           keccakParentStateRootHash = previousKeccakStateRootHash,
           conflatedExecutionTracesFile = request.tracesResponse.tracesFileName,
           tracesEngineVersion = request.tracesResponse.tracesEngineVersion,

coordinator/clients/prover-client/file-based-client/src/test/kotlin/net/consensys/zkevm/coordinator/clients/prover/ExecutionProofRequestDataDecoratorTest.kt

@@ -1,14 +1,14 @@
 package net.consensys.zkevm.coordinator.clients.prover
 
+import build.linea.clients.GetZkEVMStateMerkleProofResponse
 import com.fasterxml.jackson.databind.node.ArrayNode
+import net.consensys.ByteArrayExt
 import net.consensys.encodeHex
 import net.consensys.zkevm.coordinator.clients.BatchExecutionProofRequestV1
 import net.consensys.zkevm.coordinator.clients.GenerateTracesResponse
-import net.consensys.zkevm.coordinator.clients.GetZkEVMStateMerkleProofResponse
 import net.consensys.zkevm.coordinator.clients.L2MessageServiceLogsClient
 import net.consensys.zkevm.domain.RlpBridgeLogsData
 import net.consensys.zkevm.encoding.ExecutionPayloadV1Encoder
-import org.apache.tuweni.bytes.Bytes32
 import org.assertj.core.api.Assertions.assertThat
 import org.junit.jupiter.api.BeforeEach
 import org.junit.jupiter.api.Test
@@ -52,8 +52,8 @@ class ExecutionProofRequestDataDecoratorTest {
     val executionPayload2 = executionPayloadV1(blockNumber = 124, gasLimit = 20_000_000UL)
     val type2StateResponse = GetZkEVMStateMerkleProofResponse(
       zkStateMerkleProof = ArrayNode(null),
-      zkParentStateRootHash = Bytes32.random(),
-      zkEndStateRootHash = Bytes32.random(),
+      zkParentStateRootHash = ByteArrayExt.random32(),
+      zkEndStateRootHash = ByteArrayExt.random32(),
       zkStateManagerVersion = "2.0.0"
     )
     val generateTracesResponse = GenerateTracesResponse(
@@ -82,7 +82,7 @@ class ExecutionProofRequestDataDecoratorTest {
     val requestDto = requestDatDecorator.invoke(request).get()
 
     assertThat(requestDto.keccakParentStateRootHash).isEqualTo(stateRoot)
-    assertThat(requestDto.zkParentStateRootHash).isEqualTo(type2StateResponse.zkParentStateRootHash.toHexString())
+    assertThat(requestDto.zkParentStateRootHash).isEqualTo(type2StateResponse.zkParentStateRootHash.encodeHex())
     assertThat(requestDto.conflatedExecutionTracesFile).isEqualTo("123-114-conflated-traces.json")
     assertThat(requestDto.tracesEngineVersion).isEqualTo("1.0.0")
     assertThat(requestDto.type2StateManagerVersion).isEqualTo("2.0.0")

coordinator/clients/type2-state-manager-client/build.gradle

@@ -7,6 +7,7 @@ dependencies {
   implementation project(':jvm-libs:generic:json-rpc')
   implementation project(':jvm-libs:linea:metrics:micrometer')
   implementation project(':jvm-libs:generic:extensions:futures')
+  implementation project(':jvm-libs:linea:clients:linea-state-manager')
 
   implementation "tech.pegasys.teku.internal:bytes:${libs.versions.teku.get()}"
   implementation "io.vertx:vertx-core"

coordinator/core/build.gradle

@@ -9,6 +9,7 @@ dependencies {
   api project(':jvm-libs:linea:core:domain-models')
   api project(':jvm-libs:linea:core:metrics')
   api project(':jvm-libs:linea:core:long-running-service')
+  api project(':jvm-libs:linea:clients:linea-state-manager')
   api project(':jvm-libs:linea:core:traces')
   api project(':jvm-libs:generic:errors')
   api project(':jvm-libs:generic:extensions:kotlin')

coordinator/core/src/main/kotlin/net/consensys/zkevm/coordinator/clients/BatchExecutionProverRequestResponse.kt

@@ -7,7 +7,7 @@ import tech.pegasys.teku.ethereum.executionclient.schema.ExecutionPayloadV1
 data class BatchExecutionProofRequestV1(
   val blocks: List<ExecutionPayloadV1>,
   val tracesResponse: GenerateTracesResponse,
-  val type2StateData: GetZkEVMStateMerkleProofResponse
+  val type2StateData: build.linea.clients.GetZkEVMStateMerkleProofResponse
 ) : BlockInterval {
   override val startBlockNumber: ULong
     get() = blocks.first().blockNumber.toULong()

coordinator/core/src/main/kotlin/net/consensys/zkevm/ethereum/coordination/conflation/TracesConflationCoordinator.kt

@@ -1,8 +1,8 @@
 package net.consensys.zkevm.ethereum.coordination.conflation
 
+import build.linea.clients.GetZkEVMStateMerkleProofResponse
 import net.consensys.linea.BlockNumberAndHash
 import net.consensys.zkevm.coordinator.clients.GenerateTracesResponse
-import net.consensys.zkevm.coordinator.clients.GetZkEVMStateMerkleProofResponse
 import tech.pegasys.teku.infrastructure.async.SafeFuture
 
 data class BlocksTracesConflated(

coordinator/core/src/main/kotlin/net/consensys/zkevm/ethereum/coordination/conflation/TracesConflationCoordinatorImpl.kt

@@ -1,5 +1,9 @@
 package net.consensys.zkevm.ethereum.coordination.conflation
 
+import build.linea.clients.GetStateMerkleProofRequest
+import build.linea.clients.GetZkEVMStateMerkleProofResponse
+import build.linea.clients.StateManagerClientV1
+import build.linea.domain.BlockInterval
 import com.github.michaelbull.result.Err
 import com.github.michaelbull.result.Ok
 import com.github.michaelbull.result.Result
@@ -9,19 +13,15 @@ import com.github.michaelbull.result.mapBoth
 import net.consensys.linea.BlockNumberAndHash
 import net.consensys.linea.errors.ErrorResponse
 import net.consensys.zkevm.coordinator.clients.GenerateTracesResponse
-import net.consensys.zkevm.coordinator.clients.GetZkEVMStateMerkleProofResponse
 import net.consensys.zkevm.coordinator.clients.TracesConflationClientV1
 import net.consensys.zkevm.coordinator.clients.TracesServiceErrorType
-import net.consensys.zkevm.coordinator.clients.Type2StateManagerClient
-import net.consensys.zkevm.coordinator.clients.Type2StateManagerErrorType
 import org.apache.logging.log4j.LogManager
 import org.apache.logging.log4j.Logger
 import tech.pegasys.teku.infrastructure.async.SafeFuture
-import tech.pegasys.teku.infrastructure.unsigned.UInt64
 
 class TracesConflationCoordinatorImpl(
   private val tracesConflationClient: TracesConflationClientV1,
-  private val zkStateClient: Type2StateManagerClient
+  private val zkStateClient: StateManagerClientV1
 ) : TracesConflationCoordinator {
   private val log: Logger = LogManager.getLogger(this::class.java)
   private fun requestConflatedTraces(
@@ -45,22 +45,7 @@ class TracesConflationCoordinatorImpl(
     startBlockNumber: ULong,
     endBlockNumber: ULong
   ): SafeFuture<GetZkEVMStateMerkleProofResponse> {
-    return zkStateClient
-      .rollupGetZkEVMStateMerkleProof(
-        UInt64.valueOf(startBlockNumber.toLong()),
-        UInt64.valueOf(endBlockNumber.toLong())
-      )
-      .thenCompose { result:
-            Result<GetZkEVMStateMerkleProofResponse, ErrorResponse<Type2StateManagerErrorType>>
-        ->
-        result.mapBoth(
-          { SafeFuture.completedFuture(it) },
-          {
-            log.error("Type2State manager returned error={}", it)
-            SafeFuture.failedFuture(it.asException("State manager error: ${it.message}"))
-          }
-        )
-      }
+    return zkStateClient.makeRequest(GetStateMerkleProofRequest(BlockInterval(startBlockNumber, endBlockNumber)))
   }
 
   override fun conflateExecutionTraces(

jvm-libs/linea/clients/linea-state-manager/src/main/kotlin/build/linea/clients/StateManagerClientV1.kt

@@ -13,11 +13,12 @@ enum class StateManagerErrorType : ClientError {
   BLOCK_MISSING_IN_CHAIN
 }
 
-sealed interface StateManagerRequest
-sealed class GetChainHeadRequest() : StateManagerRequest
-data class GetStateMerkleProofRequest(
-  val blockInterval: BlockInterval
-) : StateManagerRequest, BlockInterval by blockInterval
+sealed interface StateManagerRequest<ResponseType : StateManagerResponse> : ClientRequest<ResponseType>
+sealed class GetChainHeadRequest() : StateManagerRequest<GetChainHeadResponse>
+
+data class GetStateMerkleProofRequest(val blockInterval: BlockInterval) :
+  StateManagerRequest<GetZkEVMStateMerkleProofResponse>,
+  BlockInterval by blockInterval
 
 sealed interface StateManagerResponse
 
@@ -57,14 +58,9 @@ data class GetZkEVMStateMerkleProofResponse(
   }
 }
 
-// Type alias dedicated for each method
-typealias StateManagerClientToGetStateMerkleProofV0 =
-  AsyncClient<GetStateMerkleProofRequest, GetZkEVMStateMerkleProofResponse>
-
-typealias StateManagerClientToGetChainHeadV1 =
-  AsyncClient<GetChainHeadRequest, ULong>
+data class GetChainHeadResponse(val headBlockNumber: ULong) : StateManagerResponse
 
-interface StateManagerClientV1 {
+interface StateManagerClientV1 : AsyncClient<StateManagerRequest<*>> {
   /**
    * Get the head block number of the chain.
    * @return GetZkEVMStateMerkleProofResponse
@@ -84,4 +80,15 @@ interface StateManagerClientV1 {
   ): SafeFuture<Result<GetZkEVMStateMerkleProofResponse, ErrorResponse<StateManagerErrorType>>>
 
   fun rollupGetHeadBlockNumber(): SafeFuture<ULong>
+
+  override fun <Response> makeRequest(request: ClientRequest<Response>): SafeFuture<Response> {
+    @Suppress("UNCHECKED_CAST")
+    return when (request) {
+      is GetStateMerkleProofRequest -> rollupGetStateMerkleProof(request.blockInterval) as SafeFuture<Response>
+      is GetChainHeadRequest -> rollupGetHeadBlockNumber()
+        .thenApply { GetChainHeadResponse(it) } as SafeFuture<Response>
+
+      else -> throw IllegalArgumentException("Unknown request type: $request")
+    }
+  }
 }

jvm-libs/linea/core/client-interface/src/main/kotlin/build/linea/clients/Client.kt

@@ -22,8 +22,9 @@ interface Client<Request, Response> {
   fun makeRequest(request: Request): Response
 }
 
-interface AsyncClient<Request, Response> {
-  fun makeRequest(request: Request): SafeFuture<Response>
+interface ClientRequest<Response>
+interface AsyncClient<ReqSupperType> {
+  fun <Response> makeRequest(request: ClientRequest<Response>): SafeFuture<Response>
 }
 
 fun <T, E : ClientError> SafeFuture<Result<T, ErrorResponse<E>>>.unwrapResultMonad(): SafeFuture<T> {