pangolin/server/internal/agentd/command.go

package agentd

import (
	"context"

	agentv1 "github.com/wangjia/pangolin/server/internal/pb/agentv1"
)

// runSubscribe opens the Command stream (resuming from the last durably processed
// command id) and applies each command idempotently, Ack'ing after apply
// (at-least-once). It returns when ctx is cancelled or the stream errors, so the
// session can reconnect and resume.
func (a *Agent) runSubscribe(ctx context.Context, client agentv1.AgentServiceClient) error {
	nodeUUID := a.NodeUUID()
	stream, err := client.Subscribe(ctx, &agentv1.SubscribeRequest{
		NodeUUID:      nodeUUID,
		LastCommandID: a.lastCommandID.Load(),
	})
	if err != nil {
		return err
	}
	for {
		cmd, err := stream.Recv()
		if err != nil {
			return err
		}
		a.handleCommand(cmd)
		// Ack after apply. A failed Ack just causes redelivery, which the
		// dedup-by-command-id + idempotent apply below tolerates.
		if _, err := client.Ack(ctx, &agentv1.AckRequest{
			NodeUUID:  nodeUUID,
			CommandID: cmd.CommandID,
		}); err != nil {
			return err
		}
	}
}

// handleCommand applies one command. Commands at or below the high-water mark are
// skipped (already applied) — this makes redelivery after a dropped Ack safe.
func (a *Agent) handleCommand(cmd *agentv1.Command) {
	if cmd.CommandID != 0 && cmd.CommandID <= a.lastCommandID.Load() {
		return
	}
	switch cmd.Type {
	case agentv1.CommandTypeUpsert:
		if cmd.Upsert != nil && cmd.Upsert.Credential != nil {
			a.sb.Upsert(credFromPB(cmd.Upsert.Credential))
		}
	case agentv1.CommandTypeRevoke:
		if cmd.Revoke != nil {
			a.sb.Revoke(cmd.Revoke.DpUUID)
		}
	case agentv1.CommandTypeRotateCredential:
		if cmd.Rotate != nil && cmd.Rotate.NewCredential != nil {
			a.sb.Rotate(cmd.Rotate.OldDpUUID, credFromPB(cmd.Rotate.NewCredential), cmd.Rotate.GraceUntilUnix)
		}
	case agentv1.CommandTypeApplyConfig:
		if cmd.ApplyConfig != nil {
			a.sb.ApplyConfig(cmd.ApplyConfig, cmd.ApplyConfig.Credentials != nil)
		}
	case agentv1.CommandTypeLifecycle:
		a.handleLifecycle(cmd.Lifecycle)
	default:
		logf("ignoring command id=%d with unknown type=%d", cmd.CommandID, cmd.Type)
	}

	if cmd.CommandID > a.lastCommandID.Load() {
		a.lastCommandID.Store(cmd.CommandID)
	}
}

func (a *Agent) handleLifecycle(p *agentv1.LifecyclePayload) {
	if p == nil {
		return
	}
	switch p.Action {
	case agentv1.LifecycleActionDrain:
		// Draining is steered by control-plane node weight; the data plane keeps
		// serving existing sessions. Nothing to mutate locally.
		logf("lifecycle: drain")
	case agentv1.LifecycleActionResume:
		logf("lifecycle: resume")
	case agentv1.LifecycleActionShutdown:
		logf("lifecycle: shutdown requested")
		if a.shutdown != nil {
			a.shutdown()
		}
	default:
		logf("lifecycle: unknown action %d", p.Action)
	}
}

func credFromPB(c *agentv1.Credential) *Cred {
	flow := c.Flow
	if flow == "" {
		flow = DefaultFlow
	}
	return &Cred{
		DpUUID:    c.DpUUID,
		Protocol:  c.Protocol,
		Flow:      flow,
		ExpiresAt: c.ExpiresAtUnix,
	}
}