yovision/Sense/server/app/sense/integration/brain_control/store.go

package brain_control

import (
	"encoding/json"
	"errors"
	"time"

	"gorm.io/gorm"
	"gorm.io/gorm/clause"
)

type GORMReplayStore struct{ DB *gorm.DB }

func (s GORMReplayStore) Consume(principal, tokenID string, expiresAt, now time.Time) bool {
	if s.DB == nil || principal == "" || tokenID == "" || !expiresAt.After(now) {
		return false
	}
	return s.DB.Transaction(func(tx *gorm.DB) error {
		if err := tx.Where("expires_at <= ?", now).Delete(&ReplayToken{}).Error; err != nil {
			return err
		}
		result := tx.Clauses(clause.OnConflict{DoNothing: true}).Create(&ReplayToken{Principal: principal, TokenID: tokenID, ExpiresAt: expiresAt, CreatedAt: now})
		if result.Error != nil {
			return result.Error
		}
		if result.RowsAffected != 1 {
			return errors.New("replayed")
		}
		return nil
	}) == nil
}

type RevisionStore struct{ DB *gorm.DB }

func (s RevisionStore) Next(configID string) (int64, error) {
	if s.DB == nil || !stableID.MatchString(configID) {
		return 0, errors.New("invalid revision store")
	}
	var next int64
	err := s.DB.Transaction(func(tx *gorm.DB) error {
		var row SourceRevision
		err := tx.Clauses(clause.Locking{Strength: "UPDATE"}).Where("config_id = ?", configID).First(&row).Error
		if errors.Is(err, gorm.ErrRecordNotFound) {
			row = SourceRevision{ConfigID: configID, Revision: 1}
			if err = tx.Create(&row).Error; err != nil {
				return err
			}
			next = 1
			return nil
		}
		if err != nil {
			return err
		}
		row.Revision++
		next = row.Revision
		return tx.Save(&row).Error
	})
	return next, err
}

type ProjectionStore struct {
	DB         *gorm.DB
	Clock      func() time.Time
	StaleAfter time.Duration
	FutureSkew time.Duration
}
type ProjectionView struct {
	Projection                                  RuntimeProjection
	Offline, Stale, Recovered, RevisionMismatch bool
}

func (s ProjectionStore) Ingest(raw []byte, expected map[string]int64) (ProjectionView, error) {
	if s.DB == nil {
		return ProjectionView{}, errors.New("projection database required")
	}
	now := time.Now().UTC()
	if s.Clock != nil {
		now = s.Clock().UTC()
	}
	if s.StaleAfter == 0 {
		s.StaleAfter = 90 * time.Second
	}
	if s.FutureSkew == 0 {
		s.FutureSkew = 30 * time.Second
	}
	status, err := parseRuntimeStatus(raw)
	if err != nil {
		return ProjectionView{}, err
	}
	if status.ObservedAt.After(now.Add(s.FutureSkew)) {
		return ProjectionView{}, errors.New("FUTURE_OBSERVATION")
	}
	var view ProjectionView
	err = s.DB.Transaction(func(tx *gorm.DB) error {
		var old RuntimeProjection
		find := tx.Where("brain_instance_ref = ?", status.BrainInstanceRef).First(&old).Error
		if find == nil {
			if old.StatusID == status.StatusID {
				view.Projection = old
				return nil
			}
			if status.Sequence <= old.Sequence {
				return errors.New("OUT_OF_ORDER_STATUS")
			}
			if !validRuntimeTransition(old.RuntimeState, status.Runtime.State) {
				return errors.New("INVALID_STATUS_TRANSITION")
			}
		}
		if find != nil && !errors.Is(find, gorm.ErrRecordNotFound) {
			return find
		}
		expectedJSON, _ := json.Marshal(expected)
		configs, _ := json.Marshal(status.Configurations)
		health, _ := json.Marshal(status.Health)
		inputs, _ := json.Marshal(status.Inputs)
		p := RuntimeProjection{BrainInstanceRef: status.BrainInstanceRef, StatusID: status.StatusID, Sequence: status.Sequence, ObservedAt: status.ObservedAt, ReceivedAt: now, RuntimeState: status.Runtime.State, RuntimeVersion: status.Runtime.Version, ModelRef: status.Model.ModelRef, ModelVersion: status.Model.Version, HealthOverall: status.Health.Overall, ExpectedRevisionsJSON: string(expectedJSON), ConfigurationsJSON: string(configs), HealthJSON: string(health), InputsJSON: string(inputs), WasOffline: find == nil && now.Sub(old.ObservedAt) > s.StaleAfter}
		if find == nil {
			p.CreatedAt = old.CreatedAt
			view.Recovered = p.WasOffline
		}
		if err := tx.Save(&p).Error; err != nil {
			return err
		}
		view.Projection = p
		return nil
	})
	if err != nil {
		return ProjectionView{}, err
	}
	view.Stale = now.Sub(view.Projection.ObservedAt) > s.StaleAfter
	view.Offline = view.Stale
	applied := make(map[string]*int64, len(status.Configurations))
	for _, c := range status.Configurations {
		applied[c.ConfigID] = c.AppliedRevision
	}
	for configID, want := range expected {
		got, ok := applied[configID]
		if !ok || got == nil || *got != want {
			view.RevisionMismatch = true
		}
	}
	return view, nil
}

func (s ProjectionStore) View(instance string) (ProjectionView, error) {
	var p RuntimeProjection
	if err := s.DB.First(&p, "brain_instance_ref = ?", instance).Error; err != nil {
		return ProjectionView{}, err
	}
	now := time.Now().UTC()
	if s.Clock != nil {
		now = s.Clock().UTC()
	}
	stale := s.StaleAfter
	if stale == 0 {
		stale = 90 * time.Second
	}
	return ProjectionView{Projection: p, Offline: now.Sub(p.ObservedAt) > stale, Stale: now.Sub(p.ObservedAt) > stale}, nil
}