worker/internal/workercluster/cluster.go

package workercluster

import (
	"context"
	"errors"
	"fmt"
	"io"
	"log"
	"net"
	"net/http"
	"os"
	"sync"
	"time"

	rafthttp "github.com/CanonicalLtd/raft-http"
	raftmembership "github.com/CanonicalLtd/raft-membership"
	"github.com/hashicorp/go-hclog"
	"github.com/hashicorp/raft"
)

// Peer is one other worker this node knows about at bootstrap time. It
// is also the shape used by callers that drive membership changes
// outside the bootstrap path.
type Peer struct {
	WorkerID string
	Address  string // host:port of the rafthttp endpoint
}

// Options configures a single Cluster. Defaults are applied for any
// zero-valued field; see applyDefaults for the rules.
type Options struct {
	// NodeID is the Raft ServerID this node uses for itself. Required.
	NodeID string

	// LocalAddr is the host:port this node binds for the rafthttp
	// endpoint. Required.
	LocalAddr string

	// DataDir is the directory the bbolt store and snapshot files live
	// in. Created if missing. Required.
	DataDir string

	// Creds are the basic-auth credentials shared by all voters. They
	// gate the inbound rafthttp endpoint and are required on outbound
	// dials once a future rafthttp auth hook lands. Required.
	Creds HTTPCreds

	// RaftPath is the URL path the rafthttp handler mounts on. Default
	// "/raft".
	RaftPath string

	// HeartbeatTimeout / ElectionTimeout control the raft timing.
	// Defaults are 1s / 3s which keeps the e2e tests snappy.
	HeartbeatTimeout time.Duration
	ElectionTimeout  time.Duration

	// Logger is the destination for cluster log lines. Defaults to
	// os.Stderr.
	Logger *log.Logger

	// Bootstrap seeds a one-voter cluster on first start when no
	// existing state is present. Set Bootstrap=true on the first node
	// and false on every subsequent node that joins an existing
	// cluster via Seed.
	Bootstrap bool

	// Seed is the address (host:port of the rafthttp endpoint) of an
	// existing voter this node should join before it can become a
	// voter itself. Leave empty for the bootstrap node.
	Seed Peer

	// LogOutput is the destination for rafthttp log lines. Defaults to
	// io.Discard so the test runs stay quiet.
	LogOutput io.Writer
}

// Stats is a small read-only view of the cluster's runtime state.
type Stats struct {
	NodeID    string
	LocalAddr string
	State     string
	Leader    string
	Term      uint64
	AppliedIx uint64
	LastIx    uint64
	NumPeers  int
	Members   []Member
}

// Cluster is the Raft cluster wrapper for one node. It owns the
// bbolt store, the FSM, the rafthttp transport, the raft.Raft instance,
// and the HTTP server.
type Cluster struct {
	opts Options
	log  *log.Logger

	store   *BoltStore
	fsm     *FSM
	raft    *raft.Raft
	layer   *rafthttp.Layer
	handler *rafthttp.Handler

	listener net.Listener
	server   *http.Server

	membershipWG sync.WaitGroup

	mu     sync.Mutex
	closed bool
}

// New constructs a Cluster but does not start it. Call Start to bind
// the listener and bootstrap / join the Raft group.
func New(opts *Options) (*Cluster, error) {
	if err := opts.validate(); err != nil {
		return nil, err
	}
	opts.applyDefaults()

	return &Cluster{
		opts:    *opts,
		log:     opts.Logger,
		store:   nil, // opened in Start
		fsm:     NewFSM(),
		handler: nil, // built in Start
	}, nil
}

// Start binds the listener, constructs the transport, opens the store,
// and either bootstraps a fresh cluster or joins the seed peer.
func (c *Cluster) Start(ctx context.Context) error {
	c.mu.Lock()
	defer c.mu.Unlock()
	if c.closed {
		return errors.New("workercluster: cluster is shut down")
	}

	listener, err := net.Listen("tcp", c.opts.LocalAddr)
	if err != nil {
		return fmt.Errorf("workercluster: listen %q: %w", c.opts.LocalAddr, err)
	}
	c.listener = listener

	store, err := NewBoltStore(c.opts.DataDir)
	if err != nil {
		_ = listener.Close()
		return fmt.Errorf("workercluster: open store: %w", err)
	}
	c.store = store

	c.handler = rafthttp.NewHandler()
	dial := AuthDial(rafthttp.NewDialTCP(), c.opts.Creds)

	authHandler := NewAuthHandler(c.handler, c.opts.Creds, c.log)

	layer, server, err := NewTransport(
		c.opts.RaftPath,
		listener,
		authHandler,
		dial,
		c.opts.LogOutput,
	)
	if err != nil {
		_ = store.Close()
		_ = listener.Close()
		return err
	}
	c.layer = layer
	c.server = server

	go func() {
		if err := server.Serve(listener); err != nil && !errors.Is(err, http.ErrServerClosed) {
			c.log.Printf("workercluster: http server: %v", err)
		}
	}()

	transport := raft.NewNetworkTransport(
		layer,
		2,
		10*time.Second,
		c.opts.LogOutput,
	)

	config := raft.DefaultConfig()
	config.LocalID = raft.ServerID(c.opts.NodeID)
	config.HeartbeatTimeout = c.opts.HeartbeatTimeout
	config.ElectionTimeout = c.opts.ElectionTimeout
	// LeaderLeaseTimeout must be < HeartbeatTimeout; if the caller
	// overrode the heartbeat to a small value, default the lease to a
	// safely smaller value.
	if config.LeaderLeaseTimeout == 0 || config.LeaderLeaseTimeout >= config.HeartbeatTimeout {
		config.LeaderLeaseTimeout = config.HeartbeatTimeout / 2
	}
	config.Logger = hclog.New(&hclog.LoggerOptions{
		Name:   "raft",
		Output: c.opts.LogOutput,
		Level:  hclog.DefaultLevel,
	})

	r, err := raft.NewRaft(config, c.fsm, store.LogStore(), store.StableStore(), store.SnapshotStore(), transport)
	if err != nil {
		return c.shutdownLocked(fmt.Errorf("workercluster: new raft: %w", err))
	}
	c.raft = r

	c.membershipWG.Add(1)
	go func() {
		defer c.membershipWG.Done()
		raftmembership.HandleChangeRequests(r, c.handler.Requests())
	}()

	if c.opts.Bootstrap {
		if err := c.bootstrapLocked(); err != nil {
			return c.shutdownLocked(err)
		}
	} else if (c.opts.Seed != Peer{}) {
		if err := c.joinSeedLocked(ctx); err != nil {
			return c.shutdownLocked(err)
		}
	}
	return nil
}

// Raft returns the underlying *raft.Raft. Exposed for callers that need
// access to Apply, GetConfiguration, Stats, Barrier, Snapshot, etc.
func (c *Cluster) Raft() *raft.Raft { return c.raft }

// FSM returns the cluster's FSM. Read-only: callers must not mutate the
// FSM directly outside of Apply.
func (c *Cluster) FSM() *FSM { return c.fsm }

// Stats returns a snapshot of the cluster's runtime state. Safe for
// concurrent callers.
func (c *Cluster) Stats() Stats {
	c.mu.Lock()
	defer c.mu.Unlock()

	out := Stats{
		NodeID:    c.opts.NodeID,
		LocalAddr: c.opts.LocalAddr,
		Members:   membersAsSlice(c.fsm.Membership),
	}
	if c.raft == nil {
		out.State = "uninitialized"
		return out
	}
	out.State = c.raft.State().String()
	// Leader is reported as the raft.ServerAddress the library
	// returns (a host:port). Callers that want a worker NodeID
	// should go through ClusterStats, which does the address→ID
	// lookup with a guard against stale leader caches (a follower
	// that has not yet observed a new term still reports the old
	// leader's address, and we don't want to surface that as the
	// canonical leader_id).
	out.Leader = string(c.raft.Leader())
	out.AppliedIx = c.raft.AppliedIndex()
	out.LastIx = c.raft.LastIndex()
	out.Term = parseTerm(c.raft.Stats()["term"])
	if cfgFuture := c.raft.GetConfiguration(); cfgFuture.Error() == nil {
		out.NumPeers = len(cfgFuture.Configuration().Servers)
	}

	return out
}

func parseTerm(raw string) uint64 {
	var v uint64
	for i := 0; i < len(raw); i++ {
		if raw[i] < '0' || raw[i] > '9' {
			break
		}
		v = v*10 + uint64(raw[i]-'0')
	}
	return v
}

// Apply submits a log entry to the cluster. Blocks until the entry has
// been applied (or ctx is canceled).
func (c *Cluster) Apply(_ context.Context, payload []byte, timeout time.Duration) error {
	c.mu.Lock()
	r := c.raft
	c.mu.Unlock()
	if r == nil {
		return errors.New("workercluster: not started")
	}

	future := r.Apply(payload, timeout)
	return future.Error()
}

// Join asks the seed peer (or the leader it redirects to) to add this
// node as a voter.
func (c *Cluster) Join(_ context.Context, _ Peer, id raft.ServerID, addr raft.ServerAddress, timeout time.Duration) error {
	c.mu.Lock()
	layer := c.layer
	c.mu.Unlock()
	if layer == nil {
		return errors.New("workercluster: not started")
	}
	return layer.Join(id, addr, timeout)
}

// Leave asks the seed peer (or the leader it redirects to) to remove
// this node from the cluster.
func (c *Cluster) Leave(_ context.Context, peer Peer, id raft.ServerID, timeout time.Duration) error {
	c.mu.Lock()
	layer := c.layer
	c.mu.Unlock()
	if layer == nil {
		return errors.New("workercluster: not started")
	}
	return layer.Leave(id, raft.ServerAddress(peer.Address), timeout)
}

// Snapshot asks the leader to take a snapshot now. The returned error
// is the snapshot future's error.
func (c *Cluster) Snapshot() error {
	c.mu.Lock()
	r := c.raft
	c.mu.Unlock()
	if r == nil {
		return errors.New("workercluster: not started")
	}
	return r.Snapshot().Error()
}

// ClusterID returns the NodeID this cluster was constructed with. It
// is exposed so HTTP handlers can label status responses with a
// stable identifier even when the raft library's own State() reports
// "uninitialized" (e.g. before bootstrap completes).
func (c *Cluster) ClusterID() string { return c.opts.NodeID }

// LocalAddr returns the rafthttp bind address this cluster is using.
// Exposed for the cluster status endpoint.
func (c *Cluster) LocalAddr() string { return c.opts.LocalAddr }

// ClusterStats is the cluster stats shape used by the webapp admin
// handlers. It mirrors the JSON the GET /web/api/cluster/status
// endpoint returns; keeping the mapping in one place stops drift
// between Stats() and the wire format.
//
// FSMConfigVersion / FSMOutboxLen / FSMPartition carry the FSM-side
// operator signals from plan section 6.1 (config_version, outbox
// length, partition_state). They let the JSON endpoint surface
// "what config the cluster has adopted", "how many notifications
// are still pending in the outbox" and "what the cluster thinks of
// network partition state" without the operator having to scrape a
// raft log directly.
type ClusterStats struct {
	NodeID           string
	LocalAddr        string
	State            string
	Leader           string
	Term             uint64
	AppliedIndex     uint64
	LastIndex        uint64
	NumPeers         int
	Voters           []string
	FSMChecks        int
	FSMMembers       int
	FSMConfigVersion uint64
	FSMOutboxLen     int
	FSMPartition     string
}

// ClusterStats returns a view of the cluster suitable for the webapp
// admin endpoint. It calls Stats() internally and folds in the FSM
// counters so the handler does not need to know about the FSM type.
func (c *Cluster) ClusterStats() ClusterStats {
	s := c.Stats()
	fsmStats := c.FSM().Stats()
	return ClusterStats{
		NodeID:           s.NodeID,
		LocalAddr:        s.LocalAddr,
		State:            s.State,
		Leader:           leaderIDForAddr(c, s.Leader),
		Term:             s.Term,
		AppliedIndex:     s.AppliedIx,
		LastIndex:        s.LastIx,
		NumPeers:         s.NumPeers,
		Voters:           votersFromClusterStats(c, &s),
		FSMChecks:        fsmStats.ConfigCount,
		FSMMembers:       fsmStats.Members,
		FSMConfigVersion: fsmStats.ConfigVersion,
		FSMOutboxLen:     fsmStats.OutboxLen,
		FSMPartition:     fsmStats.Partition,
	}
}

// leaderIDForAddr translates a raft.ServerAddress (host:port) the
// library returns for raft.Leader() into a worker NodeID by matching
// against the current raft configuration. The empty string is
// returned unchanged so callers can detect "no leader yet".
//
// IMPORTANT: raft.Leader() on a follower can return the address of
// the previous leader until the follower observes the new term via
// an AppendEntries. We rely on the raft configuration to do the
// translation; the configuration is updated on every membership
// change but lags the Leader() field slightly. Callers that need a
// strict guarantee should treat the result as "best-effort" — it is
// good enough for an operator-facing status endpoint, not for safety
// decisions.
func leaderIDForAddr(c *Cluster, addr string) string {
	if addr == "" {
		return ""
	}
	c.mu.Lock()
	r := c.raft
	c.mu.Unlock()
	if r == nil {
		return addr
	}
	cfgFuture := r.GetConfiguration()
	if cfgFuture.Error() != nil {
		return addr
	}
	for _, srv := range cfgFuture.Configuration().Servers {
		if string(srv.Address) == addr {
			return string(srv.ID)
		}
	}
	return addr
}

// votersFromClusterStats returns the voter WorkerIDs currently
// registered with the raft library. The FSM membership cache is the
// primary source (it mirrors what the library stores, with a denorm
// for fast reads); when the cache is empty (e.g. right after a fresh
// bootstrap, before any membership.propose_add entry has been
// committed) we fall back to the raft configuration directly so the
// endpoint never reports an empty voter list on a healthy cluster.
func votersFromClusterStats(c *Cluster, s *Stats) []string {
	out := make([]string, 0, len(s.Members))
	for _, m := range s.Members {
		if m.Role == RoleVoter {
			out = append(out, m.WorkerID)
		}
	}
	if len(out) > 0 {
		return out
	}
	c.mu.Lock()
	r := c.raft
	c.mu.Unlock()
	if r == nil {
		return out
	}
	cfgFuture := r.GetConfiguration()
	if cfgFuture.Error() != nil {
		return out
	}
	for _, srv := range cfgFuture.Configuration().Servers {
		if srv.Suffrage == raft.Voter {
			out = append(out, string(srv.ID))
		}
	}
	return out
}

// Shutdown closes the HTTP server, the rafthttp handler, and the
// underlying raft instance. After Shutdown returns the Cluster cannot be
// reused.
func (c *Cluster) Shutdown(_ context.Context) error {
	c.mu.Lock()
	defer c.mu.Unlock()
	if c.closed {
		return nil
	}
	return c.shutdownLocked(nil)
}

// shutdownLocked tears the cluster down. err is returned to the caller.
func (c *Cluster) shutdownLocked(retErr error) error {
	c.closed = true
	if c.raft != nil {
		_ = c.raft.Shutdown().Error()
	}
	c.membershipWG.Wait()
	if c.server != nil {
		shutdownErr := c.server.Shutdown(context.Background())
		if shutdownErr != nil && retErr == nil {
			retErr = shutdownErr
		}
	}
	if c.store != nil {
		_ = c.store.Close()
	}
	if c.listener != nil {
		_ = c.listener.Close()
	}
	return retErr
}

func (o *Options) validate() error {
	if o.NodeID == "" {
		return errors.New("workercluster: NodeID is required")
	}
	if o.LocalAddr == "" {
		return errors.New("workercluster: LocalAddr is required")
	}
	if o.DataDir == "" {
		return errors.New("workercluster: DataDir is required")
	}
	if !o.Creds.IsConfigured() {
		return errors.New("workercluster: Creds must have both login and password set")
	}
	return nil
}

func (o *Options) applyDefaults() {
	if o.RaftPath == "" {
		o.RaftPath = "/raft"
	}
	if o.HeartbeatTimeout == 0 {
		o.HeartbeatTimeout = 1000 * time.Millisecond
	}
	if o.ElectionTimeout == 0 {
		o.ElectionTimeout = 3000 * time.Millisecond
	}
	if o.Logger == nil {
		o.Logger = log.New(os.Stderr, "[workercluster] ", log.LstdFlags)
	}
	if o.LogOutput == nil {
		o.LogOutput = io.Discard
	}
}

// bootstrapLocked creates a one-voter configuration on this node. It is
// only called when opts.Bootstrap is true.
//
//nolint:unparam // signature reserves error return for future pre-bootstrap checks.
func (c *Cluster) bootstrapLocked() error {
	cfg := raft.Configuration{
		Servers: []raft.Server{
			{
				ID:      raft.ServerID(c.opts.NodeID),
				Address: raft.ServerAddress(c.opts.LocalAddr),
			},
		},
	}
	f := c.raft.BootstrapCluster(cfg)
	if err := f.Error(); err != nil {
		// BootstrapCluster returns an error when the cluster has
		// already been bootstrapped in a previous run. That's fine
		// for a restart: just keep going and let the existing state
		// take over.
		c.log.Printf("workercluster: bootstrap: %v (continuing with existing state)", err)
	}
	return nil
}

// joinSeedLocked asks the seed peer to add this node as a voter.
// rafthttp.Join(id, addr, timeout): id is OUR ServerID and addr is the
// peer's address we dial. The seed's address is what we want to
// contact; our own address is sent in the URL's address= param.
func (c *Cluster) joinSeedLocked(_ context.Context) error {
	if err := c.layer.Join(
		raft.ServerID(c.opts.NodeID),
		raft.ServerAddress(c.opts.Seed.Address),
		10*time.Second,
	); err != nil {
		return fmt.Errorf("workercluster: join seed %s: %w", c.opts.Seed.Address, err)
	}
	return nil
}

func membersAsSlice(m map[string]Member) []Member {
	out := make([]Member, 0, len(m))
	for _, v := range m {
		out = append(out, v)
	}
	return out
}

// EnsureDataDir creates dir and its parent directories.
func EnsureDataDir(dir string) error {
	return os.MkdirAll(dir, 0o750)
}