webui4frpc/internal/cluster/ring_handler.go

// Runtime handlers the ring engine needs, provided by the app:
//   - RuntimeLoad: sample this node's memory/network load;
//   - Claim: create the frpc worker/forward for a claimed task.
package cluster

import (
	"context"
	"os"
	"runtime"
)

// AppHandler is the app-provided bridge for ring load sampling + task claims.
type AppHandler struct {
	// LoadFn returns (mem%, net%) — lower is more idle.
	LoadFn func() (memPct, netPct float64)
	// ClaimFn creates the forward on this node (persist + spawn worker).
	ClaimFn func(ctx context.Context, tk *Task) error
	// RevokeFn cancels the forward on this node (stop worker + drop from store).
	RevokeFn func(ctx context.Context, tk *Task) error
	// RestartFn re-enables an already-claimed forward on this node (clear the
	// stopped flag + bring the worker back). Required because a stop keeps the
	// topology entry, so the normal claim path will not re-create it.
	RestartFn func(ctx context.Context, tk *Task) error
	// ReleaseFn stops this node's worker for a forward being handed back to the
	// ring, leaving the link ENABLED. When nil, Release falls back to RevokeFn —
	// which also stops the worker, but additionally marks the link stopped; that
	// is the pre-existing stop semantics and is wrong for a rebalance, so the app
	// should provide this.
	ReleaseFn func(ctx context.Context, tk *Task) error
}

// RuntimeLoad implements Handler.
func (a *AppHandler) RuntimeLoad() Load {
	mem, net := 0.0, 0.0
	if a.LoadFn != nil {
		mem, net = a.LoadFn()
	}
	return Load{MemPct: mem, NetPct: net}
}

// Claim implements Handler.
func (a *AppHandler) Claim(ctx context.Context, tk *Task) error {
	if a.ClaimFn == nil {
		return nil
	}
	return a.ClaimFn(ctx, tk)
}

// Revoke implements Handler.
func (a *AppHandler) Revoke(ctx context.Context, tk *Task) error {
	if a.RevokeFn == nil {
		return nil
	}
	return a.RevokeFn(ctx, tk)
}

// Release implements Handler.
func (a *AppHandler) Release(ctx context.Context, tk *Task) error {
	if a.ReleaseFn != nil {
		return a.ReleaseFn(ctx, tk)
	}
	// Fall back to the stop path so a host that has not provided ReleaseFn still
	// stops the worker (the forward keeps running otherwise). The link ends up
	// flagged stopped, which the claim path may then refuse — the host is warned
	// by the ring's own "re-publish was deduped" log line rather than silently
	// leaving two workers up.
	if a.RevokeFn != nil {
		return a.RevokeFn(ctx, tk)
	}
	return nil
}

// Restart implements Handler.
func (a *AppHandler) Restart(ctx context.Context, tk *Task) error {
	if a.RestartFn == nil {
		return nil
	}
	return a.RestartFn(ctx, tk)
}

// SampleMemLoad returns a cheap memory-usage percentage (0..100).
func SampleMemLoad() float64 {
	var m runtime.MemStats
	runtime.ReadMemStats(&m)
	total := m.Sys
	if total == 0 {
		return 0
	}
	return float64(m.Alloc) / float64(total) * 100
}

// SampleNetLoad is implemented in netload_linux.go for Linux and returns a
// safe zero fallback on other platforms. It is kept here as a package-level
// API because the ring handler injects it into RuntimeLoad.

// HostID returns a stable node id (hostname; empty fallback to pid).
func HostID() string {
	h, _ := os.Hostname()
	if h == "" {
		return "node"
	}
	return h
}