pangolin/server/internal/provision/service.go

package provision

import (
	"context"
	"encoding/json"
	"fmt"
	"sync"
	"time"
)

// Config wires a Service together. Only Store and Adapters are mandatory;
// everything else falls back to a safe default.
type Config struct {
	Store    Store
	Adapters AdapterFactory
	Tokens   BootstrapIssuer
	Renderer CloudInitRenderer
	Prober   Prober
	Alert    AlertHook
	Clock    Clock

	// ControlPlaneURL is injected into cloud-init so the agent knows where to
	// enroll (task #6).
	ControlPlaneURL string
	// DrainTimeout is how long a draining node is given before forced destroy
	// (doc/04 §3: default 30 minutes).
	DrainTimeout time.Duration
	// ProbeTimeout bounds the simplified probing wait during Replace.
	ProbeTimeout time.Duration
}

// Service is the ProvisionService. All operations are idempotent and write
// node_events + audit_log (doc/04 §4).
type Service struct {
	store           Store
	adapters        AdapterFactory
	tokens          BootstrapIssuer
	renderer        CloudInitRenderer
	prober          Prober
	alert           AlertHook
	clock           Clock
	controlPlaneURL string
	drainTimeout    time.Duration
	probeTimeout    time.Duration

	mu       sync.Mutex
	failPool map[Pool]int // consecutive boot/probe failures per pool
}

// NewService builds a Service, applying defaults for optional dependencies.
func NewService(cfg Config) (*Service, error) {
	if cfg.Store == nil {
		return nil, fmt.Errorf("provision: Store is required")
	}
	if cfg.Adapters == nil {
		return nil, fmt.Errorf("provision: Adapters factory is required")
	}
	s := &Service{
		store:           cfg.Store,
		adapters:        cfg.Adapters,
		tokens:          cfg.Tokens,
		renderer:        cfg.Renderer,
		prober:          cfg.Prober,
		alert:           cfg.Alert,
		clock:           cfg.Clock,
		controlPlaneURL: cfg.ControlPlaneURL,
		drainTimeout:    cfg.DrainTimeout,
		probeTimeout:    cfg.ProbeTimeout,
		failPool:        make(map[Pool]int),
	}
	if s.alert == nil {
		s.alert = noopAlert{}
	}
	if s.clock == nil {
		s.clock = realClock{}
	}
	if s.drainTimeout == 0 {
		s.drainTimeout = 30 * time.Minute
	}
	if s.probeTimeout == 0 {
		s.probeTimeout = 5 * time.Minute
	}
	return s, nil
}

// resolveAdapter loads the provider row and its adapter.
func (s *Service) resolveAdapter(ctx context.Context, providerID int64) (*Provider, CloudAdapter, error) {
	p, err := s.store.GetProvider(ctx, providerID)
	if err != nil {
		return nil, nil, err
	}
	if p == nil {
		return nil, nil, fmt.Errorf("provision: provider %d not found", providerID)
	}
	a, err := s.adapters.For(p)
	if err != nil {
		return nil, nil, fmt.Errorf("provision: resolve adapter for %s: %w", p.APIKind, err)
	}
	return p, a, nil
}

// CreateNode provisions one node. It is idempotent on idempotencyKey: replaying
// the same key returns the already-created node without booting a second
// machine (validation: "CreateNode 幂等键重放不重复开机").
//
// Flow: idempotency check → bind key→uuid → insert nodes(provisioning) →
// issue bootstrap token (task #5) → render cloud-init (task #6) → vendor boot →
// record instance ID + endpoint → node_event(provisioned) + audit_log.
// The node stays in provisioning until the agent self-registers and probing
// promotes it to up.
func (s *Service) CreateNode(ctx context.Context, spec NodeSpec, idempotencyKey string) (*Node, error) {
	if idempotencyKey == "" {
		return nil, fmt.Errorf("provision: idempotencyKey is required")
	}

	// 1. Idempotency replay.
	if uuid, found, err := s.store.LookupIdempotency(ctx, idempotencyKey); err != nil {
		return nil, err
	} else if found {
		return s.store.GetNodeByUUID(ctx, uuid)
	}

	pool := poolForTier(spec.Tier)

	// 2. Allocate UUID and bind the idempotency key BEFORE any vendor call, so a
	//    crash after this point resumes onto the same node instead of booting a
	//    duplicate.
	uuid, err := newUUID()
	if err != nil {
		return nil, err
	}
	if err := s.store.SaveIdempotency(ctx, idempotencyKey, uuid); err != nil {
		return nil, err
	}
	// Re-read in case of a race: another caller may have won the key.
	if winner, found, err := s.store.LookupIdempotency(ctx, idempotencyKey); err == nil && found && winner != uuid {
		return s.store.GetNodeByUUID(ctx, winner)
	}

	// 3. Insert the node in provisioning state.
	n := &Node{
		UUID:       uuid,
		Region:     spec.Region,
		NameZH:     spec.NameZH,
		NameEn:     spec.NameEn,
		Role:       orDefaultRole(spec.Role),
		Tier:       spec.Tier,
		Endpoint:   pendingEndpoint,
		HY2Port:    spec.HY2Port,
		RealityPBK: spec.RealityPBK,
		RealitySNI: spec.RealitySNI,
		ProviderID: spec.ProviderID,
		Tags:       spec.Tags,
		Status:     StatusProvisioning,
		Weight:     spec.Weight,
	}
	id, err := s.store.InsertNode(ctx, n)
	if err != nil {
		return nil, err
	}
	n.ID = id

	// 4. Resolve adapter.
	_, adapter, err := s.resolveAdapter(ctx, spec.ProviderID)
	if err != nil {
		s.failBoot(ctx, n, pool, err)
		return nil, err
	}

	// 5. Bootstrap token + cloud-init.
	userData := ""
	if s.renderer != nil {
		token := ""
		if s.tokens != nil {
			token, err = s.tokens.IssueToken(ctx, uuid)
			if err != nil {
				s.failBoot(ctx, n, pool, fmt.Errorf("issue bootstrap token: %w", err))
				return nil, err
			}
		}
		userData, err = s.renderer.Render(CloudInitData{
			NodeUUID:        uuid,
			BootstrapToken:  token,
			Region:          spec.Region,
			Role:            n.Role,
			Tier:            spec.Tier,
			ControlPlaneURL: s.controlPlaneURL,
		})
		if err != nil {
			s.failBoot(ctx, n, pool, fmt.Errorf("render cloud-init: %w", err))
			return nil, err
		}
	}

	// 6. Vendor boot.
	inst, err := adapter.CreateInstance(ctx, CreateInput{
		Region:   spec.Region,
		Plan:     spec.Plan,
		Label:    uuid,
		UserData: userData,
		Tags:     spec.Tags,
	})
	if err != nil {
		s.failBoot(ctx, n, pool, fmt.Errorf("vendor boot: %w", err))
		return nil, err
	}

	// 7. Record instance ID + endpoint.
	endpoint := joinEndpoint(inst.IP, spec.HY2Port)
	if err := s.store.SetNodeInstance(ctx, id, inst.ID, endpoint); err != nil {
		return nil, err
	}
	n.ProviderInstanceID = inst.ID
	n.Endpoint = endpoint

	// 8. Audit trail + reset failure counter.
	_ = s.store.WriteNodeEvent(ctx, id, EventProvisioned, jsonDetail(map[string]any{
		"instance_id": inst.ID,
		"region":      inst.Region,
	}))
	_ = s.store.WriteAuditLog(ctx, "provision", "create_node", "node:"+uuid,
		jsonDetail(map[string]any{"provider_id": spec.ProviderID, "region": spec.Region, "pool": pool}))
	s.resetFail(pool)

	return n, nil
}

// failBoot marks the node destroyed, bumps the per-pool failure counter, and
// fires a boot-failed alert (validation: "开机失败 → destroyed + 失败计数 + 告警钩子").
func (s *Service) failBoot(ctx context.Context, n *Node, pool Pool, cause error) {
	_ = s.store.UpdateNodeStatus(ctx, n.ID, StatusDestroyed)
	_ = s.store.WriteNodeEvent(ctx, n.ID, EventDestroyed, jsonDetail(map[string]any{
		"reason": "boot_failed",
		"error":  cause.Error(),
	}))
	count := s.bumpFail(pool)
	_ = s.store.WriteAuditLog(ctx, "provision", "create_node_failed", "node:"+n.UUID,
		jsonDetail(map[string]any{"error": cause.Error(), "pool": pool, "fail_count": count}))
	s.alert.Fire(ctx, Alert{
		Kind:      AlertBootFailed,
		NodeUUID:  n.UUID,
		Pool:      pool,
		Message:   cause.Error(),
		FailCount: count,
	})
}

// DestroyNode tears down a node: vendor destroy → IP release → nodes→destroyed.
// Idempotent: destroying an already-destroyed node is a no-op success.
func (s *Service) DestroyNode(ctx context.Context, nodeID int64) error {
	n, err := s.store.GetNode(ctx, nodeID)
	if err != nil {
		return err
	}
	if n == nil {
		return fmt.Errorf("provision: node %d not found", nodeID)
	}
	if n.Status == StatusDestroyed {
		return nil
	}
	if n.ProviderInstanceID != "" {
		_, adapter, err := s.resolveAdapter(ctx, n.ProviderID)
		if err != nil {
			return err
		}
		if err := adapter.DestroyInstance(ctx, n.ProviderInstanceID); err != nil {
			return fmt.Errorf("provision: destroy instance: %w", err)
		}
	}
	if err := s.store.UpdateNodeStatus(ctx, nodeID, StatusDestroyed); err != nil {
		return err
	}
	_ = s.store.WriteNodeEvent(ctx, nodeID, EventDestroyed, jsonDetail(map[string]any{"reason": "destroy"}))
	_ = s.store.WriteAuditLog(ctx, "provision", "destroy_node", "node:"+n.UUID, "")
	if _, err := s.store.BumpDirectoryVersion(ctx); err != nil {
		return err
	}
	return nil
}

// RotateIP swaps the elastic IP without re-creating the machine (doc/04 §4:
// "换 IP 不换机"). Flow: vendor IP re-bind → probe → update endpoint →
// bump directory version → node_event(ip_rotated).
func (s *Service) RotateIP(ctx context.Context, nodeID int64) (*Node, error) {
	n, err := s.store.GetNode(ctx, nodeID)
	if err != nil {
		return nil, err
	}
	if n == nil {
		return nil, fmt.Errorf("provision: node %d not found", nodeID)
	}
	_, adapter, err := s.resolveAdapter(ctx, n.ProviderID)
	if err != nil {
		return nil, err
	}
	if !adapter.SupportsElasticIP() {
		return nil, ErrElasticIPUnsupported
	}
	oldEndpoint := n.Endpoint
	newIP, err := adapter.AttachIP(ctx, n.ProviderInstanceID)
	if err != nil {
		return nil, fmt.Errorf("provision: attach IP: %w", err)
	}
	newEndpoint := joinEndpoint(newIP, n.HY2Port)

	// Probe the new endpoint before publishing it.
	probed := *n
	probed.Endpoint = newEndpoint
	if s.prober != nil {
		if err := s.prober.WaitReady(ctx, &probed); err != nil {
			s.alert.Fire(ctx, Alert{Kind: AlertProbeTimeout, NodeUUID: n.UUID, Message: err.Error()})
			return nil, fmt.Errorf("provision: probe rotated IP: %w", err)
		}
	}

	if err := s.store.UpdateNodeEndpoint(ctx, nodeID, newEndpoint); err != nil {
		return nil, err
	}
	version, err := s.store.BumpDirectoryVersion(ctx)
	if err != nil {
		return nil, err
	}
	_ = s.store.WriteNodeEvent(ctx, nodeID, EventIPRotated, jsonDetail(map[string]any{
		"old_endpoint": oldEndpoint,
		"new_endpoint": newEndpoint,
		"version":      version,
	}))
	_ = s.store.WriteAuditLog(ctx, "provision", "rotate_ip", "node:"+n.UUID,
		jsonDetail(map[string]any{"old": oldEndpoint, "new": newEndpoint}))
	n.Endpoint = newEndpoint
	return n, nil
}

// ListProviders returns enabled providers in a pool (empty pool = all pools).
func (s *Service) ListProviders(ctx context.Context, pool Pool) ([]*Provider, error) {
	return s.store.ListProviders(ctx, pool)
}

// --- failure-counter helpers (capacity protection, doc/04 §4.2) ---

func (s *Service) bumpFail(pool Pool) int {
	s.mu.Lock()
	defer s.mu.Unlock()
	s.failPool[pool]++
	return s.failPool[pool]
}

func (s *Service) resetFail(pool Pool) {
	s.mu.Lock()
	defer s.mu.Unlock()
	delete(s.failPool, pool)
}

// FailCount exposes the current consecutive-failure count for a pool (tests /
// monitoring).
func (s *Service) FailCount(pool Pool) int {
	s.mu.Lock()
	defer s.mu.Unlock()
	return s.failPool[pool]
}

// --- small helpers ---

func orDefaultRole(r Role) Role {
	if r == "" {
		return RoleEntry
	}
	return r
}

func joinEndpoint(ip string, port int) string {
	if port <= 0 {
		port = 443
	}
	return fmt.Sprintf("%s:%d", ip, port)
}

func jsonDetail(m map[string]any) string {
	b, err := json.Marshal(m)
	if err != nil {
		return "null"
	}
	return string(b)
}