registry.mbt

///|
/// A service endpoint (← go-zero's `discov` target): the host and port an
/// instance listens on, plus a routing weight (default `1`). The weight is
/// carried through registration and snapshots and is what `WeightedRoundRobin`
/// shares traffic by; round-robin and pick-first ignore it.
pub(all) struct Endpoint {
  host : String
  port : Int
  weight : Int
} derive(Eq, Debug)

///|
/// Build an endpoint; `weight` defaults to `1`, matching an unweighted instance.
pub fn Endpoint::new(host : String, port : Int, weight? : Int = 1) -> Endpoint {
  { host, port, weight, }
}

///|
/// The `host:port` dial string.
pub fn Endpoint::address(self : Endpoint) -> String {
  self.host + ":" + self.port.to_string()
}

///|
/// An in-memory service registry (← go-zero's etcd `discov` store, minus the
/// network): a two-level map of `service -> instance-id -> endpoint` and a
/// monotonic revision bumped on every mutation, mirroring etcd's store revision
/// so a watcher could detect change. Instance ids are `<service>/<n>`, the leaf of
/// the etcd key an instance would lease.
pub struct InMemoryRegistry {
  instances : Map[String, Map[String, Endpoint]]
  mut seq : Int
  mut revision : Int64
}

///|
/// A fresh, empty registry at revision `0`.
pub fn InMemoryRegistry::new() -> InMemoryRegistry {
  { instances: Map([]), seq: 0, revision: 0, }
}

///|
/// The store revision, incremented on each register/deregister — etcd's
/// mod-revision, the value a watcher compares against to see new state.
pub fn InMemoryRegistry::revision(self : InMemoryRegistry) -> Int64 {
  self.revision
}

///|
/// Register `endpoint` under `service` and return its instance key. Each call
/// mints a distinct key, so two instances of one service coexist, and bumps the
/// revision.
pub fn InMemoryRegistry::register(
  self : InMemoryRegistry,
  service : String,
  endpoint : Endpoint,
) -> String {
  let bucket = match self.instances.get(service) {
    Some(b) => b
    None => {
      let b : Map[String, Endpoint] = Map([])
      self.instances[service] = b
      b
    }
  }
  self.seq = self.seq + 1
  let key = service + "/" + self.seq.to_string()
  bucket[key] = endpoint
  self.revision = self.revision + 1L
  key
}

///|
/// Remove the instance at `key` from `service`. Returns `true` if it existed (and
/// bumps the revision), `false` if the service or key was unknown.
pub fn InMemoryRegistry::deregister(
  self : InMemoryRegistry,
  service : String,
  key : String,
) -> Bool {
  match self.instances.get(service) {
    Some(bucket) =>
      if bucket.contains(key) {
        bucket.remove(key)
        self.revision = self.revision + 1L
        true
      } else {
        false
      }
    None => false
  }
}

///|
/// The endpoints registered for `service`, in registration order.
pub fn InMemoryRegistry::resolve(
  self : InMemoryRegistry,
  service : String,
) -> Array[Endpoint] {
  match self.instances.get(service) {
    Some(bucket) => bucket.values().collect()
    None => []
  }
}

///|
/// Every service name with at least one live instance.
pub fn InMemoryRegistry::services(self : InMemoryRegistry) -> Array[String] {
  self.instances.keys().collect()
}

///|
/// Resolve `service` on the registry and pick one endpoint with `balancer` — the
/// resolve-then-balance step a zRPC client runs before each call. An etcd- or
/// consul-backed registry with the same `resolve` shape drops in unchanged.
pub fn resolve_one(
  registry : InMemoryRegistry,
  service : String,
  balancer : RoundRobin,
) -> Endpoint? {
  balancer.pick(registry.resolve(service))
}

///|
/// A round-robin balancer (← go-zero's `roundRobinBalancer`) over a resolved
/// endpoint set: successive `pick`s cycle through the instances, spreading load
/// evenly. Holds only a cursor, so it is cheap to keep per client.
pub struct RoundRobin {
  mut cursor : Int
}

///|
/// A round-robin balancer starting at the first instance.
pub fn RoundRobin::new() -> RoundRobin {
  { cursor: 0, }
}

///|
/// Pick the next endpoint in rotation, or `None` if the set is empty. The cursor
/// advances modulo the set size, so it stays valid as instances come and go.
pub fn RoundRobin::pick(
  self : RoundRobin,
  endpoints : Array[Endpoint],
) -> Endpoint? {
  let n = endpoints.length()
  if n == 0 {
    return None
  }
  let idx = self.cursor % n
  self.cursor = (self.cursor + 1) % n
  Some(endpoints[idx])
}

///|
/// A weighted balancer over `Endpoint::weight`, using smooth weighted
/// round-robin: every pick credits each instance with its own weight, serves the
/// highest-credited one, then charges it the total weight of the set. Across one
/// full cycle each instance is served exactly its share of the traffic, and the
/// picks interleave instead of arriving in runs — a weight-5 instance is not
/// handed five requests back to back.
///
/// Credit is keyed by `address()`, so an instance that leaves and returns
/// resumes where it was rather than jumping the queue, and two instances sharing
/// an address are treated as one. An instance whose weight is zero or negative is
/// never picked; a set where every weight is non-positive yields `None`.
pub struct WeightedRoundRobin {
  scores : Map[String, Int]
}

///|
/// A weighted balancer with no credit accrued yet.
pub fn WeightedRoundRobin::new() -> WeightedRoundRobin {
  { scores: Map([]), }
}

///|
/// Pick the next endpoint in weight order, or `None` if nothing is eligible.
pub fn WeightedRoundRobin::pick(
  self : WeightedRoundRobin,
  endpoints : Array[Endpoint],
) -> Endpoint? {
  let live : Array[Endpoint] = []
  let keys : Array[String] = []
  let mut total = 0
  for e in endpoints {
    if e.weight > 0 {
      live.push(e)
      keys.push(e.address())
      total = total + e.weight
    }
  }
  if live.length() == 0 {
    return None
  }
  let mut best = 0
  let mut best_score = 0
  for i = 0; i < live.length(); i = i + 1 {
    let score = self.scores.get(keys[i]).unwrap_or(0) + live[i].weight
    self.scores[keys[i]] = score
    if i == 0 || score > best_score {
      best = i
      best_score = score
    }
  }
  // credit for an instance no longer in the set would resurface stale on its
  // return, long after the weights it was earned under changed
  let stale : Array[String] = []
  for key, _ in self.scores {
    if !keys.contains(key) {
      stale.push(key)
    }
  }
  for key in stale {
    self.scores.remove(key)
  }
  self.scores[keys[best]] = best_score - total
  Some(live[best])
}

///|
/// Pick the first endpoint (← gRPC's `pick_first`), or `None` if the set is
/// empty. A stable choice that only moves when the head instance goes away.
pub fn pick_first(endpoints : Array[Endpoint]) -> Endpoint? {
  if endpoints.length() == 0 {
    None
  } else {
    Some(endpoints[0])
  }
}

///|
pub extend Endpoint with Debug::{to_repr}

///|
pub extend Endpoint with Eq::{not_equal, equal}