Skip to content
180 changes: 173 additions & 7 deletions adapter/admin_grpc.go
Original file line number Diff line number Diff line change
Expand Up @@ -41,6 +41,7 @@ type KeyVizSampler interface {
type AdminGroup interface {
Status() raftengine.Status
Configuration(ctx context.Context) (raftengine.Configuration, error)
SnapshotEvery() uint64
}

// NodeIdentity is the value form of the protobuf NodeIdentity message used for
Expand All @@ -55,6 +56,53 @@ func (n NodeIdentity) toProto() *pb.NodeIdentity {
return &pb.NodeIdentity{NodeId: n.NodeID, GrpcAddress: n.GRPCAddress}
}

// LeaderVersionProbe fetches the Admin service version for a peer address.
// Implementations must honor ctx for the 500ms async GetRaftGroups probe
// budget and any auth metadata copied from the inbound Admin request.
type LeaderVersionProbe func(ctx context.Context, grpcAddress string) (string, error)

// AdminOption adjusts optional AdminServer behavior without changing existing
// test construction call sites.
type AdminOption func(*AdminServer)

func WithAdminNodeVersion(version string) AdminOption {
return func(s *AdminServer) {
s.nodeVersion = version
}
}

func WithAdminLeaderVersionProbe(probe LeaderVersionProbe) AdminOption {
return func(s *AdminServer) {
s.leaderVersionProbe = probe
}
}

func WithAdminLeaderVersionProbeTimeout(timeout time.Duration) AdminOption {
return func(s *AdminServer) {
if timeout > 0 {
s.leaderVersionProbeTimeout = timeout
}
}
}

func WithAdminLeaderVersionCacheTTL(ttl time.Duration) AdminOption {
return func(s *AdminServer) {
if ttl > 0 {
s.leaderVersionCacheTTL = ttl
}
}
}

type versionCacheEntry struct {
version string
fetchedAt time.Time
}

const (
defaultAdminLeaderVersionProbeTimeout = 500 * time.Millisecond
defaultAdminLeaderVersionCacheTTL = 10 * time.Second
)

// AdminServer implements the node-side Admin gRPC service described in
// docs/admin_ui_key_visualizer_design.md §4 (Layer A). Phase 0 only implements
// GetClusterOverview and GetRaftGroups; remaining RPCs return Unimplemented so
Expand All @@ -78,6 +126,12 @@ type AdminServer struct {
// pairs atomically with concurrent RPC reads.
sampler KeyVizSampler

nodeVersion string
leaderVersionProbe LeaderVersionProbe
leaderVersionProbeTimeout time.Duration
leaderVersionCacheTTL time.Duration
versionCache sync.Map

pb.UnimplementedAdminServer
}

Expand All @@ -86,14 +140,22 @@ type AdminServer struct {
// snapshot shipped to the admin binary; callers that already have a membership
// source may pass nil and let the admin binary's fan-out layer discover peers
// by other means.
func NewAdminServer(self NodeIdentity, members []NodeIdentity) *AdminServer {
func NewAdminServer(self NodeIdentity, members []NodeIdentity, opts ...AdminOption) *AdminServer {
cloned := append([]NodeIdentity(nil), members...)
return &AdminServer{
self: self,
members: cloned,
groups: make(map[uint64]AdminGroup),
now: time.Now,
srv := &AdminServer{
self: self,
members: cloned,
groups: make(map[uint64]AdminGroup),
now: time.Now,
leaderVersionProbeTimeout: defaultAdminLeaderVersionProbeTimeout,
leaderVersionCacheTTL: defaultAdminLeaderVersionCacheTTL,
}
for _, opt := range opts {
if opt != nil {
opt(srv)
}
}
return srv
}

// SetClock overrides the clock used by GetRaftGroups, letting tests inject a
Expand Down Expand Up @@ -380,7 +442,7 @@ func mergeSeedMembers(seeds []NodeIdentity, selfID string, live *liveMembers) {
// GetRaftGroups returns per-group state snapshots. Phase 0 wires commit/applied
// indices only; per-follower contact and term history land in later phases.
func (s *AdminServer) GetRaftGroups(
_ context.Context,
ctx context.Context,
_ *pb.GetRaftGroupsRequest,
) (*pb.GetRaftGroupsResponse, error) {
s.groupsMu.RLock()
Expand Down Expand Up @@ -410,11 +472,115 @@ func (s *AdminServer) GetRaftGroups(
CommitIndex: st.CommitIndex,
AppliedIndex: st.AppliedIndex,
LastContactUnixMs: lastContactUnixMs,
LeaderNodeVersion: s.leaderNodeVersion(ctx, st.Leader, now),
})
}
return &pb.GetRaftGroupsResponse{Groups: out}, nil
}

func (s *AdminServer) GetNodeVersion(
context.Context,
*pb.GetNodeVersionRequest,
) (*pb.GetNodeVersionResponse, error) {
return &pb.GetNodeVersionResponse{NodeVersion: s.nodeVersion}, nil
}

func (s *AdminServer) leaderNodeVersion(ctx context.Context, leader raftengine.LeaderInfo, now time.Time) string {
if version, ok := s.localLeaderVersion(leader); ok {
return version
}
key := leaderVersionCacheKey(leader)
if version, ok := s.cachedLeaderVersion(key, now); ok {
return version
}
if s.leaderVersionProbe == nil || leader.Address == "" {
return ""
}
if version, ok := s.reserveLeaderVersionProbe(key, now); ok {
return version
}
s.probeLeaderVersionAsync(ctx, key, leader.Address)
return ""
}

func (s *AdminServer) localLeaderVersion(leader raftengine.LeaderInfo) (string, bool) {
if leader.ID == "" && leader.Address == "" {
return "", true
}
if leader.ID == s.self.NodeID || (leader.Address != "" && leader.Address == s.self.GRPCAddress) {
return s.nodeVersion, true
}
return "", false
}

func leaderVersionCacheKey(leader raftengine.LeaderInfo) string {
if leader.ID != "" {
return leader.ID
Comment thread
bootjp marked this conversation as resolved.
}
return leader.Address
}

func (s *AdminServer) cachedLeaderVersion(key string, now time.Time) (string, bool) {
if key == "" || s.leaderVersionCacheTTL <= 0 {
return "", false
}
actual, ok := s.versionCache.Load(key)
if !ok {
return "", false
}
entry, ok := actual.(versionCacheEntry)
if !ok {
s.versionCache.Delete(key)
return "", false
}
if now.Sub(entry.fetchedAt) > s.leaderVersionCacheTTL {
s.versionCache.Delete(key)
return "", false
}
return entry.version, true
}

func (s *AdminServer) reserveLeaderVersionProbe(key string, now time.Time) (string, bool) {
marker := versionCacheEntry{fetchedAt: now}
actual, loaded := s.versionCache.LoadOrStore(key, marker)
if !loaded {
return "", false
}
entry, ok := actual.(versionCacheEntry)
switch {
case !ok:
s.versionCache.Store(key, marker)
return "", false
case now.Sub(entry.fetchedAt) <= s.leaderVersionCacheTTL:
return entry.version, true
default:
s.versionCache.Store(key, marker)
return "", false
}
}

func (s *AdminServer) probeLeaderVersionAsync(ctx context.Context, key, address string) {
probe := s.leaderVersionProbe
timeout := s.leaderVersionProbeTimeout
now := s.now
var md metadata.MD
if incoming, ok := metadata.FromIncomingContext(ctx); ok {
md = incoming.Copy()
}
go func() {
probeCtx, cancel := context.WithTimeout(context.Background(), timeout)
defer cancel()
if md != nil {
probeCtx = metadata.NewOutgoingContext(probeCtx, md)
}
version, err := probe(probeCtx, address)
if err != nil {
version = ""
}
s.versionCache.Store(key, versionCacheEntry{version: version, fetchedAt: now()})

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

P2 Badge Guard cache writes from stale probe completions

When a leader-version probe is still in flight after its reservation has expired (for example with a probe timeout longer than the cache TTL, or when the injected server clock advances past the TTL), a later GetRaftGroups call can replace the expired marker and start a fresh probe, but the older goroutine still unconditionally stores its result here. If that older probe then times out or returns an older value after the fresh probe has populated the cache, it overwrites the newer entry and suppresses another probe for a full TTL; use a generation/CAS-style check against the reservation before storing the completion.

Useful? React with 👍 / 👎.

}()
}

func (s *AdminServer) snapshotLeaders() []*pb.GroupLeader {
s.groupsMu.RLock()
defer s.groupsMu.RUnlock()
Expand Down
Loading
Loading