custody.go

package dasmon

import (
	"context"
	"fmt"
	"log/slog"
	"sync"
	"time"

	"github.com/ethereum/go-ethereum/p2p/enode"
	"github.com/ethp2p/dasmon/store"
	"github.com/libp2p/go-libp2p/core/peer"
)

type MonitoredPeer struct {
	context.CancelFunc
	context.Context

	// Static fields
	peerId    peer.ID
	nodeId    enode.ID
	allGroups []uint64
	log       *slog.Logger // Cached logger with peer/node context

	// Test history
	history *store.PeerHistory

	// Dynamic fields (protected by mu)
	mu             sync.RWMutex
	custody        PeerCustody
	nextRunAt      time.Time  // Earliest time a job can run for this peer
	lastTested     time.Time  // When peer was last tested (job executed)
	nextRefreshAt  time.Time  // When to refresh custody status/metadata
	disconnectedAt *time.Time // When peer disconnected (nil if connected)
}

func newMonitoredPeer(id peer.ID, database store.Store) (*MonitoredPeer, error) {
	nodeId, err := peerIDToNodeID(id)
	if err != nil {
		return nil, fmt.Errorf("failed to convert peer ID to node ID: %v", err)
	}
	allGroups, err := computeCustodyGroups(nodeId, 128)
	if err != nil {
		return nil, fmt.Errorf("failed to compute all custody groups: %v", err)
	}
	pc := &MonitoredPeer{
		peerId:    id,
		nodeId:    nodeId,
		allGroups: allGroups,
		history:   store.NewPeerHistory(id, database),
		log:       log.With("peer", id, "node", nodeId),
	}
	return pc, nil
}

func (mp *MonitoredPeer) String() string {
	return fmt.Sprintf("Peer{id=%s, custody=%s}", mp.peerId, mp.custody.String())
}

func (mp *MonitoredPeer) refresh(info *PeerCustody) (prev, next PeerCustody) {
	if mp.custody.Equals(*info) {
		// No changes
		return mp.custody, mp.custody
	}

	diff := mp.custody.Diff(*info)
	mp.log.Info("custody updated", diff...)
	mp.custody = *info
	return mp.custody, *info
}

func (mp *MonitoredPeer) Custody() PeerCustody {
	mp.mu.RLock()
	defer mp.mu.RUnlock()
	return mp.custody
}

func (mp *MonitoredPeer) groups() []uint64 {
	mp.mu.RLock()
	defer mp.mu.RUnlock()
	return mp.allGroups[:mp.custody.CustodyGroupCount]
}

type JobConstraints struct {
	// Head is the current head
	Head uint64

	// MaxColumns to request in a single RPC call
	MaxColumns int

	// MaxSlots to request in a single RPC call
	MaxSlots uint64

	// BlockFilter restricts slot selection
	BlockFilter *BlockFilter

	// Sampler chooses slots and columns based on coverage
	Sampler *Sampler

	// CommitmentTracker provides access to our known commitments
	CommitmentTracker *CommitmentTracker

	// TestInterval is the minimum time between jobs for same peer
	TestInterval DurationSer
}

// PollJob returns a job if this peer is ready to run one, otherwise nil.
// Updates lastTested and nextRunAt on success.
func (mp *MonitoredPeer) PollJob(constraints *JobConstraints) *Job {
	mp.mu.Lock()
	defer mp.mu.Unlock()

	now := time.Now()
	if now.Before(mp.nextRunAt) {
		return nil
	}

	job := mp.createJobLocked(constraints)
	if job == nil {
		return nil
	}

	mp.lastTested = now
	mp.nextRunAt = now.Add(time.Duration(constraints.TestInterval))

	return job
}

// GetNextRefreshAt returns when this peer needs refreshing (for scheduler)
func (mp *MonitoredPeer) GetNextRefreshAt() time.Time {
	mp.mu.RLock()
	defer mp.mu.RUnlock()
	return mp.nextRefreshAt
}

// SetNextRefreshAt sets when the next refresh should occur
func (mp *MonitoredPeer) SetNextRefreshAt(t time.Time) {
	mp.mu.Lock()
	defer mp.mu.Unlock()
	mp.nextRefreshAt = t
}

// IsDisconnected returns true if peer is disconnected
func (mp *MonitoredPeer) IsDisconnected() bool {
	mp.mu.RLock()
	defer mp.mu.RUnlock()
	return mp.disconnectedAt != nil
}

// MarkDisconnected marks peer as disconnected
func (mp *MonitoredPeer) MarkDisconnected(t time.Time) {
	mp.mu.Lock()
	defer mp.mu.Unlock()
	mp.disconnectedAt = &t
}

// MarkReconnected clears disconnection status
func (mp *MonitoredPeer) MarkReconnected() {
	mp.mu.Lock()
	defer mp.mu.Unlock()
	mp.disconnectedAt = nil
}

// createJobLocked is the existing CreateJob logic but assumes lock is held
func (mp *MonitoredPeer) createJobLocked(constraints *JobConstraints) *Job {
	// Build sampling context from peer state
	ctx := SamplingContext{
		PeerID:          mp.peerId,
		CustodyRange:    RangeClosed(mp.custody.EarliestSlot, mp.custody.HeadSlot),
		CustodyColumns:  mp.groups(),
		MaxSlots:        constraints.MaxSlots,
		MaxColumns:      constraints.MaxColumns,
		BlockFilter:     constraints.BlockFilter.EffectiveRange(constraints.Head),
		CommitmentRange: constraints.CommitmentTracker.Range(),
		History:         mp.history,
	}

	// Use sampler to select slots and columns
	plan := constraints.Sampler.SelectSamples(ctx)
	if plan == nil {
		mp.log.Debug("no sampling plan generated")
		return nil
	}

	return &Job{
		PeerID:    mp.peerId,
		StartSlot: plan.StartSlot,
		EndSlot:   plan.EndSlot,
		Columns:   plan.Columns,
	}
}