notarius/internal/framework/pipeline/checkpoint.go

package pipeline

import (
	"crypto/sha256"
	"encoding/hex"
	"fmt"
	"sort"
	"strings"

	"gitea.maximumdirect.net/eric/notarius/internal/core/source"
	"gitea.maximumdirect.net/eric/notarius/internal/framework/contracts"
)

type CheckpointFingerprint struct {
	Name  string `json:"name"`
	Value string `json:"value"`
}

type CheckpointRecorder interface {
	SourceRunning(moduleKey string) error
	SourceSucceeded(moduleKey string, doc *source.SourceDocument) error
	SourceFailed(moduleKey string, err error) error
	ChunkRunning(moduleKey string, sourceDigest string) error
	ChunkSucceeded(moduleKey string, sourceDigest string, chunks []contracts.SourceChunk, warnings []contracts.Warning) error
	ChunkRejected(moduleKey string, sourceDigest string, rejected contracts.RejectedOutput) error
	ChunkFailed(moduleKey string, sourceDigest string, err error) error
	ExtractRunning(laneID string, moduleKey string, dependencies []CheckpointFingerprint) error
	ExtractSucceeded(laneID string, moduleKey string, dependencies []CheckpointFingerprint, outputs []contracts.ExtractOutput, rejected []contracts.RejectedOutput, warnings []contracts.Warning) error
	ExtractFailed(laneID string, moduleKey string, dependencies []CheckpointFingerprint, err error) error
	MergeRunning(laneID string, moduleKey string, dependencies []CheckpointFingerprint) error
	MergeSucceeded(laneID string, moduleKey string, dependencies []CheckpointFingerprint, output contracts.MergeOutput, warnings []contracts.Warning) error
	MergeRejected(laneID string, moduleKey string, dependencies []CheckpointFingerprint, rejected contracts.RejectedOutput) error
	MergeFailed(laneID string, moduleKey string, dependencies []CheckpointFingerprint, err error) error
	NormalizeRunning(laneID string, moduleKey string, dependencies []CheckpointFingerprint) error
	NormalizeSucceeded(laneID string, moduleKey string, dependencies []CheckpointFingerprint, output contracts.NormalizeOutput, warnings []contracts.Warning) error
	NormalizeRejected(laneID string, moduleKey string, dependencies []CheckpointFingerprint, rejected contracts.RejectedOutput) error
	NormalizeFailed(laneID string, moduleKey string, dependencies []CheckpointFingerprint, err error) error
}

type noopCheckpointRecorder struct{}

func NoopCheckpointRecorder() CheckpointRecorder { return noopCheckpointRecorder{} }

func (noopCheckpointRecorder) SourceRunning(string) error                           { return nil }
func (noopCheckpointRecorder) SourceSucceeded(string, *source.SourceDocument) error { return nil }
func (noopCheckpointRecorder) SourceFailed(string, error) error                     { return nil }
func (noopCheckpointRecorder) ChunkRunning(string, string) error                    { return nil }
func (noopCheckpointRecorder) ChunkSucceeded(string, string, []contracts.SourceChunk, []contracts.Warning) error {
	return nil
}
func (noopCheckpointRecorder) ChunkRejected(string, string, contracts.RejectedOutput) error {
	return nil
}
func (noopCheckpointRecorder) ChunkFailed(string, string, error) error { return nil }
func (noopCheckpointRecorder) ExtractRunning(string, string, []CheckpointFingerprint) error {
	return nil
}
func (noopCheckpointRecorder) ExtractSucceeded(string, string, []CheckpointFingerprint, []contracts.ExtractOutput, []contracts.RejectedOutput, []contracts.Warning) error {
	return nil
}
func (noopCheckpointRecorder) ExtractFailed(string, string, []CheckpointFingerprint, error) error {
	return nil
}
func (noopCheckpointRecorder) MergeRunning(string, string, []CheckpointFingerprint) error { return nil }
func (noopCheckpointRecorder) MergeSucceeded(string, string, []CheckpointFingerprint, contracts.MergeOutput, []contracts.Warning) error {
	return nil
}
func (noopCheckpointRecorder) MergeRejected(string, string, []CheckpointFingerprint, contracts.RejectedOutput) error {
	return nil
}
func (noopCheckpointRecorder) MergeFailed(string, string, []CheckpointFingerprint, error) error {
	return nil
}
func (noopCheckpointRecorder) NormalizeRunning(string, string, []CheckpointFingerprint) error {
	return nil
}
func (noopCheckpointRecorder) NormalizeSucceeded(string, string, []CheckpointFingerprint, contracts.NormalizeOutput, []contracts.Warning) error {
	return nil
}
func (noopCheckpointRecorder) NormalizeRejected(string, string, []CheckpointFingerprint, contracts.RejectedOutput) error {
	return nil
}
func (noopCheckpointRecorder) NormalizeFailed(string, string, []CheckpointFingerprint, error) error {
	return nil
}

func rawOutputDigests(payloads []contracts.RawPayload) []CheckpointFingerprint {
	values := make([]CheckpointFingerprint, 0, len(payloads))
	for i, payload := range payloads {
		values = append(values, CheckpointFingerprint{
			Name:  fmt.Sprintf("payload[%d]", i),
			Value: checkpointContentDigest(payload.Content),
		})
	}
	return normalizeCheckpointFingerprints(values)
}

func extractPayloads(outputs []contracts.ExtractOutput) []contracts.RawPayload {
	if len(outputs) == 0 {
		return nil
	}
	payloads := make([]contracts.RawPayload, 0, len(outputs))
	for _, output := range outputs {
		payloads = append(payloads, output.Payload)
	}
	return payloads
}

func digestFingerprints(name string, digest string) []CheckpointFingerprint {
	digest = strings.TrimSpace(digest)
	if digest == "" {
		return nil
	}
	return []CheckpointFingerprint{{Name: name, Value: digest}}
}

func joinedChunkDigest(chunks []contracts.SourceChunk) string {
	if len(chunks) == 0 {
		return ""
	}
	values := make([]string, 0, len(chunks))
	for _, chunk := range chunks {
		values = append(values, chunk.ID+"="+checkpointContentDigest(chunk.Content))
	}
	sort.Strings(values)
	sum := sha256.Sum256([]byte(strings.Join(values, "\n")))
	return "sha256:" + hex.EncodeToString(sum[:])
}

func normalizeCheckpointFingerprints(values []CheckpointFingerprint) []CheckpointFingerprint {
	if len(values) == 0 {
		return nil
	}
	byName := make(map[string]string, len(values))
	for _, value := range values {
		name := strings.TrimSpace(value.Name)
		fingerprint := strings.TrimSpace(value.Value)
		if name == "" || fingerprint == "" {
			continue
		}
		byName[name] = fingerprint
	}
	if len(byName) == 0 {
		return nil
	}
	names := make([]string, 0, len(byName))
	for name := range byName {
		names = append(names, name)
	}
	sort.Strings(names)
	out := make([]CheckpointFingerprint, 0, len(names))
	for _, name := range names {
		out = append(out, CheckpointFingerprint{Name: name, Value: byName[name]})
	}
	return out
}

func checkpointContentDigest(content []byte) string {
	sum := sha256.Sum256(content)
	return "sha256:" + hex.EncodeToString(sum[:])
}