90CNServer/server/internal/platform/nativepatch/etcd_control.go

package nativepatch

import (
	"context"
	"crypto/rand"
	"encoding/base64"
	"encoding/hex"
	"encoding/json"
	"errors"
	"fmt"
	"sort"
	"strings"
	"sync"
	"time"

	clientv3 "go.etcd.io/etcd/client/v3"
)

const defaultEtcdNamespace = "longheng/nativepatch"

var (
	ErrEtcdControlInstanceID = errors.New("native patch etcd control instance id generation failed")
	ErrEtcdControlStore      = errors.New("native patch etcd control store is required")
	etcdControlRandRead      = rand.Read
)

type EtcdClient interface {
	Get(context.Context, string, ...clientv3.OpOption) (*clientv3.GetResponse, error)
	Put(context.Context, string, string, ...clientv3.OpOption) (*clientv3.PutResponse, error)
	Watch(context.Context, string, ...clientv3.OpOption) clientv3.WatchChan
	Close() error
}

type EtcdControlOptions struct {
	Client      EtcdClient
	Namespace   string
	CloseClient bool
	Now         func() time.Time
}

type EtcdControlStore struct {
	client      EtcdClient
	namespace   string
	closeClient bool
	instanceID  string
	now         func() time.Time

	mu       sync.Mutex
	sequence uint64
}

func NewEtcdControlStore(options EtcdControlOptions) (*EtcdControlStore, error) {
	if options.Client == nil {
		return nil, ErrEtcdControlStore
	}
	namespace := normEtcdNamespace(options.Namespace)
	now := options.Now
	if now == nil {
		now = time.Now
	}
	instanceID, err := newEtcdInstID()
	if err != nil {
		return nil, err
	}
	return &EtcdControlStore{
		client:      options.Client,
		namespace:   namespace,
		closeClient: options.CloseClient,
		instanceID:  instanceID,
		now:         now,
	}, nil
}

func (s *EtcdControlStore) Publish(ctx context.Context, target string, intent ControlIntent) error {
	if ctx == nil {
		ctx = context.Background()
	}
	if err := ctx.Err(); err != nil {
		return err
	}
	client, namespace, _, _, err := s.ready()
	if err != nil {
		return err
	}
	target = strings.TrimSpace(target)
	if target == "" {
		return fmt.Errorf("native patch control target is required")
	}
	intent.Target = target
	if intent.Action == "" {
		intent.Action = ControlApply
	}
	if intent.Action != ControlApply && intent.Action != ControlRestore {
		return fmt.Errorf("unsupported native patch control action %s", intent.Action)
	}
	if strings.TrimSpace(intent.Version) == "" {
		return ErrVersionRequired
	}
	if intent.Action == ControlApply && strings.TrimSpace(intent.PackageDir) == "" {
		return ErrPlanRequired
	}
	data, err := json.Marshal(intent)
	if err != nil {
		return err
	}
	_, err = client.Put(ctx, intentKey(namespace, target), string(data))
	return err
}

func (s *EtcdControlStore) Watch(ctx context.Context, target string) (<-chan ControlIntent, error) {
	if ctx == nil {
		ctx = context.Background()
	}
	if err := ctx.Err(); err != nil {
		return nil, err
	}
	client, namespace, _, _, err := s.ready()
	if err != nil {
		return nil, err
	}
	target = strings.TrimSpace(target)
	if target == "" {
		return nil, fmt.Errorf("native patch control target is required")
	}
	key := intentKey(namespace, target)
	out := make(chan ControlIntent, 1)
	resp, err := client.Get(ctx, key)
	if err != nil {
		close(out)
		return nil, err
	}
	if len(resp.Kvs) > 0 {
		if intent, err := decodeControlIntent(resp.Kvs[len(resp.Kvs)-1].Value); err == nil {
			if !sendControlIntent(ctx, out, intent) {
				return out, nil
			}
		}
	}
	watchCh := client.Watch(ctx, key)
	go func() {
		defer close(out)
		for {
			select {
			case <-ctx.Done():
				return
			case watchResp, ok := <-watchCh:
				if !ok {
					return
				}
				for _, event := range watchResp.Events {
					if event == nil || event.Kv == nil {
						continue
					}
					intent, err := decodeControlIntent(event.Kv.Value)
					if err != nil {
						continue
					}
					if !sendControlIntent(ctx, out, intent) {
						return
					}
				}
			}
		}
	}()
	return out, nil
}

func sendControlIntent(ctx context.Context, out chan ControlIntent, intent ControlIntent) bool {
	select {
	case out <- intent:
		return true
	case <-ctx.Done():
		return false
	default:
	}
	select {
	case <-out:
	default:
	}
	select {
	case out <- intent:
		return true
	case <-ctx.Done():
		return false
	default:
		return true
	}
}

func (s *EtcdControlStore) Report(ctx context.Context, progress ControlProgress) error {
	if ctx == nil {
		ctx = context.Background()
	}
	if err := ctx.Err(); err != nil {
		return err
	}
	client, namespace, now, instanceID, err := s.ready()
	if err != nil {
		return err
	}
	progress.Target = strings.TrimSpace(progress.Target)
	progress.NodeID = strings.TrimSpace(progress.NodeID)
	if progress.Target == "" || progress.NodeID == "" {
		return fmt.Errorf("native patch progress target and node_id are required")
	}
	if progress.UpdatedAt.IsZero() {
		progress.UpdatedAt = now().UTC()
	}
	data, err := json.Marshal(progress)
	if err != nil {
		return err
	}
	_, err = client.Put(ctx, s.progressKey(namespace, instanceID, progress), string(data))
	return err
}

func (s *EtcdControlStore) ListProgress(ctx context.Context, target string) ([]ControlProgress, error) {
	if ctx == nil {
		ctx = context.Background()
	}
	if err := ctx.Err(); err != nil {
		return nil, err
	}
	client, namespace, _, _, err := s.ready()
	if err != nil {
		return nil, err
	}
	resp, err := client.Get(ctx, progressPrefix(namespace, target), clientv3.WithPrefix())
	if err != nil {
		return nil, err
	}
	sort.Slice(resp.Kvs, func(i, j int) bool {
		return string(resp.Kvs[i].Key) < string(resp.Kvs[j].Key)
	})
	out := make([]ControlProgress, 0, len(resp.Kvs))
	for _, kv := range resp.Kvs {
		var progress ControlProgress
		if err := json.Unmarshal(kv.Value, &progress); err != nil {
			return nil, err
		}
		out = append(out, progress)
	}
	return out, nil
}

func (s *EtcdControlStore) ready() (EtcdClient, string, func() time.Time, string, error) {
	if s == nil {
		return nil, "", nil, "", ErrEtcdControlStore
	}
	s.mu.Lock()
	defer s.mu.Unlock()
	if s.client == nil {
		return nil, "", nil, "", ErrEtcdControlStore
	}
	now := s.now
	if now == nil {
		now = time.Now
	}
	instanceID := s.instanceID
	if instanceID == "" {
		generated, err := newEtcdInstID()
		if err != nil {
			return nil, "", nil, "", err
		}
		s.instanceID = generated
		instanceID = generated
	}
	return s.client, normEtcdNamespace(s.namespace), now, instanceID, nil
}

func (s *EtcdControlStore) Close() error {
	if s == nil || s.client == nil || !s.closeClient {
		return nil
	}
	return s.client.Close()
}

func normEtcdNamespace(namespace string) string {
	namespace = strings.Trim(strings.TrimSpace(namespace), "/")
	if namespace == "" {
		return defaultEtcdNamespace
	}
	return namespace
}

func intentKey(namespace string, target string) string {
	return namespace + "/intent/" + encodeEtcdPart(target)
}

func progressPrefix(namespace string, target string) string {
	return namespace + "/progress/" + encodeEtcdPart(target) + "/"
}

func (s *EtcdControlStore) progressKey(namespace string, instanceID string, progress ControlProgress) string {
	s.mu.Lock()
	s.sequence++
	sequence := s.sequence
	s.mu.Unlock()
	return fmt.Sprintf("%s%s/%020d/%s/%020d", progressPrefix(namespace, progress.Target), encodeEtcdPart(progress.NodeID), progress.UpdatedAt.UnixNano(), instanceID, sequence)
}

func newEtcdInstID() (string, error) {
	var b [8]byte
	if _, err := etcdControlRandRead(b[:]); err != nil {
		return "", fmt.Errorf("%w: %w", ErrEtcdControlInstanceID, err)
	}
	return hex.EncodeToString(b[:]), nil
}

func encodeEtcdPart(value string) string {
	value = strings.TrimSpace(value)
	if value == "" {
		return "_"
	}
	return base64.RawURLEncoding.EncodeToString([]byte(value))
}

func decodeControlIntent(data []byte) (ControlIntent, error) {
	var intent ControlIntent
	if err := json.Unmarshal(data, &intent); err != nil {
		return ControlIntent{}, err
	}
	return intent, nil
}