data-mesher/pkg/cluster/cluster.go

// Package cluster provides types and methods for working with a libp2p cluster.
package cluster

import (
	"context"
	"errors"
	"fmt"
	"log/slog"
	"math"
	"time"

	"git.clan.lol/clan/data-mesher/pkg/config"
	filetransferv1 "git.clan.lol/clan/data-mesher/pkg/libp2p/services/file_transfer/v1"
	stateexchangev1 "git.clan.lol/clan/data-mesher/pkg/libp2p/services/state_exchange/v1"
	dmnet "git.clan.lol/clan/data-mesher/pkg/net"
	"git.clan.lol/clan/data-mesher/pkg/state"
	dht "github.com/libp2p/go-libp2p-kad-dht"
	"github.com/libp2p/go-libp2p/core/host"
	"github.com/libp2p/go-libp2p/core/peerstore"
)

var ErrNoPeersAvailable = errors.New("no peers discovered in DHT")

// Cluster represents a distributed peer-to-peer network, managed via libp2p.
type Cluster struct {
	cfg   *config.Config
	files *state.Files

	ps   peerstore.Peerstore
	host host.Host

	dht *dht.IpfsDHT

	fileServer     *filetransferv1.Server
	stateExchanger *stateexchangev1.Exchanger
}

// New initializes and returns a new Cluster instance with the provided configuration and store.
func New(cfg *config.Config, files *state.Files, resolver dmnet.InterfaceResolver) (*Cluster, error) {
	// create libp2p host
	hst, ps, err := newHost(cfg, resolver)
	if err != nil {
		return nil, fmt.Errorf("failed to create libp2p host: %w", err)
	}

	// log host info
	slog.Info(
		"libp2p host created",
		"peer_id", hst.ID().String(),
		"addresses", hst.Addrs(),
	)

	cluster := &Cluster{
		cfg:   cfg,
		ps:    ps,
		host:  hst,
		files: files,
	}

	return cluster, nil
}

// Join connects to bootstrap peers and starts the push-pull loop.
func (c *Cluster) Join(ctx context.Context) error {
	var err error

	// set up network notifiee for connection/disconnection events
	c.host.Network().Notify(&notify{cluster: c})

	// create the DHT for peer discovery
	if c.dht, err = NewDHT(ctx, c.host, &c.cfg.Cluster); err != nil {
		return fmt.Errorf("join failure: %w", err)
	}

	// register a file server
	c.fileServer = filetransferv1.NewServer(c.host, c.cfg, c.files)

	// register a state exchanger
	c.stateExchanger = stateexchangev1.NewExchanger(
		c.cfg,
		c.ps,
		c.host,
		c.files,
	)

	c.exchangeStateOnJoin(ctx)

	return nil
}

// Leave gracefully shuts down the cluster by closing the host.
func (c *Cluster) Leave(_ context.Context) error {
	slog.Info("leaving cluster")

	// cancel any in-flight downloads
	if err := c.stateExchanger.Close(); err != nil {
		slog.Error("failed to cancel in-flight downloads", slog.Any("err", err))
	}

	// close the DHT
	if err := c.dht.Close(); err != nil {
		slog.Error("failed to close DHT", slog.Any("err", err))
	}

	// close the host
	if err := c.host.Close(); err != nil {
		slog.Error("failed to close libp2p host", slog.Any("err", err))
	}

	slog.Info("finished leaving cluster")

	return nil
}

// PushPullLoop periodically selects one random connected peer and exchanges state.
func (c *Cluster) PushPullLoop(ctx context.Context) error {
	peerCount := c.dht.RoutingTable().Size()
	interval := PushPullInterval(peerCount, c.cfg.Cluster.PushPullInterval)

	ticker := time.NewTicker(interval)
	defer ticker.Stop()

	for {
		select {
		case <-ctx.Done():
			return nil
		case <-ticker.C:
			slog.Debug("attempting push/pull", slog.Int("peer_count", peerCount))

			if err := c.exchangeState(ctx, 3); err != nil {
				slog.Error(
					"push/pull failure, will try again later",
					slog.String("err", err.Error()),
					slog.Duration("interval", c.cfg.Cluster.PushPullInterval),
				)
			}

			// reset the ticker to the next interval
			nextInterval := PushPullInterval(peerCount, c.cfg.Cluster.PushPullInterval)

			ticker.Reset(nextInterval)

			slog.Debug("push/pull successful", slog.Duration("interval", nextInterval))
		}
	}
}

// PushPullInterval returns the push/pull interval for the given network size.
// We scale the push/pull interval based on the network size for clusters > 32 nodes to prevent saturating the
// network with state sync requests, same as memberlist.
func PushPullInterval(n int, interval time.Duration) time.Duration {
	if n < 32 {
		return interval
	}

	result := interval.Seconds() * (math.Log2(float64(n)) - math.Log2(float64(32)) + 1)

	return time.Duration(result) * time.Second
}