Better type for payload status

UnitsNetwork · Sep 18, 2024 · d4b0531 · d4b0531
1 parent a898b18
commit d4b0531
Show file tree

Hide file tree

Showing 8 changed files with 49 additions and 29 deletions.
diff --git a/src/main/scala/units/ELUpdater.scala b/src/main/scala/units/ELUpdater.scala
@@ -1471,12 +1471,12 @@ class ELUpdater(
       }
       .map(_ => ())
 
-  private def confirmBlock(block: L2BlockLike, finalizedBlock: L2BlockLike): JobResult[String] = {
+  private def confirmBlock(block: L2BlockLike, finalizedBlock: L2BlockLike): JobResult[PayloadStatus] = {
     val finalizedBlockHash = if (finalizedBlock.height > block.height) block.hash else finalizedBlock.hash
     engineApiClient.forkChoiceUpdate(block.hash, finalizedBlockHash)
   }
 
-  private def confirmBlock(hash: BlockHash, finalizedBlockHash: BlockHash): JobResult[String] =
+  private def confirmBlock(hash: BlockHash, finalizedBlockHash: BlockHash): JobResult[PayloadStatus] =
     engineApiClient.forkChoiceUpdate(hash, finalizedBlockHash)
 
   private def confirmBlockAndStartMining(
@@ -1486,7 +1486,7 @@ class ELUpdater(
       suggestedFeeRecipient: EthAddress,
       prevRandao: String,
       withdrawals: Vector[Withdrawal]
-  ): JobResult[String] = {
+  ): JobResult[PayloadId] = {
     val finalizedBlockHash = if (finalizedBlock.height > lastBlock.height) lastBlock.hash else finalizedBlock.hash
     engineApiClient
       .forkChoiceUpdateWithPayloadId(

diff --git a/src/main/scala/units/client/engine/EngineApiClient.scala b/src/main/scala/units/client/engine/EngineApiClient.scala
@@ -7,7 +7,7 @@ import units.eth.EthAddress
 import units.{BlockHash, JobResult}
 
 trait EngineApiClient {
-  def forkChoiceUpdate(blockHash: BlockHash, finalizedBlockHash: BlockHash): JobResult[String] // TODO Replace String with an appropriate type
+  def forkChoiceUpdate(blockHash: BlockHash, finalizedBlockHash: BlockHash): JobResult[PayloadStatus]
 
   def forkChoiceUpdateWithPayloadId(
       lastBlockHash: BlockHash,

diff --git a/src/main/scala/units/client/engine/HttpEngineApiClient.scala b/src/main/scala/units/client/engine/HttpEngineApiClient.scala
@@ -10,6 +10,7 @@ import units.client.engine.EngineApiClient.PayloadId
 import units.client.engine.HttpEngineApiClient.*
 import units.client.engine.model.*
 import units.client.engine.model.ForkChoiceUpdatedRequest.ForkChoiceAttributes
+import units.client.engine.model.PayloadStatus.{Syncing, Valid}
 import units.eth.EthAddress
 import units.{BlockHash, ClientConfig, ClientError, JobResult}
 
@@ -19,16 +20,16 @@ class HttpEngineApiClient(val config: ClientConfig, val backend: SttpBackend[Ide
 
   val apiUrl: Uri = uri"${config.executionClientAddress}"
 
-  def forkChoiceUpdate(blockHash: BlockHash, finalizedBlockHash: BlockHash): JobResult[String] = {
+  def forkChoiceUpdate(blockHash: BlockHash, finalizedBlockHash: BlockHash): JobResult[PayloadStatus] = {
     sendEngineRequest[ForkChoiceUpdatedRequest, ForkChoiceUpdatedResponse](
       ForkChoiceUpdatedRequest(blockHash, finalizedBlockHash, None),
       BlockExecutionTimeout
     )
       .flatMap {
-        case ForkChoiceUpdatedResponse(PayloadStatus(status, _, _), None) if status == "SYNCING" || status == "VALID" => Right(status)
-        case ForkChoiceUpdatedResponse(PayloadStatus(_, _, Some(validationError)), _) =>
+        case ForkChoiceUpdatedResponse(ps @ PayloadState(Valid | Syncing, _, _), None) => Right(ps.status)
+        case ForkChoiceUpdatedResponse(PayloadState(_, _, Some(validationError)), _) =>
           Left(ClientError(s"Payload validation error: $validationError"))
-        case ForkChoiceUpdatedResponse(payloadStatus, _) => Left(ClientError(s"Unexpected payload status ${payloadStatus.status}"))
+        case ForkChoiceUpdatedResponse(payloadState, _) => Left(ClientError(s"Unexpected payload status ${payloadState.status}"))
       }
   }
 
@@ -48,14 +49,14 @@ class HttpEngineApiClient(val config: ClientConfig, val backend: SttpBackend[Ide
       ),
       BlockExecutionTimeout
     ).flatMap {
-      case ForkChoiceUpdatedResponse(payloadStatus, Some(payloadId)) if payloadStatus.status == "VALID" =>
+      case ForkChoiceUpdatedResponse(PayloadState(Valid, _, _), Some(payloadId)) =>
         Right(payloadId)
       case ForkChoiceUpdatedResponse(_, None) =>
         Left(ClientError(s"Payload id for $lastBlockHash is not defined"))
-      case ForkChoiceUpdatedResponse(PayloadStatus(_, _, Some(validationError)), _) =>
+      case ForkChoiceUpdatedResponse(PayloadState(_, _, Some(validationError)), _) =>
         Left(ClientError(s"Payload validation error for $lastBlockHash: $validationError"))
-      case ForkChoiceUpdatedResponse(payloadStatus, _) =>
-        Left(ClientError(s"Unexpected payload status for $lastBlockHash: ${payloadStatus.status}"))
+      case ForkChoiceUpdatedResponse(payloadState, _) =>
+        Left(ClientError(s"Unexpected payload status for $lastBlockHash: ${payloadState.status}"))
     }
   }
 
@@ -64,12 +65,12 @@ class HttpEngineApiClient(val config: ClientConfig, val backend: SttpBackend[Ide
   }
 
   def applyNewPayload(payload: JsObject): JobResult[Option[BlockHash]] = {
-    sendEngineRequest[NewPayloadRequest, PayloadStatus](NewPayloadRequest(payload), BlockExecutionTimeout).flatMap {
-      case PayloadStatus(_, _, Some(validationError))                           => Left(ClientError(s"Payload validation error: $validationError"))
-      case PayloadStatus(status, Some(latestValidHash), _) if status == "VALID" => Right(Some(latestValidHash))
-      case PayloadStatus(status, latestValidHash, _) if status == "SYNCING"     => Right(latestValidHash)
-      case PayloadStatus(status, None, _) => Left(ClientError(s"Latest valid hash is not defined at status $status"))
-      case PayloadStatus(status, _, _)    => Left(ClientError(s"Unexpected payload status: $status"))
+    sendEngineRequest[NewPayloadRequest, PayloadState](NewPayloadRequest(payload), BlockExecutionTimeout).flatMap {
+      case PayloadState(_, _, Some(validationError))     => Left(ClientError(s"Payload validation error: $validationError"))
+      case PayloadState(Valid, Some(latestValidHash), _) => Right(Some(latestValidHash))
+      case PayloadState(Syncing, latestValidHash, _)     => Right(latestValidHash)
+      case PayloadState(status, None, _)                 => Left(ClientError(s"Latest valid hash is not defined at status $status"))
+      case PayloadState(status, _, _)                    => Left(ClientError(s"Unexpected payload status: $status"))
     }
   }
 

diff --git a/src/main/scala/units/client/engine/LoggedEngineApiClient.scala b/src/main/scala/units/client/engine/LoggedEngineApiClient.scala
@@ -13,9 +13,9 @@ import java.util.concurrent.ThreadLocalRandom
 import scala.util.chaining.scalaUtilChainingOps
 
 class LoggedEngineApiClient(underlying: EngineApiClient) extends EngineApiClient {
-  protected val log = LoggerFacade(LoggerFactory.getLogger(underlying.getClass))
+  protected val log: LoggerFacade = LoggerFacade(LoggerFactory.getLogger(underlying.getClass))
 
-  override def forkChoiceUpdate(blockHash: BlockHash, finalizedBlockHash: BlockHash): JobResult[String] =
+  override def forkChoiceUpdate(blockHash: BlockHash, finalizedBlockHash: BlockHash): JobResult[PayloadStatus] =
     wrap(s"forkChoiceUpdate($blockHash, f=$finalizedBlockHash)", underlying.forkChoiceUpdate(blockHash, finalizedBlockHash))
 
   override def forkChoiceUpdateWithPayloadId(

diff --git a/src/main/scala/units/client/engine/model/ForkChoiceUpdatedResponse.scala b/src/main/scala/units/client/engine/model/ForkChoiceUpdatedResponse.scala
@@ -3,7 +3,7 @@ package units.client.engine.model
 import units.client.engine.EngineApiClient.PayloadId
 import play.api.libs.json.{Json, Reads}
 
-case class ForkChoiceUpdatedResponse(payloadStatus: PayloadStatus, payloadId: Option[PayloadId])
+case class ForkChoiceUpdatedResponse(payloadStatus: PayloadState, payloadId: Option[PayloadId])
 
 object ForkChoiceUpdatedResponse {
   implicit val reads: Reads[ForkChoiceUpdatedResponse] = Json.reads

diff --git a/src/main/scala/units/client/engine/model/PayloadState.scala b/src/main/scala/units/client/engine/model/PayloadState.scala
@@ -0,0 +1,10 @@
+package units.client.engine.model
+
+import units.BlockHash
+import play.api.libs.json.{Json, Reads}
+
+case class PayloadState(status: PayloadStatus, latestValidHash: Option[BlockHash], validationError: Option[String])
+
+object PayloadState {
+  implicit val reads: Reads[PayloadState] = Json.reads
+}
diff --git a/src/main/scala/units/client/engine/model/PayloadStatus.scala b/src/main/scala/units/client/engine/model/PayloadStatus.scala
@@ -1,10 +1,19 @@
 package units.client.engine.model
 
-import units.BlockHash
-import play.api.libs.json.{Json, Reads}
+import play.api.libs.json.Reads
 
-case class PayloadStatus(status: String, latestValidHash: Option[BlockHash], validationError: Option[String])
+sealed abstract class PayloadStatus(val value: String) {
+  override def toString: String = value
+}
 
 object PayloadStatus {
-  implicit val reads: Reads[PayloadStatus] = Json.reads
+  case object Valid                                 extends PayloadStatus("VALID")
+  case object Syncing                               extends PayloadStatus("SYNCING")
+  case class Unexpected(override val value: String) extends PayloadStatus(value)
+
+  implicit val reads: Reads[PayloadStatus] = Reads.of[String].map {
+    case "VALID"   => Valid
+    case "SYNCING" => Syncing
+    case other     => Unexpected(other)
+  }
 }
diff --git a/src/test/scala/units/client/TestEcClients.scala b/src/test/scala/units/client/TestEcClients.scala
@@ -64,18 +64,18 @@ class TestEcClients private (
 
   val engineApi = new LoggedEngineApiClient(
     new EngineApiClient {
-      override def forkChoiceUpdate(blockHash: BlockHash, finalizedBlockHash: BlockHash): JobResult[String] = {
+      override def forkChoiceUpdate(blockHash: BlockHash, finalizedBlockHash: BlockHash): JobResult[PayloadStatus] = {
         knownBlocks.get().get(blockHash) match {
           case Some(cid) =>
             currChainIdValue.set(cid)
             chains.transform { chains =>
               chains.updated(cid, chains(cid).dropWhile(_.hash != blockHash))
             }
             log.debug(s"Curr chain: ${currChain.map(_.hash).mkString(" <- ")}")
-            "VALID"
+            PayloadStatus.Valid
           case None =>
             log.warn(s"Can't find a block $blockHash during forkChoiceUpdate call")
-            "INVALID" // Generally this is wrong, but enough for now
+            PayloadStatus.Unexpected("INVALID") // Generally this is wrong, but enough for now
         }
       }.asRight
 
@@ -161,7 +161,7 @@ class TestEcClients private (
   )
 
   protected def notImplementedMethodJob[A](text: String): JobResult[A] = throw new NotImplementedMethod(text)
-  protected def notImplementedCase(text: String): Throwable      = new NotImplementedCase(text)
+  protected def notImplementedCase(text: String): Throwable            = new NotImplementedCase(text)
 }
 
 object TestEcClients {