Key-value store

// Preference walks the ring clockwise from the key's hash and lists each node once, in the order
// it meets them. The first N are the key's home replicas; the rest are stand-ins, in order.
func (c *Cluster) Preference(key string) []string {
  h := sharding.Hash(key)
  start := sort.Search(len(c.tokens), func(i int) bool { return c.tokens[i].Pos >= h })
  var out []string
  for i := range c.tokens {
    t := c.tokens[(start+i)%len(c.tokens)] // wrap past the last token
    if !slices.Contains(out, t.Node) {
      out = append(out, t.Node)
    }
    if len(out) == len(c.nodes) {
      break
    }
  }
  return out
}

// Join adds a node to the ring. Its tokens take over parts of other nodes' ranges. For every
// key whose home replicas changed, the new home copies the value from the old replicas (merged),
// and a node that is no longer a home replica drops its copy. It returns how many key copies
// moved and how many copies the cluster holds.
func (c *Cluster) Join(name string) (moved, copies int, err error) {
  if _, ok := c.nodes[name]; ok {
    return 0, 0, fmt.Errorf("join %s: already a member", name)
  }
  keys := c.Keys()
  before := map[string][]string{}
  for _, k := range keys {
    before[k] = c.Home(k)
  }
  c.addTokens(name)
  for _, k := range keys {
    merged := Value{}
    for _, h := range before[k] {
      merged = merged.Merge(c.nodes[h].Data[k])
    }
    after := c.Home(k)
    for _, h := range after {
      if !slices.Contains(before[k], h) {
        c.nodes[h].Data[k] = merged // streamed from the old replicas
        moved++
      }
    }
    for _, h := range before[k] {
      if !slices.Contains(after, h) {
        delete(c.nodes[h].Data, k)
      }
    }
    copies += len(after)
  }
  return moved, copies, nil
}

// Put writes val under key. ctx is the value the client read before (empty for a blind write):
// the new version supersedes everything in it. The coordinator is the first healthy home
// replica. It sends the write to each home replica; with a sloppy quorum, a down replica's copy
// goes to the next healthy node outside the N, with a hint. The write succeeds at W acks.
func (c *Cluster) Put(key, val string, ctx Value) (PutResult, error) {
  coord := c.firstUp(c.Preference(key))
  if coord == nil {
    return PutResult{}, fmt.Errorf("put %s: %w: every node is down", key, ErrNoQuorum)
  }
  return c.PutAt(coord.Name, key, val, ctx)
}

// PutAt is Put with a chosen coordinator, as when a load balancer sends the request to any node.
func (c *Cluster) PutAt(coordinator, key, val string, ctx Value) (PutResult, error) {
  pref := c.Preference(key)
  coord, err := c.Node(coordinator)
  if err != nil {
    return PutResult{}, fmt.Errorf("put %s: %w", key, err)
  }
  if !coord.Up {
    return PutResult{}, fmt.Errorf("put %s: coordinator %s is down: %w", key, coordinator, ErrNoQuorum)
  }
  vv := crdt.VV{}
  for _, e := range ctx.Entries {
    vv = vv.Join(e.VV)
  }
  for _, e := range coord.Data[key].Entries { // never reuse a counter the coordinator has issued
    vv[coord.Name] = max(vv[coord.Name], e.VV[coord.Name])
  }
  vv[coord.Name]++
  w := Value{Entries: []crdt.MVEntry[string]{{Value: val, VV: vv}}}

  res := PutResult{Coordinator: coord.Name, Version: vv}
  used := map[string]bool{}
  for _, home := range pref[:c.cfg.N] {
    if n := c.nodes[home]; n.Up {
      n.Data[key] = n.Data[key].Merge(w)
      res.Acks = append(res.Acks, Ack{Node: home})
      continue
    }
    if !c.cfg.Sloppy {
      continue
    }
    for _, s := range pref[c.cfg.N:] { // the next healthy node not used yet
      if n := c.nodes[s]; n.Up && !used[s] {
        used[s] = true
        n.Hints = append(n.Hints, Hint{For: home, Key: key, Val: w})
        res.Acks = append(res.Acks, Ack{Node: s, For: home})
        break
      }
    }
  }
  if len(res.Acks) < c.cfg.W {
    return res, fmt.Errorf("put %s: %w: %d of W=%d acks", key, ErrNoQuorum, len(res.Acks), c.cfg.W)
  }
  return res, nil
}

// Get reads key from the healthy home replicas. Replies arrive in preference order; the client
// gets the merge of the first R. The coordinator then merges every reply and writes the merge
// back to each replica that was behind (read repair).
func (c *Cluster) Get(key string) (GetResult, error) {
  var res GetResult
  for _, home := range c.Home(key) {
    if n := c.nodes[home]; n.Up {
      res.Replies = append(res.Replies, Reply{Node: home, Val: n.Data[key]})
    }
  }
  if len(res.Replies) < c.cfg.R {
    return res, fmt.Errorf("get %s: %w: %d of R=%d replies", key, ErrNoQuorum, len(res.Replies), c.cfg.R)
  }
  all := Value{}
  for i, r := range res.Replies {
    if i < c.cfg.R {
      res.Val = res.Val.Merge(r.Val)
    }
    all = all.Merge(r.Val)
  }
  for _, r := range res.Replies {
    if !r.Val.Equal(all) {
      c.nodes[r.Node].Data[key] = all
      res.Repaired = append(res.Repaired, r.Node)
    }
  }
  return res, nil
}

// Handoff runs on every healthy node: each hint whose home replica is up again is delivered,
// merged into that replica, and deleted. Hints for replicas still down stay.
func (c *Cluster) Handoff() []Ack {
  var done []Ack
  for _, name := range c.Names() {
    n := c.nodes[name]
    if !n.Up {
      continue
    }
    var keep []Hint
    for _, h := range n.Hints {
      home := c.nodes[h.For]
      if !home.Up {
        keep = append(keep, h)
        continue
      }
      home.Data[h.Key] = home.Data[h.Key].Merge(h.Val)
      done = append(done, Ack{Node: name, For: h.For})
    }
    n.Hints = keep
  }
  return done
}

// Put writes val under key. ctx is the value the client read before (empty for a blind write):
// the new version supersedes everything in it. The coordinator is the first healthy home
// replica. It sends the write to each home replica; with a sloppy quorum, a down replica's copy
// goes to the next healthy node outside the N, with a hint. The write succeeds at W acks.
func (c *Cluster) Put(key, val string, ctx Value) (PutResult, error) {
  coord := c.firstUp(c.Preference(key))
  if coord == nil {
    return PutResult{}, fmt.Errorf("put %s: %w: every node is down", key, ErrNoQuorum)
  }
  return c.PutAt(coord.Name, key, val, ctx)
}

// PutAt is Put with a chosen coordinator, as when a load balancer sends the request to any node.
func (c *Cluster) PutAt(coordinator, key, val string, ctx Value) (PutResult, error) {
  pref := c.Preference(key)
  coord, err := c.Node(coordinator)
  if err != nil {
    return PutResult{}, fmt.Errorf("put %s: %w", key, err)
  }
  if !coord.Up {
    return PutResult{}, fmt.Errorf("put %s: coordinator %s is down: %w", key, coordinator, ErrNoQuorum)
  }
  vv := crdt.VV{}
  for _, e := range ctx.Entries {
    vv = vv.Join(e.VV)
  }
  for _, e := range coord.Data[key].Entries { // never reuse a counter the coordinator has issued
    vv[coord.Name] = max(vv[coord.Name], e.VV[coord.Name])
  }
  vv[coord.Name]++
  w := Value{Entries: []crdt.MVEntry[string]{{Value: val, VV: vv}}}

  res := PutResult{Coordinator: coord.Name, Version: vv}
  used := map[string]bool{}
  for _, home := range pref[:c.cfg.N] {
    if n := c.nodes[home]; n.Up {
      n.Data[key] = n.Data[key].Merge(w)
      res.Acks = append(res.Acks, Ack{Node: home})
      continue
    }
    if !c.cfg.Sloppy {
      continue
    }
    for _, s := range pref[c.cfg.N:] { // the next healthy node not used yet
      if n := c.nodes[s]; n.Up && !used[s] {
        used[s] = true
        n.Hints = append(n.Hints, Hint{For: home, Key: key, Val: w})
        res.Acks = append(res.Acks, Ack{Node: s, For: home})
        break
      }
    }
  }
  if len(res.Acks) < c.cfg.W {
    return res, fmt.Errorf("put %s: %w: %d of W=%d acks", key, ErrNoQuorum, len(res.Acks), c.cfg.W)
  }
  return res, nil
}

// Set writes v at replica r with timestamp ts from r's clock.
func (l LWW[T]) Set(r string, v T, ts int64) (LWW[T], LWW[T]) {
  w := LWW[T]{Value: v, TS: ts, Replica: r}
  return l.Merge(w), w
}

// Merge keeps the later write. Equal timestamps fall back to the replica name, so every
// replica picks the same winner. The loser is gone, even if it happened later in real time.
func (l LWW[T]) Merge(o LWW[T]) LWW[T] {
  if o.TS > l.TS || (o.TS == l.TS && o.Replica > l.Replica) {
    return o
  }
  return l
}

// MV is a multi-value register. Each write carries a version vector. A write that has seen
// another write replaces it; two concurrent writes are both kept, as siblings, until a later
// write that has seen both replaces them.
type MV[T comparable] struct {
  Entries []MVEntry[T] `json:"entries"`
}

// MVEntry is one sibling: a value and the version vector of its write.
type MVEntry[T comparable] struct {
  Value T  `json:"value"`
  VV    VV `json:"vv"`
}

// Set writes v at replica r. Its vector is the join of every sibling's vector plus one for r,
// so it supersedes every sibling this replica has seen.
func (m MV[T]) Set(r string, v T) (MV[T], MV[T]) {
  vv := VV{}
  for _, e := range m.Entries {
    vv = vv.Join(e.VV)
  }
  vv[r]++
  w := MV[T]{Entries: []MVEntry[T]{{Value: v, VV: vv}}}
  return m.Merge(w), w
}

// Merge keeps every entry that no other entry's vector dominates.
func (m MV[T]) Merge(o MV[T]) MV[T] {
  all := append(slices.Clone(m.Entries), o.Entries...)
  var out []MVEntry[T]
  for i, e := range all {
    keep := true
    for j, f := range all {
      if i != j && (e.VV.Before(f.VV) || (e.VV.Equal(f.VV) && j < i)) {
        keep = false // dominated, or a duplicate of an earlier entry
        break
      }
    }
    if keep {
      out = append(out, e)
    }
  }
  sortEntries(out)
  return MV[T]{Entries: out}
}

// BuildMerkle hashes the given keys of one node into a tree of the given depth.
func BuildMerkle(n *Node, keys []string, depth int) *Merkle {
  m := &Merkle{Depth: depth, keys: make([][]string, 1<<depth)}
  for _, k := range keys {
    l := Leaf(k, depth)
    m.keys[l] = append(m.keys[l], k)
  }
  leaves := make([][32]byte, 1<<depth)
  for i, ks := range m.keys {
    slices.Sort(ks)
    h := sha256.New()
    for _, k := range ks {
      if v, ok := n.Data[k]; ok { // a key this node lacks adds nothing
        fmt.Fprintf(h, "%s:%s\n", k, canonical(v))
      }
    }
    copy(leaves[i][:], h.Sum(nil))
  }
  m.Levels = make([][][32]byte, depth+1)
  m.Levels[depth] = leaves
  for lv := depth - 1; lv >= 0; lv-- {
    below := m.Levels[lv+1]
    m.Levels[lv] = make([][32]byte, len(below)/2)
    for i := range m.Levels[lv] {
      m.Levels[lv][i] = sha256.Sum256(append(below[2*i][:], below[2*i+1][:]...))
    }
  }
  return m
}

// Compare walks both trees from the root. It descends only into children whose hashes differ,
// so matching subtrees cost one comparison each.
func Compare(a, b *Merkle) Diff {
  var d Diff
  var walk func(lv, i int)
  walk = func(lv, i int) {
    d.Compared++
    d.Path = append(d.Path, [2]int{lv, i})
    if a.Levels[lv][i] == b.Levels[lv][i] {
      return // the whole range matches
    }
    d.Differ = append(d.Differ, [2]int{lv, i})
    if lv == a.Depth {
      d.Leaves = append(d.Leaves, i)
      return
    }
    walk(lv+1, 2*i)
    walk(lv+1, 2*i+1)
  }
  walk(0, 0)
  return d
}

// AntiEntropy compares two replicas with Merkle trees and repairs only the differing leaves:
// for each key in them, both sides store the merge of the two values.
func (c *Cluster) AntiEntropy(a, b string, depth int) (SyncResult, error) {
  na, err := c.Node(a)
  if err != nil {
    return SyncResult{}, err
  }
  nb, err := c.Node(b)
  if err != nil {
    return SyncResult{}, err
  }
  shared := c.SharedKeys(a, b)
  ta, tb := BuildMerkle(na, shared, depth), BuildMerkle(nb, shared, depth)
  res := SyncResult{Diff: Compare(ta, tb), Shared: len(shared), LeafCount: 1 << depth}
  for _, l := range res.Diff.Leaves {
    for _, k := range ta.KeysIn(l) {
      res.Keys++
      va, vb := na.Data[k], nb.Data[k]
      if va.Equal(vb) {
        continue
      }
      m := va.Merge(vb)
      na.Data[k], nb.Data[k] = m, m
      res.Repaired = append(res.Repaired, k)
    }
  }
  return res, nil
}