llmgw/internal/mirror/peer.go

package mirror

import (
	"context"
	"log/slog"
	"path/filepath"
	"strings"
	"time"
)

// Mirrors kept in step, so the cluster has its models while one machine is
// down: each reads its peers every while, copies the versions a peer has
// that it lacks, checksum checked, and removes those a peer removed after
// they came here. Neither leads; what was done last stands.

// PeerInterval is how often each peer is read.
const PeerInterval = 30 * time.Second

// Peers keeps one mirror in step with others.
type Peers struct {
	M       *Manager
	Clients []*Client
	Log     *slog.Logger
}

// Run reads the peers until ctx ends.
func (p *Peers) Run(ctx context.Context) {
	// Files fetched before checksums were kept have none, and a peer copies
	// only what it can check: they are given theirs, beside the sync.
	go func() {
		if n := p.M.FillChecksums(ctx); n > 0 {
			p.Log.Info("recorded the checksums of files that had none, so peer mirrors can copy them", "files", n)
		}
	}()
	t := time.NewTicker(PeerInterval)
	defer t.Stop()
	for {
		for _, c := range p.Clients {
			p.SyncOnce(ctx, c)
		}
		select {
		case <-ctx.Done():
			return
		case <-t.C:
		}
	}
}

// SyncOnce brings this mirror in step with one peer.
func (p *Peers) SyncOnce(ctx context.Context, c *Client) {
	ctx, cancel := context.WithTimeout(ctx, 20*time.Second)
	defer cancel()
	theirs, err := c.List(ctx)
	if err != nil {
		p.Log.Warn("could not read a peer mirror; trying again shortly", "peer", c.Addr(), "err", err)
		return
	}
	mine, err := p.M.All()
	if err != nil {
		p.Log.Warn("could not read this mirror", "err", err)
		return
	}
	here := map[string]Status{}
	for _, st := range mine {
		here[st.Model+"/"+st.Version+"/"+st.File] = st
	}
	removed := p.M.Tombstones()

	// What the peer removed after it came here goes here too.
	for key, at := range theirs.Removed {
		model, version, ok := strings.Cut(key, "/")
		if !ok {
			continue
		}
		newest := time.Time{}
		has := false
		for _, st := range mine {
			if st.Model == model && st.Version == version {
				has = true
				if st.AddedAt.After(newest) {
					newest = st.AddedAt
				}
			}
		}
		if has && newest.UnixMilli() < at {
			if _, err := p.M.RemoveVersion(model, version); err == nil {
				p.Log.Info("removed a version a peer mirror removed", "peer", c.Addr(), "model", model, "version", version)
			}
		} else if !has && at > removed[key] {
			p.M.recordTombstone(model, version, time.UnixMilli(at))
		}
	}

	// Versions whole there and missing here: every file present, each with
	// its checksum, so a copy can be checked.
	type version struct {
		files []Status
		whole bool
	}
	versions := map[string]*version{}
	for _, st := range theirs.Models {
		key := st.Model + "/" + st.Version
		v := versions[key]
		if v == nil {
			v = &version{whole: true}
			versions[key] = v
		}
		v.files = append(v.files, st)
		v.whole = v.whole && st.State == Present && st.SHA256 != ""
	}
	for key, v := range versions {
		if !v.whole {
			continue
		}
		for _, st := range v.files {
			if _, ok := here[key+"/"+st.File]; ok {
				continue
			}
			// Removed here after the peer had it: not copied back. The
			// peer removes its own when it reads this one.
			if at, gone := removed[key]; gone && st.AddedAt.UnixMilli() < at {
				continue
			}
			source := strings.TrimRight(theirs.Base, "/") + "/" + st.Model + "/" + st.Version + "/" + st.File
			e := Entry{Model: st.Model, Version: st.Version, File: st.File}
			if _, err := p.M.startCopy(source, e, st.SHA256); err == nil {
				p.Log.Info("copying a version from a peer mirror", "peer", c.Addr(), "model", st.Model,
					"version", st.Version, "file", st.File)
			}
		}
	}
}

// FillChecksums records the sha256 of every file the mirror holds without
// one, and says how many it recorded.
func (m *Manager) FillChecksums(ctx context.Context) int {
	entries, err := List(m.root)
	if err != nil {
		return 0
	}
	n := 0
	for _, e := range entries {
		if ctx.Err() != nil {
			break
		}
		path := filepath.Join(m.root, e.Model, e.Version, e.File)
		if readChecksum(path) != "" || m.isRunning(e.Model, e.Version, e.File) {
			continue
		}
		sum, err := sha256File(path)
		if err != nil {
			m.log.Warn("could not read a file's checksum", "file", path, "err", err)
			continue
		}
		if err := writeChecksum(path, sum); err != nil {
			m.log.Warn("could not record a file's checksum", "file", path, "err", err)
			continue
		}
		n++
	}
	return n
}