Object storage

// CreateUpload starts a multipart upload for key. Nothing is visible under key until Complete.
func (s *ObjectStore) CreateUpload(key string) string {
  s.mu.Lock()
  defer s.mu.Unlock()
  s.nextID++
  id := "u" + strconv.Itoa(s.nextID)
  s.uploads[id] = &upload{key: key, parts: map[int][]byte{}}
  return id
}

// UploadPart stores one part. Parts can arrive in any order and in parallel; sending a part again
// replaces it. The client sends the SHA-256 of the part, and a corrupted part is refused.
func (s *ObjectStore) UploadPart(id string, n int, data []byte, sum string) error {
  if digest(data) != sum {
    return fmt.Errorf("part %d: %w", n, ErrBadChecksum)
  }
  s.mu.Lock()
  defer s.mu.Unlock()
  u, ok := s.uploads[id]
  if !ok {
    return fmt.Errorf("upload %s: %w", id, ErrNoUpload)
  }
  u.parts[n] = slices.Clone(data)
  return nil
}

// ListParts returns the part numbers the store holds, so a client can resume after a crash.
func (s *ObjectStore) ListParts(id string) ([]int, error) {
  s.mu.Lock()
  defer s.mu.Unlock()
  u, ok := s.uploads[id]
  if !ok {
    return nil, fmt.Errorf("upload %s: %w", id, ErrNoUpload)
  }
  nums := make([]int, 0, len(u.parts))
  for n := range u.parts {
    nums = append(nums, n)
  }
  slices.Sort(nums)
  return nums, nil
}

// Complete joins the listed parts in order into one object. The object appears all at once.
func (s *ObjectStore) Complete(id string, parts []int) (Object, error) {
  s.mu.Lock()
  defer s.mu.Unlock()
  u, ok := s.uploads[id]
  if !ok {
    return Object{}, fmt.Errorf("upload %s: %w", id, ErrNoUpload)
  }
  var data []byte
  for i, n := range parts {
    p, ok := u.parts[n]
    if !ok {
      return Object{}, fmt.Errorf("part %d: %w", n, ErrMissingPart)
    }
    if i < len(parts)-1 && len(p) < s.MinPart {
      return Object{}, fmt.Errorf("part %d is %d bytes: %w", n, len(p), ErrSmallPart)
    }
    data = append(data, p...)
  }
  o := Object{Data: data, SHA256: digest(data), ContentType: "application/octet-stream"}
  s.objects[u.key] = o
  delete(s.uploads, id)
  return o, nil
}

// place picks a node for one more fragment of v. Spread only takes a node whose domain holds
// fewer than Cap fragments of v, and prefers the domain that holds the fewest. Both policies
// skip down nodes and nodes that already hold a fragment of v.
func (c *Cluster) place(v *Version) (int, error) {
  used := map[int]bool{}
  inDomain := make([]int, c.Domains)
  for _, f := range v.Frags {
    if f.Node >= 0 {
      used[f.Node] = true
      inDomain[c.Nodes[f.Node].Domain]++
    }
  }
  var cands []int
  best := c.Code.K + c.Code.M + 1
  for i, n := range c.Nodes {
    if !n.Up || used[i] {
      continue
    }
    if c.Policy == Spread {
      d := inDomain[n.Domain]
      if d >= c.Cap() || d > best {
        continue
      }
      if d < best { // a domain with fewer fragments: start the list again
        best, cands = d, cands[:0]
      }
    }
    cands = append(cands, i)
  }
  if len(cands) == 0 {
    return -1, fmt.Errorf("fragment of %s: %w", v.Key, ErrNoPlacement)
  }
  return cands[c.rng.IntN(len(cands))], nil
}

// Put encodes data into K+M fragments and writes each to its own node. The version is not
// visible until Commit: a read in between still returns the previous version.
func (c *Cluster) Put(key string, data []byte) (*Version, error) {
  shards := c.Code.Encode(data)
  c.nextID++
  v := &Version{Key: key, ID: c.nextID, Size: len(data), FragSize: len(shards[0]), Frags: make([]Frag, len(shards))}
  for i := range v.Frags {
    v.Frags[i].Node = -1
  }
  for i, s := range shards {
    n, err := c.place(v)
    if err != nil {
      c.drop(v) // fragments already written are garbage now
      return nil, err
    }
    c.Nodes[n].frags[fragID{v.ID, i}] = s
    v.Frags[i] = Frag{Node: n, Sum: sha256.Sum256(s)}
  }
  c.versions[v.ID] = v
  return v, nil
}

// Commit makes v the latest version of its key: one write to the index. The version it
// replaces stays on the nodes until GC.
func (c *Cluster) Commit(v *Version) {
  v.Committed = true
  c.latest[v.Key] = v
}

// Get reads the latest version of key. It reads the K data fragments; for each one it cannot
// read, it reads a parity fragment instead and decodes. With more than M fragments gone it
// returns ErrUnreadable, never wrong bytes.
func (c *Cluster) Get(key string) ([]byte, error) {
  v, ok := c.latest[key]
  if !ok || v.Deleted {
    return nil, fmt.Errorf("%s: %w", key, ErrNotFound)
  }
  shards, degraded, err := c.readK(v)
  if err != nil {
    return nil, err
  }
  c.Traffic.Reads++
  c.Traffic.ReadBytes += int64(c.Code.K * v.FragSize)
  if degraded {
    c.Traffic.DegradedReads++
  }
  return slices.Concat(shards[:c.Code.K]...)[:v.Size], nil
}

// readK collects K good fragments, data fragments first, and rebuilds the rest.
func (c *Cluster) readK(v *Version) (shards [][]byte, degraded bool, err error) {
  shards = make([][]byte, len(v.Frags))
  have := 0
  for i := range v.Frags {
    if have == c.Code.K {
      break
    }
    if b, ok := c.fragment(v, i); ok {
      shards[i] = slices.Clone(b)
      have++
    } else if i < c.Code.K {
      degraded = true
    }
  }
  if have < c.Code.K {
    return nil, false, fmt.Errorf("%s: %d of %d fragments: %w", v.Key, have, len(v.Frags), ErrUnreadable)
  }
  if degraded {
    if err := c.Code.Reconstruct(shards); err != nil {
      return nil, false, fmt.Errorf("decode %s: %w", v.Key, err)
    }
  }
  return shards, degraded, nil
}

// Reconstruct fills in the missing (nil) shards. Any K surviving shards are K known rows of the
// encoding matrix times the data, so inverting those K rows recovers the data shards. The lost
// parity shards are then encoded again.
func (c *Code) Reconstruct(shards [][]byte) error {
  var rows matrix
  var have [][]byte
  for i, s := range shards {
    if s != nil && len(rows) < c.K {
      rows = append(rows, c.enc[i])
      have = append(have, s)
    }
  }
  if len(rows) < c.K {
    return fmt.Errorf("%d of %d shards left, need %d: %w", len(rows), c.K+c.M, c.K, ErrTooFewShards)
  }
  dec, err := rows.invert()
  if err != nil {
    return fmt.Errorf("decode: %w", err)
  }
  data := make([][]byte, c.K)
  for j := range c.K {
    data[j] = c.combine(dec[j], have)
  }
  for i := range shards {
    if shards[i] == nil {
      shards[i] = c.combine(c.enc[i], data)
    }
  }
  return nil
}

// Repair rebuilds every missing fragment of every committed version. For one version it reads
// K good fragments once, decodes, and writes each missing fragment to a new node that the
// placement rule allows. The old copy, if its node returns, is an orphan for GC.
func (c *Cluster) Repair() (RepairResult, error) {
  var res RepairResult
  for _, id := range c.ids() {
    v := c.versions[id]
    if !v.Committed || v.Deleted {
      continue
    }
    var gone []int
    for i := range v.Frags {
      if _, ok := c.fragment(v, i); !ok {
        gone = append(gone, i)
      }
    }
    if len(gone) == 0 {
      continue
    }
    res.Objects++
    if len(gone) > c.Code.M {
      res.Lost++
      continue
    }
    shards, _, err := c.readK(v)
    if err != nil {
      return res, err
    }
    if err := c.Code.Reconstruct(shards); err != nil { // fills the parity readK skipped
      return res, fmt.Errorf("rebuild %s: %w", v.Key, err)
    }
    c.Traffic.RepairRead += int64(c.Code.K * v.FragSize)
    for _, i := range gone {
      v.Frags[i].Node = -1 // the old nodes no longer count in the placement rule
    }
    for _, i := range gone {
      n, err := c.place(v)
      if err != nil {
        return res, err
      }
      c.Nodes[n].frags[fragID{v.ID, i}] = shards[i]
      v.Frags[i].Node = n
      c.Traffic.RepairWrite += int64(v.FragSize)
      c.Traffic.Rebuilt++
      res.Rebuilt++
    }
  }
  return res, nil
}

// Scrub reads every fragment on every healthy node and checks it against the checksum in the
// index. A fragment that fails is deleted, so the next repair rebuilds it from the others.
func (c *Cluster) Scrub() int {
  bad := 0
  for _, id := range c.ids() {
    v := c.versions[id]
    for i, f := range v.Frags {
      if f.Node < 0 || !c.Nodes[f.Node].Up {
        continue
      }
      key := fragID{v.ID, i}
      if b, ok := c.Nodes[f.Node].frags[key]; ok && sha256.Sum256(b) != f.Sum {
        delete(c.Nodes[f.Node].frags, key)
        bad++
      }
    }
  }
  return bad
}

// GC deletes what no reader can reach: every committed version that is not the latest of its
// key (overwritten, or hidden by a delete marker), and every fragment a node holds that the
// index places elsewhere (left behind when repair moved it). It returns the bytes freed.
func (c *Cluster) GC() int64 {
  var freed int64
  for _, id := range c.ids() {
    v := c.versions[id]
    if v.Committed && c.latest[v.Key] != v {
      for i, f := range v.Frags {
        if f.Node >= 0 {
          freed += int64(len(c.Nodes[f.Node].frags[fragID{v.ID, i}]))
        }
      }
      c.drop(v)
    }
  }
  for ni, n := range c.Nodes {
    if !n.Up {
      continue // swept when the node returns
    }
    for k, b := range n.frags {
      v, ok := c.versions[k.version]
      if !ok || v.Frags[k.idx].Node != ni {
        freed += int64(len(b))
        delete(n.frags, k)
      }
    }
  }
  return freed
}

// Durable works out the yearly chance to lose an object. The model: disks fail independently at
// rate afr a year, and a failed disk is rebuilt within repairHours. An object is lost only if more
// than M of its K+M disks fail inside one repair window.
func Durable(s Scheme, afr, repairHours float64) Durability {
  n := s.K + s.M
  p := afr * repairHours / 8760 // one disk, one window
  var perWindow float64
  for i := s.M + 1; i <= n; i++ {
    perWindow += binom(n, i) * math.Pow(p, float64(i)) * math.Pow(1-p, float64(n-i))
  }
  windows := 8760 / repairHours
  annual := windows * perWindow
  return Durability{
    Overhead:  float64(n) / float64(s.K),
    Tolerates: s.M,
    P:         p,
    Windows:   windows,
    Lead:      binom(n, s.M+1) * math.Pow(p, float64(s.M+1)),
    PerWindow: perWindow,
    Annual:    annual,
    Nines:     -math.Log10(annual),
  }
}

-- The newest version wins, whatever order two writers commit in.
WITH v AS (
  UPDATE versions SET committed = true WHERE id = $1 RETURNING bucket_id, key, id
)
INSERT INTO latest (bucket_id, key, version_id)
SELECT bucket_id, key, id FROM v
ON CONFLICT (bucket_id, key) DO UPDATE SET version_id = EXCLUDED.version_id
WHERE latest.version_id < EXCLUDED.version_id
RETURNING version_id;

-- One page of keys under a prefix, in key order, after the last key of the previous page.
SELECT l.key, v.size_bytes, v.etag
FROM latest l JOIN versions v ON v.id = l.version_id
WHERE l.bucket_id = $1
  AND l.key >= $2 AND ($3::text IS NULL OR l.key < $3)
  AND l.key > $4
  AND NOT v.is_delete
ORDER BY l.key
LIMIT $5;

-- Versions no reader can reach, older than the grace period: overwritten, hidden by a delete
-- marker, or written and never committed. A versioned bucket keeps its old versions.
SELECT v.id
FROM versions v JOIN buckets b ON b.id = v.bucket_id
WHERE NOT EXISTS (SELECT 1 FROM latest l WHERE l.version_id = v.id)
  AND v.created_at < now() - make_interval(secs => $1)
  AND (NOT b.versioned OR NOT v.committed)
ORDER BY v.id
LIMIT $2;

// Begin records a version and where its fragments went, in one transaction, after the data
// nodes acknowledged them. Nothing reads it until Commit.
func (m Meta) Begin(ctx context.Context, bucket int64, key string, size int64, etag string, frags []Placed) (int64, error) {
  var id int64
  err := pgx.BeginFunc(ctx, m.DB, func(tx pgx.Tx) error {
    if err := tx.QueryRow(ctx, stmts["begin_version"], bucket, key, size, etag, false).Scan(&id); err != nil {
      return fmt.Errorf("add version: %w", err)
    }
    for i, f := range frags {
      if _, err := tx.Exec(ctx, stmts["add_fragment"], id, i, f.Node, f.Checksum); err != nil {
        return fmt.Errorf("add fragment %d: %w", i, err)
      }
    }
    return nil
  })
  if err != nil {
    return 0, fmt.Errorf("begin %s: %w", key, err)
  }
  return id, nil
}

// Commit points the key at version id, unless a newer version already holds it. It returns
// false when a newer write won; the caller's version is then garbage.
func (m Meta) Commit(ctx context.Context, id int64) (bool, error) {
  var got int64
  err := m.DB.QueryRow(ctx, stmts["commit"], id).Scan(&got)
  if errors.Is(err, pgx.ErrNoRows) {
    return false, nil
  }
  if err != nil {
    return false, fmt.Errorf("commit version %d: %w", id, err)
  }
  return true, nil
}

// List returns up to limit keys that start with prefix and sort after the key `after`. The next
// page passes the last key of this one. Keys sort by byte value, so one prefix is one range of
// the primary key index.
func (m Meta) List(ctx context.Context, bucket int64, prefix, after string, limit int) ([]Entry, error) {
  var upper *string
  if u, ok := prefixEnd(prefix); ok {
    upper = &u
  }
  rows, err := m.DB.Query(ctx, stmts["list"], bucket, prefix, upper, after, limit)
  if err != nil {
    return nil, fmt.Errorf("list %q: %w", prefix, err)
  }
  var out []Entry
  for rows.Next() {
    var e Entry
    var etag string
    if err := rows.Scan(&e.Key, &e.Size, &etag); err != nil {
      return nil, fmt.Errorf("list %q: %w", prefix, err)
    }
    out = append(out, e)
  }
  return out, rows.Err()
}

// prefixEnd returns the smallest string above every string that starts with prefix: the
// prefix with its last byte raised by one. An empty prefix has no end.
func prefixEnd(prefix string) (string, bool) {
  b := []byte(prefix)
  for i := len(b) - 1; i >= 0; i-- {
    if b[i] < 0x7f {
      b[i]++
      return string(b[:i+1]), true
    }
  }
  return "", false
}

// Garbage returns versions that no reader can reach and that are older than grace. The grace
// period protects a write between Begin and Commit.
func (m Meta) Garbage(ctx context.Context, grace time.Duration, limit int) ([]int64, error) {
  rows, err := m.DB.Query(ctx, stmts["garbage"], grace.Seconds(), limit)
  if err != nil {
    return nil, fmt.Errorf("find garbage: %w", err)
  }
  ids, err := pgx.CollectRows(rows, pgx.RowTo[int64])
  if err != nil {
    return nil, fmt.Errorf("find garbage: %w", err)
  }
  return ids, nil
}

// Purge deletes a version's rows after its fragments are gone from the nodes. Postgres refuses
// if a key still points to it.
func (m Meta) Purge(ctx context.Context, id int64) error {
  if _, err := m.DB.Exec(ctx, stmts["purge"], id); err != nil {
    return fmt.Errorf("purge version %d: %w", id, err)
  }
  return nil
}