A Distributed System from scratch, with Scala 3 - Part 3: Job submission, worker scaling, and leader election & consensus with Raft

// Abstract
case class SampleBridgeFourJob[F[_]: Async](cfg: AssignedTaskConfig) extends BridgeFourJob[F]
// Concrete
case class IOBridgeFourJob(cfg: AssignedTaskConfig) extends BridgeFourJob[IO]

SampleBridgeFourJob(com.chollinger.bridgefour.shared.models.Task$AssignedTaskConfig,cats.effect.kernel.Async)

  def makeJob(className: String, cfg: AssignedTaskConfig): BridgeFourJob[F] = {
    val cls      = Class.forName(className)
    val ctor     = cls.getDeclaredConstructor(classOf[AssignedTaskConfig], classOf[Async[F]])
    val instance = ctor.newInstance(cfg, summon[Async[F]])
    instance.asInstanceOf[BridgeFourJob[F]]
  }

val cls = Class.forName(className)
val params: Array[Type] =
      cls.getAnnotatedInterfaces.head.getType.asInstanceOf[ParameterizedType].getActualTypeArguments

// deps: e.g. IO
def makeJob[F[_]](className: String, cfg: AssignedTaskConfig, deps: AnyRef*): BridgeFourJob[F] = {
  val cls = Class.forName(className)

  // Build out the arguments
  val args = cfg +: deps
  val argTypes = AssignedTaskConfig.getClass +: deps.map {
    case null => classOf[Object]
    case x    => x.getClass
  }

  // ...and search for a matching constructor
  val ctor = cls.getConstructors.find { c =>
    val paramTypes = c.getParameterTypes
    paramTypes.length == argTypes.length &&
    paramTypes.zip(argTypes).forall { case (p, a) => p.isAssignableFrom(a) }
  }.getOrElse {
    val argsStr = argTypes.map(_.getName).mkString(", ")
    throw new NoSuchMethodException(s"No constructor found for $className with args: $argsStr")
  }

  val instance = ctor.newInstance(args*)
  instance.asInstanceOf[BridgeFourJob[F]]
}
// ...
def myProgram[F[_]: Async](): F[Unit] = {
  // Async[F] as AnyRef
  makeJob(TestJobs.sampleJobClass, Config(), summon[Async[F]])
}

myProgram[IO]()

SampleBridgeFourJob[F[_]: Async](cfg: AssignedTaskConfig) extends BridgeFourJob[F]

case class IOBridgeFourJob(cfg: AssignedTaskConfig) extends BridgeFourJob[IO]

trait BridgeFourJobCreator[F[_]] {
  def makeJob(className: String, cfg: AssignedTaskConfig): BridgeFourJob[F] = {
    val cls      = Class.forName(className)
    val ctor     = cls.getDeclaredConstructor(classOf[AssignedTaskConfig])
    val instance = ctor.newInstance(cfg)
    instance.asInstanceOf[BridgeFourJob[F]]
  }
}

/** The base job interface. Since these jobs are summoned, they need a concrete effect type in their implementation
  *
  * @tparam F
  *   Effect, e.g. IO
  */
trait BridgeFourJob[F[_]: Async]

lazy val bridgefourShared = ProjectRef(
  uri("https://github.com/chollinger93/bridgefour.git"),
  "shared"
)

lazy val root = (project in file("."))
  .settings(
    commonSettings,
    name := "wordcount",
    libraryDependencies ++= sharedDeps,
  )
  .dependsOn(bridgefourShared)
lazy val root = (project in file("src"))
  .settings(
    commonSettings,
    name := "wordcount",
    libraryDependencies ++= sharedDeps,
  ).dependsOn(RootProject(uri("https://github.com/chollinger93/bridgefour.git")))

case class DelayedWordCountLeaderJob(job: JobDetails) extends LeaderJob[IO] {
  // ...logger, encoder
  private def readFileAsCsv(in: FilePath): IO[Map[String, Int]] = // ...

  def collectResults(): IO[Json] = for {
    counts <- allFiles.parTraverse(f => readFileAsCsv(f))
    res = counts.foldLeft(Map[String, Int]()) { case (m1, m2) =>
            m1 ++ m2.map { case (k, v) => k -> (v + m1.getOrElse(k, 0)) }
          }
  } yield res.asJson
}

lazy val dockerSettings = Seq(
  dockerUpdateLatest   := true,
  dockerBaseImage      := "eclipse-temurin:17-jdk",
  dockerExposedVolumes += "/jars",
  scriptClasspath     ++= Seq("/jars/*")
)

{
    "name": "Example job",
    "jobClass": "com.chollinger.bridgefour.jobs.example.DelayedWordCountBridgeFourJob",
    "input": "/tmp/in",
    "output": "/tmp/out",
    "userSettings": {
        "timeout": "2"
    }
}

2025-05-09 19:38:16 spren-01-1  | [io-compute-2] INFO  o.h.s.m.Logger - service raised an error: class java.lang.ClassNotFoundException

 cfg.workers.parTraverse(f: A => M[B]))

case class ClusterControllerImpl[F[_]: ThrowableMonadError: Concurrent: Async: Logger](
    // ...
    cfg: ServiceConfig,
    workerConfig: Persistence[F, WorkerId, WorkerConfig]

sealed trait WorkerCache[F[_]] {
  def get(id: WorkerId): F[Option[WorkerConfig]]
  def add(cfg: WorkerConfig): F[Unit]
  def list(): F[Map[WorkerId, WorkerConfig]]
}
// ...
def makeF[F[_]: Async](
    cfg: ServiceConfig,
    workers: Persistence[F, WorkerId, WorkerConfig]
): F[WorkerCache[F]] = ???

override def getClusterState(): F[ClusterState] =
  for {
    _ <- Logger[F].debug(s"Configured workers: ${cfg.workers}")
    state <- cfg.workers
    workerMap <- workerConfig.list()
    workers    = workerMap.values.toList
    _         <- Logger[F].debug(s"Configured workers: $workers")
    state <- workers
               .parTraverse(c => checkWorkerState(c))
    _      <- Logger[F].debug(s"Worker responses: $state")
    cluster = ClusterState(cfg.workers, state)
    cluster = ClusterState(workers, state)
    _      <- Logger[F].debug(s"Cluster: $cluster")
  } yield cluster

cfg.workers.find(_.id == wId)

private def getTaskStatus(tc: AssignedTaskConfig): F[(TaskId, ExecutionStatus)] = {
  val wId = tc.slotId.workerId
  for {
    cfg <- workers.get(wId)
    wCfg <- cfg match {
              case Some(c) => sF.blocking(c)
              case _       => err.raiseError(InvalidWorkerConfigException(s"Invalid worker config for id $wId"))
            }
// ...

def initWorkerState[F[_]: Async](
    cfg: ServiceConfig,
    workers: Persistence[F, WorkerId, WorkerConfig]
): F[Unit] = {
  cfg.workers
    .foldLeft(Async[F].unit) { case (_, w) =>
      workers.put(w.id, w)
    }
}

  case req @ POST -> Root / "cluster" / "addWorker" =>
    Ok(for {
      workerCfg <- req.as[WorkerConfig]
      res       <- clusterOverseer.addWorker(workerCfg)
    } yield res)

{
    "status": {
        "type": "Healthy"
    },
    "workers": {
       // ...
    },
    "slots": {
        "open": 2,
        "total": 2
    }
}

{
    "status": {
        "type": "Degraded"
    },
    "workers": {
// ..
        "2": {
            "id": 2,
            "slots": [],
            "status": {
                "type": "Dead"
            }
        }
    },
    "slots": {
        "open": 2,
        "total": 2
    }
}

"slots": {
    "open": 4,
    "total": 4
}

override def removeWorker(workerId: WorkerId): F[Boolean] = {
  for {
    _ <- Logger[F].info(s"Trying to remove worker: $workerId")
    w <- workers.remove(workerId)
  } yield w
}

spren-02:
  image: spren:latest
depends_on:
    - kaladin

services:
  kaladin:
    # ...
    depends_on:
      - spren-01
      - spren-02
      - spren-03

private def bgThread(): F[Unit] = for {
  _ <- err.handleErrorWith(updateClusterStatus().void)(t => Logger[F].error(s"Failed to update cluster status: $t"))
  _ <- err.handleErrorWith(rebalanceUnassignedTasks().void)(t => Logger[F].error(s"Failed to balance tasks: $t"))
  _ <- err.handleErrorWith(updateAllJobStates().void)(t => Logger[F].error(s"Failed to update job states: $t"))
  _ <- Logger[F].debug("Kaladin Background threads finished")
} yield ()

N = 2k + 1, k ∈ ℤ

trait RaftService[F[_]] {
  // Runs the loop for sending heartbeats as a leader and checking if we need an election as a follower/candidate
  def runFibers(): F[Unit]
  // Handles a vote request from other Candidates
  def handleVote(req: RequestVote): F[RequestVoteResponse]
  // Handle a heartbeat from the leader; this is used to reset the election timer
  def handleHeartbeat(req: HeartbeatRequest): F[Unit]
  // Returns the current view of the world
  def getState: F[RaftElectionState]
  // Handles HTTP requests
  def withLeaderRedirect(routes: HttpRoutes[F]): HttpRoutes[F]
}

case class RaftElectionState(
    ownId: Int,
    term: Int = 0,
    ownState: RaftState = Follower,
    votedFor: Option[Int] = None,
    lastElectionEpoch: Option[Long] = None,
    lastHeartbeatEpoch: Option[Long] = None,
    currentLeader: Option[Int] = None,
    peers: List[LeaderConfig] = List.empty
)

def make[F[_]: ThrowableMonadError: Async: Clock: Logger: Concurrent](
    client: Client[F],
    lock: Mutex[F],
    state: AtomicCell[F, RaftElectionState],
    cfg: RaftConfig = RaftConfig()
): RaftService[F]

class RaftServiceIntegrationTest extends CatsEffectSuite {
  private val testTimeout: FiniteDuration = 5.seconds

  private def raftNodeResource(l: LeaderConfig, leaders: List[LeaderConfig], cfg: RaftConfig): Resource[IO, Unit] =
    for {
      client   <- EmberClientBuilder.default[IO].build
      raftLock <- Resource.eval(Mutex[IO])
      raftState <- Resource.eval(
                     AtomicCell[IO].of(
                       RaftElectionState(
                         ownId = l.id,
                         peers = leaders.filter(_.id != l.id)
                       )
                     )
                   )
      raftSvc = RaftService.make[IO](client, raftLock, raftState, cfg)
      httpApp = RaftRoutes[IO](raftSvc).routes.orNotFound
      _ <- EmberServerBuilder
             .default[IO]
             .withHost(Host.fromString(l.host).get)
             .withPort(Port.fromInt(l.port).get)
             .withHttpApp(Http4sLogger.httpApp(false, false)(httpApp))
             .build
      _ <- Resource.make(raftSvc.runFibers())(_ => IO.unit)
    } yield ()

  private def testState(client: Client[IO], leaders: List[LeaderConfig]): IO[Unit] = {

    for {
      _ <- IO.sleep(1.second)
      resps <- leaders.traverse { l =>
                 val uri = Uri.unsafeFromString(s"http://${l.host}:${l.port}/raft/state")
                 val req = Request[IO](method = Method.GET, uri = uri)
                 logger.info(s"Testing state for ${l.id}") >> client.expect[RaftElectionState](req)
               }
      _ <- logger.info(s"Got response: $resps")
      // Exactly one leader
      _ = assertEquals(
            resps.count(_.ownState == Leader),
            1
          )
      // 2 Followers
      _ = assertEquals(
            resps.count(_.ownState == Follower),
            2
          )
    } yield ()
  }

  test("Rafts elect a leader in a reasonable timeframe") {
    val leaders = List( // ...
    val cfg = RaftConfig( // ...
    val prog: Resource[IO, Unit] = for {

      _ <- leaders.parTraverse(raftNodeResource(_, leaders, cfg))
    } yield ()

    for {
      _     <- prog.useForever.start
      _     <- IO.sleep(testTimeout)
      client = EmberClientBuilder.default[IO].build
      _ <- {
        client.use { client =>
          testState(client, leaders)
        }
      }
    } yield ()
  }

}

override def runFibers(): F[Unit] =
  for {
    _ <- Logger[F].info("Starting Raft service")
    _ <- heartbeatLoop().start
    _ <- electionLoop().start
  } yield ()

  private def heartbeatLoop(): F[Unit] = {
    for {
      now <- Temporal[F].realTime
      s   <- state.get
      _ <- s.ownState match {
             case Leader => sendHeartbeatToFollowers(now)
             case _      => Async[F].unit
           }
      _ <- Temporal[F].sleep(cfg.heartbeatInterval)
      _ <- heartbeatLoop()
    } yield ()
  }

@CalledLocked
private def startElection: F[Unit] = {
  for {
    now <- Temporal[F].realTime
    _   <- Logger[F].info("Starting election")
    // ...

 @PartiallyLocked
  override def handleVote(req: RequestVote): F[RequestVoteResponse] = {
    for {
      _  <- logger.debug(s"Received vote request from ${req.candidateId} for term ${req.term}")
      s  <- state.get
      ts <- Temporal[F].realTime

      // Update state in case we became a follower
      s <- if (req.term > s.term) {
             logger.debug(
               s"Candidate ${req.candidateId} has a higher term than us (${req.term}/${s.term}, accepting leader"
             )
               >> state.update(
                 _.copy(
                   ownState = Follower, lastElectionEpoch = Some(ts.toMillis), votedFor = None,
                   currentLeader = Some(req.candidateId), term = req.term
                 )
               ) >> state.get
             // No state change
           } else Async[F].blocking(s)
      votedForCandidate = s.votedFor
      ourTerm           = s.term
      voteGranted       = req.term == ourTerm && (votedForCandidate.isEmpty || votedForCandidate.get == req.candidateId)
      _ <-
        logger.debug(s"Our vote for ${req.candidateId}: $voteGranted, ourTerm: $ourTerm, theirTerm: ${req.term}, votedForCandidate: $votedForCandidate")
      term <- if (voteGranted) {
                lock.lock.surround {
                  state.update(
                    _.copy(
                      ownState = Follower, lastHeartbeatEpoch = Some(ts.toMillis),
                      lastElectionEpoch = Some(ts.toMillis), votedFor = Some(req.candidateId),
                      currentLeader = Some(req.candidateId), term = req.term
                    )
                  )
                }
                  >> Async[F].blocking(req.term)
              } else Async[F].blocking(ourTerm)
      s <- state.get
      _ <- logger.info(s"State after vote: $s")
    } yield RequestVoteResponse(
      term = term,
      voteGranted = voteGranted
    )
  }

  @CalledLocked
  private def startElection: F[Unit] = {
    for {
      now <- Temporal[F].realTime
      _   <- Logger[F].info("Starting election")
      // We're now a candidate for the new term - this is the only time the term gets updated
      _ <- state.update { s =>
             val newTerm = s.term + 1
             s.copy(
               ownState = Candidate,
               term = newTerm,
               votedFor = Some(s.ownId),
               lastElectionEpoch = Some(now.toMillis)
             )
           }
      s <- state.get
      voteReq = RequestVote(
                  term = s.term,
                  candidateId = s.ownId
                )
      _ <- Logger[F].debug(s"Requesting vote $voteReq from peers: ${s.peers}")
      res <- s.peers.parTraverse { lCfg =>
               val req =
                 Request[F](method = Method.POST, uri = Uri.unsafeFromString(s"${lCfg.uri()}/raft/requestVote"))
                   .withEntity(voteReq)
               Logger[F].debug(s"Requesting election on $lCfg") >> err
                 .handleErrorWith(client.expect[RequestVoteResponse](req))(e =>
                   Logger[F].error(e)(s"Election on leader ${lCfg.id} failed") >>
                     Async[F].blocking(
                       RequestVoteResponse(
                         term = s.term,
                         voteGranted = false
                       )
                     )
                 )
             }
      _ <- Logger[F].debug(s"Received votes: $res")
      // One vote for ourselves
      _ <- res.map(_.voteGranted).count(_ == true) + 1 match {
             case count if count > s.peers.size / 2 =>
               Logger[F].info(s"Received majority of votes, becoming leader") >>
                 state.update(_.copy(ownState = Leader, currentLeader = Some(s.ownId)))
             case _ =>
               Logger[F].debug(s"Did not receive majority of votes, remaining candidate")
           }
      s <- state.get
      _ <- Logger[F].debug(s"State after vote: $s")
    } yield ()
  }

#!/bin/bash

DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" >/dev/null 2>&1 && pwd)"
cd "$DIR/.."

function checkRaftState() {
  hasLeader=false
  for i in {0..2}; do
    raw=$(curl -s --location --request GET "http://localhost:655$i/raft/state")
    st=$(echo "$raw" | jq '.ownState | to_entries | .[].key')
    leader=$(echo "$raw" | jq '.currentLeader')
    term=$(echo "$raw" | jq '.term')
    echo "$(date +%s),$i,$st,$leader,$term"
    if [[ $st == *"Leader"* ]]; then
      hasLeader=true
    fi
  done
  if [[ -n $RUN_FOREVER ]]; then
    sleep 1
    checkRaftState
  fi
  if [[ $hasLeader == false ]]; then
    sleep 1
    checkRaftState
  fi
}

docker compose up -d
echo "ts,node,state,leader,term"
checkRaftState

{
    "ownId": 1,
    "term": 1,
    "ownState": {
        "Leader": {}
    },
    "votedFor": 1,
    "lastElectionEpoch": 1747601542995,
    "lastHeartbeatEpoch": null,
    "currentLeader": 1,
    "peers": [
        {
            "id": 1,
            "schema": "http",
            "host": "kaladin-01",
            "port": 6551
        },
        {
            "id": 2,
            "schema": "http",
            "host": "kaladin-02",
            "port": 6552
        }
    ]
}

{
    "ownId": 0,
    "term": 1,
    "ownState": {
        "Follower": {}
    },
    "votedFor": 1,
    "lastElectionEpoch": 1747601543035,
    "lastHeartbeatEpoch": 1747601559492,
    "currentLeader": 0
}

def withLeaderRedirect(routes: HttpRoutes[F]): HttpRoutes[F] = HttpRoutes.of[F] { req =>
        if (req.uri.path.renderString.contains("/raft/")) {
          logger.debug(s"Request: ${req.uri.path.renderString}") >>
            routes(req).value.flatMap {
              case Some(response) => response.pure[F]
              case None           => Response[F](Status.NotFound).pure[F]
            }
        } else {
          state.get.flatMap {
            case s if s.ownState == Leader =>
              routes(req).value.flatMap {
                case Some(response) => response.pure[F]
                case None           => Response[F](Status.NotFound).pure[F]
              }
            case s =>
              s.currentLeader match {
                case Some(leaderId) =>
                  s.peers.find(_.id == leaderId) match {
                    case Some(leaderCfg) =>
                      val redirectUri = {
                        Uri.unsafeFromString(
                          s"${leaderCfg.schema}://${leaderCfg.host}:${leaderCfg.port}${req.uri.path.renderString}"
                        )
                      Response[F](Status.TemporaryRedirect)
                        .putHeaders(Header.Raw(ci"Location", redirectUri.renderString))
                        .pure[F]
                    case None =>
                      Response[F](Status.ServiceUnavailable)
                        .withEntity("Leader config missing")
                        .pure[F]
                  }
                case None =>
                  Response[F](Status.ServiceUnavailable)
                    .withEntity("No known leader")
                    .pure[F]
              }
          }
        }
      }
    }

  leaderRoutes = LeaderRoutes[F](jobController, healthMonitor).routes
  ratftRoutes  = RaftRoutes[F](raftSvc).routes
  // Ensure Raft wraps the main routes, i.e. we always talk to the leader
  httpApp: Kleisli[F, Request[F], Response[F]] =
    raftSvc.withLeaderRedirect(leaderRoutes <+> ratftRoutes).orNotFound