fission-src/pkg/utils/namespace_manager.go

package utils

import (
	"context"
	"errors"
	"sort"
	"sync"
	"time"

	corev1 "k8s.io/api/core/v1"
)

type NamespaceSubscriber interface {
	Name() string
	OnNamespaceAdd(ctx context.Context, record NamespaceRecord) error
	OnNamespaceRemove(ctx context.Context, record NamespaceRecord) error
	OnNamespaceResync(ctx context.Context, record NamespaceRecord) error
}

type NamespaceSubscriberFuncs struct {
	SubscriberName   string
	AddFunc          func(ctx context.Context, record NamespaceRecord) error
	RemoveFunc       func(ctx context.Context, record NamespaceRecord) error
	ResyncFunc       func(ctx context.Context, record NamespaceRecord) error
}

func (s NamespaceSubscriberFuncs) Name() string {
	return s.SubscriberName
}

func (s NamespaceSubscriberFuncs) OnNamespaceAdd(ctx context.Context, record NamespaceRecord) error {
	if s.AddFunc == nil {
		return nil
	}
	return s.AddFunc(ctx, record)
}

func (s NamespaceSubscriberFuncs) OnNamespaceRemove(ctx context.Context, record NamespaceRecord) error {
	if s.RemoveFunc == nil {
		return nil
	}
	return s.RemoveFunc(ctx, record)
}

func (s NamespaceSubscriberFuncs) OnNamespaceResync(ctx context.Context, record NamespaceRecord) error {
	if s.ResyncFunc == nil {
		return nil
	}
	return s.ResyncFunc(ctx, record)
}

type NamespaceManager interface {
	Snapshot() []string
	SnapshotRecords() []NamespaceRecord
	Get(name string) (NamespaceRecord, bool)
	Bootstrap(namespaces []string, source NamespaceSource, observedAt time.Time) []NamespaceRecord
	BootstrapAndDispatch(ctx context.Context, namespaces []string, source NamespaceSource, observedAt time.Time) ([]NamespaceRecord, error)
	Subscribe(subscriber NamespaceSubscriber)
	SnapshotSubscribers() []string
	Upsert(event NamespaceEvent) NamespaceRecord
	DispatchAdd(ctx context.Context, namespace string) (NamespaceRecord, bool, error)
	DispatchRemove(ctx context.Context, namespace string) (NamespaceRecord, bool, error)
	DispatchResync(ctx context.Context, namespace string) (NamespaceRecord, bool, error)
	MarkPartState(namespace string, part string, state NamespacePartState) (NamespaceRecord, bool)
	MarkPartRegistering(namespace string, part string) (NamespaceRecord, bool)
	MarkPartActive(namespace string, part string) (NamespaceRecord, bool)
	MarkPartFailed(namespace string, part string, err error) (NamespaceRecord, bool)
	Remove(name string) bool
}

type inMemoryNamespaceManager struct {
	mu      sync.RWMutex
	records map[string]NamespaceRecord
	subs    map[string]NamespaceSubscriber
}

func NewNamespaceManager() NamespaceManager {
	return &inMemoryNamespaceManager{
		records: make(map[string]NamespaceRecord),
		subs:    make(map[string]NamespaceSubscriber),
	}
}

func NewBootstrappedNamespaceManager(resolver *NamespaceResolver, source NamespaceSource, observedAt time.Time) NamespaceManager {
	manager := NewNamespaceManager()
	if resolver == nil {
		return manager
	}
	manager.Bootstrap(resolver.Snapshot(), source, observedAt)
	return manager
}

func NewWatcherNamespaceManager(ctx context.Context, namespaces []string, source NamespaceSource, observedAt time.Time, subscribers ...NamespaceSubscriber) (NamespaceManager, error) {
	manager := NewNamespaceManager()
	for _, subscriber := range subscribers {
		if subscriber == nil {
			continue
		}
		manager.Subscribe(subscriber)
	}
	_, err := manager.BootstrapAndDispatch(ctx, namespaces, source, observedAt)
	return manager, err
}

func NamespaceBecameUnmanaged(oldNamespace *corev1.Namespace, newNamespace *corev1.Namespace) bool {
	if oldNamespace == nil || newNamespace == nil {
		return false
	}
	return IsManagedNamespace(oldNamespace.Labels) && !IsManagedNamespace(newNamespace.Labels)
}

func DispatchNamespaceAdd(ctx context.Context, manager NamespaceManager, namespace *corev1.Namespace, source NamespaceSource, observedAt time.Time) (NamespaceRecord, bool, error) {
	if manager == nil || namespace == nil || namespace.Name == "" {
		return NamespaceRecord{}, false, nil
	}
	manager.Upsert(NamespaceEventFromNamespace(NamespaceEventAdd, namespace, source, observedAt))
	return manager.DispatchAdd(ctx, namespace.Name)
}

func DispatchNamespaceResync(ctx context.Context, manager NamespaceManager, namespace *corev1.Namespace, source NamespaceSource, observedAt time.Time) (NamespaceRecord, bool, error) {
	if manager == nil || namespace == nil || namespace.Name == "" {
		return NamespaceRecord{}, false, nil
	}
	manager.Upsert(NamespaceEventFromNamespace(NamespaceEventUpdate, namespace, source, observedAt))
	return manager.DispatchResync(ctx, namespace.Name)
}

func RecordNamespaceRemoval(manager NamespaceManager, obj interface{}, source NamespaceSource, observedAt time.Time) (NamespaceRecord, bool) {
	if manager == nil {
		return NamespaceRecord{}, false
	}
	event := NamespaceEventFromObject(NamespaceEventRemove, obj, source, observedAt)
	if event.Name == "" {
		return NamespaceRecord{}, false
	}
	record := manager.Upsert(event)
	return record, true
}

func (m *inMemoryNamespaceManager) Subscribe(subscriber NamespaceSubscriber) {
	m.mu.Lock()
	defer m.mu.Unlock()
	m.subs[subscriber.Name()] = subscriber
}

func (m *inMemoryNamespaceManager) Bootstrap(namespaces []string, source NamespaceSource, observedAt time.Time) []NamespaceRecord {
	records := make([]NamespaceRecord, 0, len(namespaces))
	for _, namespace := range namespaces {
		records = append(records, m.Upsert(NamespaceEvent{
			Type:       NamespaceEventAdd,
			Name:       namespace,
			Source:     source,
			ObservedAt: observedAt,
		}))
	}
	sort.Slice(records, func(i, j int) bool {
		return records[i].Name < records[j].Name
	})
	return records
}

func (m *inMemoryNamespaceManager) BootstrapAndDispatch(ctx context.Context, namespaces []string, source NamespaceSource, observedAt time.Time) ([]NamespaceRecord, error) {
	records := m.Bootstrap(namespaces, source, observedAt)
	var joinErr error
	for _, record := range records {
		updatedRecord, ok, err := m.DispatchAdd(ctx, record.Name)
		if ok {
			record = updatedRecord
		}
		if err != nil {
			joinErr = errors.Join(joinErr, err)
		}
		for index := range records {
			if records[index].Name == record.Name {
				records[index] = record
				break
			}
		}
	}
	return records, joinErr
}

func (m *inMemoryNamespaceManager) SnapshotSubscribers() []string {
	m.mu.RLock()
	defer m.mu.RUnlock()

	names := make([]string, 0, len(m.subs))
	for name := range m.subs {
		names = append(names, name)
	}
	sort.Strings(names)
	return names
}

func (m *inMemoryNamespaceManager) snapshotSubscriberObjects() []NamespaceSubscriber {
	m.mu.RLock()
	defer m.mu.RUnlock()

	names := make([]string, 0, len(m.subs))
	for name := range m.subs {
		names = append(names, name)
	}
	sort.Strings(names)

	subscribers := make([]NamespaceSubscriber, 0, len(names))
	for _, name := range names {
		subscribers = append(subscribers, m.subs[name])
	}
	return subscribers
}

func (m *inMemoryNamespaceManager) Snapshot() []string {
	m.mu.RLock()
	defer m.mu.RUnlock()

	namespaces := make([]string, 0, len(m.records))
	for name, record := range m.records {
		if record.Phase == NamespacePhaseRemoved {
			continue
		}
		namespaces = append(namespaces, name)
	}
	sort.Strings(namespaces)
	return namespaces
}

func (m *inMemoryNamespaceManager) SnapshotRecords() []NamespaceRecord {
	m.mu.RLock()
	defer m.mu.RUnlock()

	records := make([]NamespaceRecord, 0, len(m.records))
	for _, record := range m.records {
		records = append(records, record.Clone())
	}
	sort.Slice(records, func(i, j int) bool {
		return records[i].Name < records[j].Name
	})
	return records
}

func (m *inMemoryNamespaceManager) Get(name string) (NamespaceRecord, bool) {
	m.mu.RLock()
	defer m.mu.RUnlock()

	record, ok := m.records[name]
	if !ok {
		return NamespaceRecord{}, false
	}
	return record.Clone(), true
}

func (m *inMemoryNamespaceManager) Upsert(event NamespaceEvent) NamespaceRecord {
	m.mu.Lock()
	defer m.mu.Unlock()

	now := event.ObservedAt
	if now.IsZero() {
		now = time.Now().UTC()
	}

	record, exists := m.records[event.Name]
	if !exists {
		record = NamespaceRecord{
			Name:            event.Name,
			RegisteredParts: make(map[string]NamespacePartState),
		}
		if event.Type == NamespaceEventRemove {
			record.Generation = 0
		} else {
			record.Generation = 1
		}
	} else if event.Type != NamespaceEventRemove {
		record.Generation++
	}

	record.Name = event.Name
	record.Source = event.Source
	record.Labels = cloneStringMap(event.Labels)
	record.UpdatedAt = now

	switch event.Type {
	case NamespaceEventAdd, NamespaceEventUpdate, NamespaceEventResync:
		record.Phase = NamespacePhaseDiscovered
		record.LastError = ""
	case NamespaceEventRemove:
		record.Phase = NamespacePhaseRemoved
	}

	if record.RegisteredParts == nil {
		record.RegisteredParts = make(map[string]NamespacePartState)
	}

	m.records[event.Name] = record
	return record.Clone()
}

func (m *inMemoryNamespaceManager) DispatchAdd(ctx context.Context, namespace string) (NamespaceRecord, bool, error) {
	return m.dispatch(ctx, namespace, func(subscriber NamespaceSubscriber, record NamespaceRecord) error {
		return subscriber.OnNamespaceAdd(ctx, record)
	})
}

func (m *inMemoryNamespaceManager) DispatchRemove(ctx context.Context, namespace string) (NamespaceRecord, bool, error) {
	record, ok := m.Get(namespace)
	if !ok {
		return NamespaceRecord{}, false, nil
	}

	var firstErr error
	for _, subscriber := range m.snapshotSubscriberObjects() {
		err := subscriber.OnNamespaceRemove(ctx, record)
		if err != nil && firstErr == nil {
			firstErr = err
		}
	}

	removedRecord := m.Upsert(NamespaceEvent{
		Type:       NamespaceEventRemove,
		Name:       record.Name,
		Labels:     record.Labels,
		Source:     record.Source,
		ObservedAt: time.Now().UTC(),
	})
	return removedRecord, true, firstErr
}

func (m *inMemoryNamespaceManager) DispatchResync(ctx context.Context, namespace string) (NamespaceRecord, bool, error) {
	return m.dispatch(ctx, namespace, func(subscriber NamespaceSubscriber, record NamespaceRecord) error {
		return subscriber.OnNamespaceResync(ctx, record)
	})
}

func (m *inMemoryNamespaceManager) dispatch(ctx context.Context, namespace string, handler func(NamespaceSubscriber, NamespaceRecord) error) (NamespaceRecord, bool, error) {
	record, ok := m.Get(namespace)
	if !ok {
		return NamespaceRecord{}, false, nil
	}

	var firstErr error
	for _, subscriber := range m.snapshotSubscriberObjects() {
		_, _ = m.MarkPartRegistering(namespace, subscriber.Name())
		currentRecord, _ := m.Get(namespace)
		err := handler(subscriber, currentRecord)
		if err != nil {
			_, _ = m.MarkPartFailed(namespace, subscriber.Name(), err)
			if firstErr == nil {
				firstErr = err
			}
			continue
		}
		_, _ = m.MarkPartActive(namespace, subscriber.Name())
	}

	record, _ = m.Get(namespace)
	return record, true, firstErr
}

func (m *inMemoryNamespaceManager) MarkPartState(namespace string, part string, state NamespacePartState) (NamespaceRecord, bool) {
	m.mu.Lock()
	defer m.mu.Unlock()

	record, ok := m.records[namespace]
	if !ok {
		return NamespaceRecord{}, false
	}
	if record.RegisteredParts == nil {
		record.RegisteredParts = make(map[string]NamespacePartState)
	}
	if state.UpdatedAt.IsZero() {
		state.UpdatedAt = time.Now().UTC()
	}
	record.RegisteredParts[part] = state
	record.Phase = deriveNamespacePhase(record)
	record.UpdatedAt = state.UpdatedAt
	m.records[namespace] = record
	return record.Clone(), true
}

func (m *inMemoryNamespaceManager) MarkPartRegistering(namespace string, part string) (NamespaceRecord, bool) {
	return m.MarkPartState(namespace, part, NamespacePartState{State: NamespacePartStateRegistering})
}

func (m *inMemoryNamespaceManager) MarkPartActive(namespace string, part string) (NamespaceRecord, bool) {
	return m.MarkPartState(namespace, part, NamespacePartState{State: NamespacePartStateActive})
}

func (m *inMemoryNamespaceManager) MarkPartFailed(namespace string, part string, err error) (NamespaceRecord, bool) {
	lastError := ""
	if err != nil {
		lastError = err.Error()
	}
	return m.MarkPartState(namespace, part, NamespacePartState{State: NamespacePartStateFailed, LastError: lastError})
}

func (m *inMemoryNamespaceManager) Remove(name string) bool {
	m.mu.Lock()
	defer m.mu.Unlock()

	if _, ok := m.records[name]; !ok {
		return false
	}
	delete(m.records, name)
	return true
}

func deriveNamespacePhase(record NamespaceRecord) NamespacePhase {
	if len(record.RegisteredParts) == 0 {
		return record.Phase
	}

	hasRegistering := false
	for _, part := range record.RegisteredParts {
		switch part.State {
		case NamespacePartStateFailed:
			return NamespacePhaseFailed
		case NamespacePartStateRegistering:
			hasRegistering = true
		case NamespacePartStateActive:
			continue
		default:
			hasRegistering = true
		}
	}

	if hasRegistering {
		return NamespacePhaseRegistering
	}

	return NamespacePhaseActive
}