iop/apps/agent/internal/clientprocess/store.go

package clientprocess

import (
	"bytes"
	"context"
	"encoding/json"
	"errors"
	"fmt"
	"io"
	"strings"
	"time"

	"iop/packages/go/agenttask"
)

const clientRecordKeyPrefix = "client-process/"

type durableStore struct {
	store StateStore
}

func newDurableStore(store StateStore) (*durableStore, error) {
	if store == nil {
		return nil, fmt.Errorf("clientprocess: state store is required")
	}
	return &durableStore{store: store}, nil
}

func (s *durableStore) load(
	ctx context.Context,
	kind ClientKind,
) (Record, string, bool, error) {
	payload, revision, found, err := s.store.LoadIntegrationRecord(
		ctx,
		clientRecordKey(kind),
	)
	if err != nil {
		return Record{}, "", false, err
	}
	if !found {
		return initialRecord(kind), "", false, nil
	}
	record, err := decodeRecord(payload, kind)
	if err != nil {
		return Record{}, "", false, err
	}
	return record, revision, true, nil
}

func (s *durableStore) save(
	ctx context.Context,
	record Record,
	expected string,
) (string, error) {
	if err := validateRecord(record, record.Kind); err != nil {
		return "", err
	}
	payload, err := json.Marshal(record)
	if err != nil {
		return "", fmt.Errorf("clientprocess: encode durable record: %w", err)
	}
	revision, err := s.store.CompareAndSwapIntegrationRecord(
		ctx,
		clientRecordKey(record.Kind),
		expected,
		payload,
	)
	if errors.Is(err, agenttask.ErrRevisionConflict) {
		return "", ErrStateConflict
	}
	if err != nil {
		return "", err
	}
	return revision, nil
}

func clientRecordKey(kind ClientKind) string {
	return clientRecordKeyPrefix + string(kind)
}

func decodeRecord(payload []byte, kind ClientKind) (Record, error) {
	decoder := json.NewDecoder(bytes.NewReader(payload))
	decoder.DisallowUnknownFields()
	var record Record
	if err := decoder.Decode(&record); err != nil {
		return Record{}, fmt.Errorf("clientprocess: decode durable record: %w", err)
	}
	var trailing any
	if err := decoder.Decode(&trailing); !errors.Is(err, io.EOF) {
		return Record{}, fmt.Errorf("clientprocess: durable record has trailing data")
	}
	if err := validateRecord(record, kind); err != nil {
		return Record{}, err
	}
	return record, nil
}

func validateRecord(record Record, kind ClientKind) error {
	if record.SchemaVersion != RecordSchemaVersion {
		return fmt.Errorf(
			"clientprocess: unsupported record schema %d",
			record.SchemaVersion,
		)
	}
	if record.Kind != kind || !validKind(kind) {
		return fmt.Errorf("clientprocess: durable client identity is invalid")
	}
	switch record.State {
	case StateStopped, StateCrashed:
		if record.Identity != nil {
			return fmt.Errorf("clientprocess: stopped or crashed state cannot have process identity")
		}
		if record.Connected {
			return fmt.Errorf("clientprocess: stopped or crashed state cannot be connected")
		}
	case StateConnected:
		if !record.Connected {
			return fmt.Errorf("clientprocess: connected state must have Connected=true")
		}
		if record.Identity == nil {
			return fmt.Errorf("clientprocess: connected state must have process identity")
		}
	case StateStarting:
		if record.Connected {
			return fmt.Errorf("clientprocess: starting state cannot be connected")
		}
	default:
		return fmt.Errorf("clientprocess: durable lifecycle state is invalid")
	}
	if record.RestartAttempts < 0 {
		return fmt.Errorf("clientprocess: durable restart attempts are invalid")
	}
	for _, identity := range []*ProcessIdentity{
		record.Identity,
		record.LastIdentity,
	} {
		if identity != nil &&
			(identity.PID <= 0 ||
				strings.TrimSpace(identity.StartToken) == "" ||
				len(identity.StartToken) > 512) {
			return fmt.Errorf("clientprocess: durable process identity is invalid")
		}
	}
	if len(record.Commands) > maxCommandReceipts {
		return fmt.Errorf("clientprocess: durable command receipts exceed limit")
	}
	for k, receipt := range record.Commands {
		if strings.TrimSpace(k) == "" || k != receipt.CommandID {
			return fmt.Errorf("clientprocess: durable command receipt key mismatch")
		}
		if !validCommandAction(receipt.Action) {
			return fmt.Errorf("clientprocess: durable command receipt action is invalid")
		}
		switch receipt.Status {
		case CommandReceiptPending:
			if !validPendingReceipt(receipt) {
				return fmt.Errorf("clientprocess: pending command receipt has a result")
			}
		case CommandReceiptCompleted:
			if !validCompletedReceipt(receipt) {
				return fmt.Errorf("clientprocess: durable command receipt result is invalid")
			}
		default:
			return fmt.Errorf("clientprocess: durable command receipt status is invalid")
		}
	}
	if record.UpdatedAt.IsZero() {
		return fmt.Errorf("clientprocess: durable update time is required")
	}
	return nil
}

func initialRecord(kind ClientKind) Record {
	return Record{
		SchemaVersion: RecordSchemaVersion,
		Kind:          kind,
		State:         StateStopped,
		UpdatedAt:     time.Now().UTC(),
	}
}

func validKind(kind ClientKind) bool {
	return kind == ClientFlutter || kind == ClientUnity
}

func cloneRecord(record Record) Record {
	out := record
	if record.Identity != nil {
		identity := *record.Identity
		out.Identity = &identity
	}
	if record.LastIdentity != nil {
		identity := *record.LastIdentity
		out.LastIdentity = &identity
	}
	if record.Commands != nil {
		out.Commands = make(map[string]CommandReceipt, len(record.Commands))
		for k, v := range record.Commands {
			out.Commands[k] = v
		}
	}
	return out
}