From 91678e9cc0b029b461e1002f75c6dfeada79b5a0 Mon Sep 17 00:00:00 2001 From: Tjark Gunnar Rasche Date: Thu, 20 Aug 2026 21:58:06 +0000 Subject: [PATCH] feat(sdk): expose snapshot diff facade Signed-off-by: Tjark Gunnar Rasche --- docs/integrator/go-library.md | 51 +++- docs/integrator/public-api.md | 6 +- pkg/cli/diff.go | 20 +- pkg/cli/diff_test.go | 16 +- pkg/client/v1/aicr.go | 10 +- pkg/client/v1/diff.go | 218 +++++++++++++++ pkg/client/v1/diff_test.go | 467 ++++++++++++++++++++++++++++++++ pkg/client/v1/example_test.go | 37 +++ pkg/client/v1/stability_test.go | 42 +++ pkg/diff/diff.go | 404 ++++++++++++++++++++++----- pkg/diff/diff_test.go | 127 +++++++++ pkg/diff/topology.go | 85 +++++- 12 files changed, 1369 insertions(+), 114 deletions(-) create mode 100644 pkg/client/v1/diff.go create mode 100644 pkg/client/v1/diff_test.go diff --git a/docs/integrator/go-library.md b/docs/integrator/go-library.md index ab6463854..d117f7c33 100644 --- a/docs/integrator/go-library.md +++ b/docs/integrator/go-library.md @@ -39,6 +39,7 @@ than in yours. | `Example_criteriaDimensions` | The coverage dimensions | yes | | `Example_committedConfig` | `AICRConfig` → source → catalog → criteria, in the required order | no | | `Example_resolveFromSnapshot` | `LoadSnapshot` plus snapshot criteria relaxation | no | +| `ExampleClient_DiffSnapshots` | In-memory drift detection between two loaded snapshots | no | | `ExampleClient_LoadRecipe` | Reading a previously emitted recipe | no | | `ExampleClient_CollectSnapshot` | Capturing cluster state via the snapshotter Job | no | | `ExampleClient_ValidateState` | Selecting validation phases, and `--no-cluster` mode | no | @@ -128,10 +129,10 @@ func main() { ## Snapshotting and validation Beyond recipe resolution, the facade exposes the rest of the -Snapshot → Validate workflow. Both methods are stateless w.r.t. the -Client's recipe source; they are surfaced through the Client only to -keep the facade uniform and leave room for future per-Client -telemetry hooks. +Snapshot → Validate workflow, including comparison of two snapshots for +configuration drift. These operations are stateless w.r.t. the Client's recipe +source; they are surfaced through the Client to keep the facade uniform and +leave room for future per-Client telemetry hooks. ### Loading a snapshot you already have @@ -174,6 +175,46 @@ identity with the loaded snapshot matters, such as hashing what you validated, capture the source contents yourself and load from that capture instead of re-reading afterwards. +### Comparing snapshots for drift + +`DiffSnapshots` compares the measurement payloads already held by two facade +snapshots. The comparison is in memory: it does not read a cluster or revisit +the file, URL, or ConfigMap the snapshots came from. + +```go +baseline, err := client.LoadSnapshot(ctx, "before.yaml", "") +if err != nil { + log.Fatalf("load baseline: %v", err) +} +target, err := client.LoadSnapshot(ctx, "after.yaml", "") +if err != nil { + log.Fatalf("load target: %v", err) +} + +result, err := client.DiffSnapshots(ctx, baseline, target, aicr.SnapshotDiffOptions{ + BaselineSource: "before.yaml", + TargetSource: "after.yaml", +}) +if err != nil { + log.Fatalf("diff snapshots: %v", err) +} +if result.HasDrift() { + log.Printf("detected %d change(s)", result.Summary.Total) +} +``` + +Drift is returned as data, not as an error. `SnapshotDiff.Changes` preserves +added, removed, and modified values, while `Summary` provides aggregate counts. +The source labels are optional output metadata and do not affect comparison. +Use `aicr.WriteSnapshotDiffTable` for the same human-readable table format as +`aicr diff`; JSON and YAML serializers can consume the facade-owned result +directly. + +Inputs must retain at least one typed measurement through `LoadSnapshot`, +`CollectSnapshot`, or `WrapSnapshot`. A hand-constructed `&aicr.Snapshot{}` or +a wrapped snapshot with no usable measurement is rejected instead of being +reported as no drift. + ### Capturing a snapshot from a live cluster ```go @@ -1069,6 +1110,8 @@ Per-operation caps: load whatever the source: a local file read, an HTTP(S) fetch, or a `cm://` ConfigMap read against the Kubernetes API. Distinct from `SnapshotOperationTimeout` below, which bounds deploying an agent Job. +- `DiffSnapshots`: **no facade cap** — comparison is in memory and the caller's + context governs unchanged. - `CollectSnapshot`: caller-controlled via `AgentConfig.Timeout` (falling back to `defaults.SnapshotOperationTimeout` when unset), plus `defaults.SnapshotOperationGrace`. The grace exists because diff --git a/docs/integrator/public-api.md b/docs/integrator/public-api.md index b2965df99..8d04d7bf0 100644 --- a/docs/integrator/public-api.md +++ b/docs/integrator/public-api.md @@ -31,7 +31,7 @@ in the [Go library integration guide](./go-library.md). | `pkg/bom` | Internal | Bill-of-materials / image inventory generation. | | `pkg/config` | Internal | Config-file loading and flag/spec resolution. | | `pkg/corroborate` | Internal | Cross-source corroboration of observed state. | -| `pkg/diff` | Internal | Structural diff between two snapshots. | +| `pkg/diff` | Internal | Structural snapshot comparison implementation. External consumers use `Client.DiffSnapshots` and `aicr.WriteSnapshotDiffTable`. | | `pkg/fingerprint` | Internal | Cluster/provider fingerprint detection. | | `pkg/health` | Internal | Health-check orchestration. | | `pkg/helm` | Internal | Helm chart rendering helpers. | @@ -65,6 +65,10 @@ unrelated exports in their evolving packages remain free to change. | Facade symbol | Translates to/from | Notes | |---|---|---| | `aicr.Snapshot` | `pkg/snapshotter.Snapshot` | **Facade-owned struct**. Public fields are identifying metadata; full measurement payload is preserved in an unexported field for round-trip through `ValidateState`. Obtain one from `Client.LoadSnapshot` (file, URL, or `cm://` ConfigMap) or `Client.CollectSnapshot` (live capture) — neither requires importing `pkg/snapshotter`. `aicr.WrapSnapshot` remains for the narrower case of lifting a `*snapshotter.Snapshot` you already hold from a direct `pkg/snapshotter` call. | +| `aicr.SnapshotDiff`, `aicr.SnapshotChange`, `aicr.SnapshotDiffSummary` | `pkg/diff` result shapes | **Facade-owned structs** returned by `Client.DiffSnapshots`. They preserve the CLI's JSON/YAML schema without exposing `pkg/diff` types. Drift is data (`SnapshotDiff.HasDrift`), while invalid or payload-less inputs and context cancellation are errors. | +| `aicr.SnapshotDiffOptions` | `Client.DiffSnapshots` input | **Facade-owned input struct** carrying optional baseline and target source labels. The labels are copied to output metadata and do not affect comparison semantics. | +| `aicr.SnapshotChangeKind` and its constants | `pkg/diff.ChangeKind` | **Facade-owned string enum** whose values describe added, removed, and modified readings. | +| `aicr.SnapshotChangeSeverity` and `aicr.SnapshotChangeSeverityInfo` | `pkg/diff.Severity` | **Facade-owned string enum** classifying change impact; informational is the currently defined severity. | | `aicr.AgentConfig` | `pkg/snapshotter.AgentConfig` | **Facade-owned struct** covering the deployment-time agent fields. `Tolerations` keeps `k8s.io/api/core/v1.Toleration` since `k8s.io` is itself a stable contract. It does **not** mirror every `pkg/snapshotter.AgentConfig` field — the network-collector fields `ClusterConfigPath` and `DiscoverNetwork` are not surfaced on the facade type. `AKSGPUPoolsPath` **is** surfaced (controller-side pool projection input, required for AKS profile-qualified resolution from a collected snapshot). | | `aicr.PhaseResult` | `pkg/validator.PhaseResult` | **Facade-owned struct**. Exposes `Summary` (CTRF counts) and `RawReport` (CTRF JSON bytes); `Report *ctrf.Report` is retained for in-tree consumers that merge per-phase reports. | | `aicr.Phase`, `aicr.PhaseDeployment` / `PhasePerformance` / `PhaseConformance` | string consts | **Facade-owned**. Values match `pkg/validator/v1` constants verbatim for byte-identical wire round-trip. | diff --git a/pkg/cli/diff.go b/pkg/cli/diff.go index 05dc2c076..360f2ff0c 100644 --- a/pkg/cli/diff.go +++ b/pkg/cli/diff.go @@ -24,8 +24,8 @@ import ( "github.com/urfave/cli/v3" + aicr "github.com/NVIDIA/aicr/pkg/client/v1" "github.com/NVIDIA/aicr/pkg/defaults" - "github.com/NVIDIA/aicr/pkg/diff" "github.com/NVIDIA/aicr/pkg/errors" "github.com/NVIDIA/aicr/pkg/serializer" ) @@ -122,13 +122,13 @@ func runDiffCmd(ctx context.Context, cmd *cli.Command) error { return err } - // Unwrap to reach pkg/diff, which still takes the internal shape. This is - // the last direct hop left in this command; exposing diff on the facade - // (#2025) is what removes it, and it was deliberately sequenced after - // LoadSnapshot so it can take facade snapshots rather than paths. - result := diff.Snapshots(baseline.Unwrap(), target.Unwrap()) - result.BaselineSource = baselinePath - result.TargetSource = targetPath + result, err := client.DiffSnapshots(ctx, baseline, target, aicr.SnapshotDiffOptions{ + BaselineSource: baselinePath, + TargetSource: targetPath, + }) + if err != nil { + return err + } slog.Info("snapshot diff complete", slog.Int("added", result.Summary.Added), @@ -154,7 +154,7 @@ func runDiffCmd(ctx context.Context, cmd *cli.Command) error { // // kubeconfig is propagated through to ConfigMap writers so multi-cluster // workflows write back to the same cluster the snapshots were read from. -func writeDiffResult(ctx context.Context, cmd *cli.Command, outFormat serializer.Format, kubeconfig string, result *diff.Result) (err error) { +func writeDiffResult(ctx context.Context, cmd *cli.Command, outFormat serializer.Format, kubeconfig string, result *aicr.SnapshotDiff) (err error) { output := cmd.String("output") // Use custom table writer for human-readable output @@ -176,7 +176,7 @@ func writeDiffResult(ctx context.Context, cmd *cli.Command, outFormat serializer }() w = f } - return diff.WriteTable(w, result) + return aicr.WriteSnapshotDiffTable(w, result) } // JSON/YAML use standard serializer; thread kubeconfig so ConfigMap diff --git a/pkg/cli/diff_test.go b/pkg/cli/diff_test.go index 8456065e2..4ce331a77 100644 --- a/pkg/cli/diff_test.go +++ b/pkg/cli/diff_test.go @@ -23,7 +23,7 @@ import ( "github.com/urfave/cli/v3" - "github.com/NVIDIA/aicr/pkg/diff" + aicr "github.com/NVIDIA/aicr/pkg/client/v1" "github.com/NVIDIA/aicr/pkg/serializer" ) @@ -111,13 +111,13 @@ func TestWriteTable_ToFile(t *testing.T) { tmpDir := t.TempDir() outFile := filepath.Join(tmpDir, "out.txt") - result := &diff.Result{Changes: make([]diff.Change, 0)} + result := &aicr.SnapshotDiff{Changes: make([]aicr.SnapshotChange, 0)} f, err := os.Create(outFile) if err != nil { t.Fatalf("failed to create output file: %v", err) } - err = diff.WriteTable(f, result) + err = aicr.WriteSnapshotDiffTable(f, result) if closeErr := f.Close(); closeErr != nil && err == nil { err = closeErr } @@ -135,10 +135,10 @@ func TestWriteTable_ToFile(t *testing.T) { } func TestWriteTable_ToStdout(t *testing.T) { - result := &diff.Result{Changes: make([]diff.Change, 0)} + result := &aicr.SnapshotDiff{Changes: make([]aicr.SnapshotChange, 0)} // WriteTable to stdout should not error. - err := diff.WriteTable(os.Stdout, result) + err := aicr.WriteSnapshotDiffTable(os.Stdout, result) if err != nil { t.Errorf("WriteTable to stdout failed: %v", err) } @@ -336,7 +336,7 @@ func TestWriteDiffResult_TableToFile(t *testing.T) { outFile := filepath.Join(tmpDir, "out.txt") cmd := buildDiffCommandWithOutput(t, outFile) - result := &diff.Result{Changes: make([]diff.Change, 0)} + result := &aicr.SnapshotDiff{Changes: make([]aicr.SnapshotChange, 0)} if err := writeDiffResult(t.Context(), cmd, serializer.FormatTable, "", result); err != nil { t.Fatalf("writeDiffResult failed: %v", err) @@ -358,7 +358,7 @@ func TestWriteDiffResult_CreateFails(t *testing.T) { bogusPath := filepath.Join(t.TempDir(), "does-not-exist", "out.txt") cmd := buildDiffCommandWithOutput(t, bogusPath) - result := &diff.Result{Changes: make([]diff.Change, 0)} + result := &aicr.SnapshotDiff{Changes: make([]aicr.SnapshotChange, 0)} err := writeDiffResult(t.Context(), cmd, serializer.FormatTable, "", result) if err == nil { @@ -380,7 +380,7 @@ func TestWriteDiffResult_KubeconfigPropagatesToConfigMap(t *testing.T) { bogusKubeconfig := filepath.Join(tmpDir, "missing-kubeconfig.yaml") cmd := buildDiffCommandWithOutput(t, "cm://aicr/test") - result := &diff.Result{Changes: make([]diff.Change, 0)} + result := &aicr.SnapshotDiff{Changes: make([]aicr.SnapshotChange, 0)} err := writeDiffResult(t.Context(), cmd, serializer.FormatJSON, bogusKubeconfig, result) if err == nil { diff --git a/pkg/client/v1/aicr.go b/pkg/client/v1/aicr.go index a56fdbe23..4af52457f 100644 --- a/pkg/client/v1/aicr.go +++ b/pkg/client/v1/aicr.go @@ -32,6 +32,8 @@ // - LoadSnapshot — read a previously captured *Snapshot from a file, // URL, or cm:// ConfigMap, for the common case where the snapshot // already exists and no cluster is needed. +// - DiffSnapshots — compare two loaded or collected snapshots in memory and +// return facade-owned field-level changes for drift detection. // - ValidateState — evaluate a resolved recipe against a snapshot, // running deployment / conformance / performance phases. // - LoadConfig — read and validate the AICRConfig a team commits, from a @@ -62,10 +64,10 @@ // - VerifyBinaryAttestation — package-level; prove an aicr binary was // built by NVIDIA CI. // -// All facade types (Snapshot, AgentConfig, Criteria, RecipeRequest, -// RecipeResult, ComponentBundle, ComponentRef, PhaseResult, and AllowLists) -// are facade-owned structs translated to and from the upstream pkg/* -// shapes, so internal field renames don't churn external callers. +// All facade types (Snapshot, SnapshotDiff, SnapshotChange, AgentConfig, +// Criteria, RecipeRequest, RecipeResult, ComponentBundle, ComponentRef, +// PhaseResult, AllowLists) are facade-owned structs translated to and from the +// upstream pkg/* shapes, so internal field renames don't churn external callers. // // Seven types remain deliberate transparent aliases: BundleConfig, // BundleAttester, BundleArtifact, OIDCResolveOptions, CriteriaRegistry, diff --git a/pkg/client/v1/diff.go b/pkg/client/v1/diff.go new file mode 100644 index 000000000..cb8976d08 --- /dev/null +++ b/pkg/client/v1/diff.go @@ -0,0 +1,218 @@ +// Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package aicr + +import ( + "context" + stderrors "errors" + "io" + + "github.com/NVIDIA/aicr/pkg/diff" + "github.com/NVIDIA/aicr/pkg/errors" + "github.com/NVIDIA/aicr/pkg/snapshotter" +) + +// SnapshotChangeKind describes how a snapshot value changed. +type SnapshotChangeKind string + +const ( + // SnapshotChangeAdded indicates a value exists only in the target snapshot. + SnapshotChangeAdded SnapshotChangeKind = "added" + // SnapshotChangeRemoved indicates a value exists only in the baseline snapshot. + SnapshotChangeRemoved SnapshotChangeKind = "removed" + // SnapshotChangeModified indicates a value differs between the snapshots. + SnapshotChangeModified SnapshotChangeKind = "modified" +) + +// SnapshotChangeSeverity classifies the impact of a snapshot change. +type SnapshotChangeSeverity string + +const ( + // SnapshotChangeSeverityInfo indicates an informational snapshot change. + SnapshotChangeSeverityInfo SnapshotChangeSeverity = "info" +) + +// SnapshotDiffOptions configures labels attached to a snapshot diff result. +// The labels identify the inputs in serialized output; they do not affect the +// comparison. +type SnapshotDiffOptions struct { + BaselineSource string + TargetSource string +} + +// SnapshotChange is one field-level difference between two snapshots. +// Baseline and Target are pointers so an absent side remains distinguishable +// from a present value whose string representation is empty. +type SnapshotChange struct { + Kind SnapshotChangeKind `json:"kind" yaml:"kind"` + Severity SnapshotChangeSeverity `json:"severity" yaml:"severity"` + Path string `json:"path" yaml:"path"` + Baseline *string `json:"baseline,omitempty" yaml:"baseline,omitempty"` + Target *string `json:"target,omitempty" yaml:"target,omitempty"` +} + +// SnapshotDiffSummary contains aggregate snapshot change counts. +type SnapshotDiffSummary struct { + Added int `json:"added" yaml:"added"` + Removed int `json:"removed" yaml:"removed"` + Modified int `json:"modified" yaml:"modified"` + Total int `json:"total" yaml:"total"` +} + +// SnapshotDiff contains the complete field-level comparison of two snapshots. +type SnapshotDiff struct { + BaselineSource string `json:"baselineSource,omitempty" yaml:"baselineSource,omitempty"` + TargetSource string `json:"targetSource,omitempty" yaml:"targetSource,omitempty"` + Changes []SnapshotChange `json:"changes" yaml:"changes"` + Summary SnapshotDiffSummary `json:"summary" yaml:"summary"` +} + +// HasDrift reports whether the diff contains any field-level changes. +func (r *SnapshotDiff) HasDrift() bool { + return r != nil && len(r.Changes) > 0 +} + +// DiffSnapshots compares two facade snapshots in memory. +// +// Both snapshots must carry a usable measurement payload from LoadSnapshot, +// CollectSnapshot, or WrapSnapshot. A hand-constructed Snapshot has no such +// payload; a wrapped payload containing no typed measurement is equally +// unusable. Both are rejected rather than being reported as a false no-drift +// result. Source labels in opts are copied to the result for JSON, YAML, and +// table consumers; they do not influence comparison semantics. +// +// The operation performs no cluster, filesystem, or recipe-catalog I/O and +// adds no facade timeout. The caller's context governs unchanged. +func (c *Client) DiffSnapshots( + ctx context.Context, + baseline, target *Snapshot, + opts SnapshotDiffOptions, +) (*SnapshotDiff, error) { + + if c == nil { + return nil, errors.New(errors.ErrCodeInvalidRequest, "aicr client not initialized") + } + if ctx == nil { + return nil, errors.New(errors.ErrCodeInvalidRequest, "context is required (got nil)") + } + baselineInternal, err := requireSnapshotDiffPayload(baseline, "baseline") + if err != nil { + return nil, err + } + targetInternal, err := requireSnapshotDiffPayload(target, "target") + if err != nil { + return nil, err + } + if openErr := c.assertOpen(); openErr != nil { + return nil, openErr + } + result, err := diff.SnapshotsWithContext(ctx, baselineInternal, targetInternal) + if err != nil { + return nil, err + } + + out := &SnapshotDiff{ + BaselineSource: opts.BaselineSource, + TargetSource: opts.TargetSource, + Changes: make([]SnapshotChange, len(result.Changes)), + Summary: SnapshotDiffSummary{ + Added: result.Summary.Added, + Removed: result.Summary.Removed, + Modified: result.Summary.Modified, + Total: result.Summary.Total, + }, + } + for i := range result.Changes { + if err := snapshotDiffContextError(ctx); err != nil { + return nil, err + } + out.Changes[i] = SnapshotChange{ + Kind: SnapshotChangeKind(result.Changes[i].Kind), + Severity: SnapshotChangeSeverity(result.Changes[i].Severity), + Path: result.Changes[i].Path, + Baseline: copySnapshotDiffString(result.Changes[i].Baseline), + Target: copySnapshotDiffString(result.Changes[i].Target), + } + } + return out, nil +} + +// WriteSnapshotDiffTable writes a human-readable snapshot diff table. +func WriteSnapshotDiffTable(w io.Writer, result *SnapshotDiff) error { + if w == nil { + return errors.New(errors.ErrCodeInvalidRequest, "snapshot diff table writer is required (got nil)") + } + if result == nil { + return errors.New(errors.ErrCodeInvalidRequest, "snapshot diff result is required (got nil)") + } + + internal := &diff.Result{ + BaselineSource: result.BaselineSource, + TargetSource: result.TargetSource, + Changes: make([]diff.Change, len(result.Changes)), + Summary: diff.Summary{ + Added: result.Summary.Added, + Removed: result.Summary.Removed, + Modified: result.Summary.Modified, + Total: result.Summary.Total, + }, + } + for i := range result.Changes { + internal.Changes[i] = diff.Change{ + Kind: diff.ChangeKind(result.Changes[i].Kind), + Severity: diff.Severity(result.Changes[i].Severity), + Path: result.Changes[i].Path, + Baseline: copySnapshotDiffString(result.Changes[i].Baseline), + Target: copySnapshotDiffString(result.Changes[i].Target), + } + } + return diff.WriteTable(w, internal) +} + +func requireSnapshotDiffPayload(s *Snapshot, role string) (*snapshotter.Snapshot, error) { + if s == nil { + return nil, errors.New(errors.ErrCodeInvalidRequest, role+" snapshot is required (got nil)") + } + if s.internal == nil { + return nil, errors.New(errors.ErrCodeInvalidRequest, + role+" snapshot has no measurement payload; use Client.LoadSnapshot, Client.CollectSnapshot, or aicr.WrapSnapshot") + } + for _, m := range s.internal.Measurements { + if m != nil && m.Type != "" { + return s.internal, nil + } + } + return nil, errors.New(errors.ErrCodeInvalidRequest, + role+" snapshot has no usable measurement payload; use Client.LoadSnapshot, Client.CollectSnapshot, or aicr.WrapSnapshot") +} + +func snapshotDiffContextError(ctx context.Context) error { + err := ctx.Err() + if err == nil { + return nil + } + if stderrors.Is(err, context.Canceled) { + return errors.Wrap(errors.ErrCodeCanceled, "snapshot diff canceled", err) + } + return errors.Wrap(errors.ErrCodeTimeout, "snapshot diff deadline exceeded", err) +} + +func copySnapshotDiffString(value *string) *string { + if value == nil { + return nil + } + copied := *value + return &copied +} diff --git a/pkg/client/v1/diff_test.go b/pkg/client/v1/diff_test.go new file mode 100644 index 000000000..2d0632a7a --- /dev/null +++ b/pkg/client/v1/diff_test.go @@ -0,0 +1,467 @@ +// Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package aicr_test + +import ( + "bytes" + "context" + stderrors "errors" + "fmt" + "io" + "reflect" + "strings" + "testing" + "time" + + aicr "github.com/NVIDIA/aicr/pkg/client/v1" + aicrerrors "github.com/NVIDIA/aicr/pkg/errors" + "github.com/NVIDIA/aicr/pkg/measurement" + "github.com/NVIDIA/aicr/pkg/snapshotter" +) + +func TestDiffSnapshots_Guards(t *testing.T) { + client := newVerifyClient(t) + closed := newClosedClient(t) + valid := diffTestSnapshot(nil, nil, nil) + wrappedEmpty := aicr.WrapSnapshot(&snapshotter.Snapshot{}) + wrappedNil := aicr.WrapSnapshot(&snapshotter.Snapshot{ + Measurements: []*measurement.Measurement{nil}, + }) + wrappedTypeless := aicr.WrapSnapshot(&snapshotter.Snapshot{ + Measurements: []*measurement.Measurement{{}}, + }) + + tests := []struct { + name string + client *aicr.Client + ctx context.Context + baseline *aicr.Snapshot + target *aicr.Snapshot + }{ + {name: "nil client", client: nil, ctx: t.Context(), baseline: valid, target: valid}, + {name: "nil context", client: client, ctx: nil, baseline: valid, target: valid}, + {name: "nil baseline", client: client, ctx: t.Context(), baseline: nil, target: valid}, + {name: "nil target", client: client, ctx: t.Context(), baseline: valid, target: nil}, + {name: "hand-constructed baseline has no payload", client: client, ctx: t.Context(), baseline: &aicr.Snapshot{}, target: valid}, + {name: "hand-constructed target has no payload", client: client, ctx: t.Context(), baseline: valid, target: &aicr.Snapshot{}}, + {name: "wrapped empty baseline has no measurements", client: client, ctx: t.Context(), baseline: wrappedEmpty, target: valid}, + {name: "wrapped nil baseline has no usable measurements", client: client, ctx: t.Context(), baseline: wrappedNil, target: valid}, + {name: "wrapped typeless target has no usable measurements", client: client, ctx: t.Context(), baseline: valid, target: wrappedTypeless}, + {name: "closed client", client: closed, ctx: t.Context(), baseline: valid, target: valid}, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + _, err := tt.client.DiffSnapshots(tt.ctx, tt.baseline, tt.target, aicr.SnapshotDiffOptions{}) + if err == nil { + t.Fatal("DiffSnapshots() error = nil, want ErrCodeInvalidRequest") + } + if !stderrors.Is(err, aicrerrors.New(aicrerrors.ErrCodeInvalidRequest, "")) { + t.Errorf("DiffSnapshots() error = %v, want ErrCodeInvalidRequest", err) + } + }) + } +} + +func TestDiffSnapshots_CoreChanges(t *testing.T) { + client := newVerifyClient(t) + emptyData := map[string]measurement.Reading{} + + tests := []struct { + name string + baseline *aicr.Snapshot + target *aicr.Snapshot + wantChanges []aicr.SnapshotChange + wantSummary aicr.SnapshotDiffSummary + }{ + { + name: "no drift", + baseline: diffTestSnapshot(map[string]measurement.Reading{"version": measurement.Str("1.31.0")}, nil, nil), + target: diffTestSnapshot(map[string]measurement.Reading{"version": measurement.Str("1.31.0")}, nil, nil), + wantChanges: []aicr.SnapshotChange{}, + }, + { + name: "added", + baseline: diffTestSnapshot(emptyData, nil, nil), + target: diffTestSnapshot(map[string]measurement.Reading{"version": measurement.Str("1.32.0")}, nil, nil), + wantChanges: []aicr.SnapshotChange{{ + Kind: aicr.SnapshotChangeAdded, + Severity: aicr.SnapshotChangeSeverityInfo, + Path: "K8s.server.version", + Target: diffTestString("1.32.0"), + }}, + wantSummary: aicr.SnapshotDiffSummary{Added: 1, Total: 1}, + }, + { + name: "removed", + baseline: diffTestSnapshot(map[string]measurement.Reading{"version": measurement.Str("1.31.0")}, nil, nil), + target: diffTestSnapshot(emptyData, nil, nil), + wantChanges: []aicr.SnapshotChange{{ + Kind: aicr.SnapshotChangeRemoved, + Severity: aicr.SnapshotChangeSeverityInfo, + Path: "K8s.server.version", + Baseline: diffTestString("1.31.0"), + }}, + wantSummary: aicr.SnapshotDiffSummary{Removed: 1, Total: 1}, + }, + { + name: "modified", + baseline: diffTestSnapshot(map[string]measurement.Reading{"version": measurement.Str("1.31.0")}, nil, nil), + target: diffTestSnapshot(map[string]measurement.Reading{"version": measurement.Str("1.32.0")}, nil, nil), + wantChanges: []aicr.SnapshotChange{{ + Kind: aicr.SnapshotChangeModified, + Severity: aicr.SnapshotChangeSeverityInfo, + Path: "K8s.server.version", + Baseline: diffTestString("1.31.0"), + Target: diffTestString("1.32.0"), + }}, + wantSummary: aicr.SnapshotDiffSummary{Modified: 1, Total: 1}, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + result, err := client.DiffSnapshots(t.Context(), tt.baseline, tt.target, aicr.SnapshotDiffOptions{ + BaselineSource: "before.yaml", + TargetSource: "after.yaml", + }) + if err != nil { + t.Fatalf("DiffSnapshots() error = %v", err) + } + if result.BaselineSource != "before.yaml" || result.TargetSource != "after.yaml" { + t.Errorf("sources = %q, %q, want before.yaml, after.yaml", result.BaselineSource, result.TargetSource) + } + if !reflect.DeepEqual(result.Changes, tt.wantChanges) { + t.Errorf("changes = %#v, want %#v", result.Changes, tt.wantChanges) + } + if result.Summary != tt.wantSummary { + t.Errorf("summary = %#v, want %#v", result.Summary, tt.wantSummary) + } + if result.HasDrift() != (len(tt.wantChanges) > 0) { + t.Errorf("HasDrift() = %v, want %v", result.HasDrift(), len(tt.wantChanges) > 0) + } + }) + } +} + +func TestDiffSnapshots_PreservesEmptyAndStructuredValues(t *testing.T) { + client := newVerifyClient(t) + + t.Run("explicit empty remains present", func(t *testing.T) { + result, err := client.DiffSnapshots(t.Context(), + diffTestSnapshot(map[string]measurement.Reading{"version": measurement.Str("1.31.0")}, nil, nil), + diffTestSnapshot(map[string]measurement.Reading{"version": measurement.Str("")}, nil, nil), + aicr.SnapshotDiffOptions{}) + if err != nil { + t.Fatalf("DiffSnapshots() error = %v", err) + } + if len(result.Changes) != 1 { + t.Fatalf("changes = %d, want 1", len(result.Changes)) + } + change := result.Changes[0] + if change.Target == nil || *change.Target != "" { + t.Errorf("target = %#v, want non-nil pointer to empty string", change.Target) + } + if change.Baseline == nil || *change.Baseline != "1.31.0" { + t.Errorf("baseline = %#v, want 1.31.0", change.Baseline) + } + }) + + t.Run("context and items remain field-level changes", func(t *testing.T) { + baselineItems := []measurement.ItemEntry{{ + Context: map[string]string{"name": "pf0"}, + Data: map[string]measurement.Reading{"mtu": measurement.Int(1500)}, + }} + targetItems := []measurement.ItemEntry{{ + Context: map[string]string{"name": "pf1"}, + Data: map[string]measurement.Reading{"mtu": measurement.Int(9000)}, + }} + result, err := client.DiffSnapshots(t.Context(), + diffTestSnapshot(nil, map[string]string{"node": "n1"}, baselineItems), + diffTestSnapshot(nil, map[string]string{"node": "n2"}, targetItems), + aicr.SnapshotDiffOptions{}) + if err != nil { + t.Fatalf("DiffSnapshots() error = %v", err) + } + wantPaths := []string{ + "K8s.server.context.node", + "K8s.server.items[0].context.name", + "K8s.server.items[0].data.mtu", + } + gotPaths := make([]string, len(result.Changes)) + for i := range result.Changes { + gotPaths[i] = result.Changes[i].Path + } + if !reflect.DeepEqual(gotPaths, wantPaths) { + t.Errorf("paths = %v, want %v", gotPaths, wantPaths) + } + }) +} + +func TestDiffSnapshots_ContextCancellation(t *testing.T) { + client := newVerifyClient(t) + valid := diffTestSnapshot(nil, nil, nil) + canceled, cancel := context.WithCancel(t.Context()) + cancel() + expired, expire := context.WithDeadline(t.Context(), time.Now().Add(-time.Second)) + defer expire() + + tests := []struct { + name string + ctx context.Context + wantCode aicrerrors.ErrorCode + wantCause error + }{ + {name: "canceled", ctx: canceled, wantCode: aicrerrors.ErrCodeCanceled, wantCause: context.Canceled}, + {name: "deadline", ctx: expired, wantCode: aicrerrors.ErrCodeTimeout, wantCause: context.DeadlineExceeded}, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + _, err := client.DiffSnapshots(tt.ctx, valid, valid, aicr.SnapshotDiffOptions{}) + if err == nil { + t.Fatalf("DiffSnapshots() error = nil, want %s", tt.wantCode) + } + if !stderrors.Is(err, aicrerrors.New(tt.wantCode, "")) { + t.Errorf("error = %v, want code %s", err, tt.wantCode) + } + if !stderrors.Is(err, tt.wantCause) { + t.Errorf("error = %v, want cause %v", err, tt.wantCause) + } + }) + } +} + +func TestDiffSnapshots_MidTraversalContextCancellation(t *testing.T) { + client := newVerifyClient(t) + baselineData := make(map[string]measurement.Reading, 64) + targetData := make(map[string]measurement.Reading, 64) + for i := 0; i < 64; i++ { + key := fmt.Sprintf("reading-%02d", i) + baselineData[key] = measurement.Int(i) + targetData[key] = measurement.Int(i + 1) + } + baseline := diffTestSnapshot(baselineData, nil, nil) + target := diffTestSnapshot(targetData, nil, nil) + probeCtx := &snapshotDiffCountingContext{Context: t.Context()} + probeResult, err := client.DiffSnapshots(probeCtx, baseline, target, aicr.SnapshotDiffOptions{}) + if err != nil { + t.Fatalf("DiffSnapshots() probe error = %v", err) + } + // Exclude the facade's one mapping checkpoint per change, then cancel + // halfway through the comparison's final summary traversal. This proves + // the internal partial result is discarded rather than failing in mapping. + comparisonChecks := probeCtx.checks - len(probeResult.Changes) + cancelAt := comparisonChecks - probeResult.Summary.Total/2 + if cancelAt <= 0 || cancelAt >= comparisonChecks { + t.Fatalf("derived cancellation checkpoint = %d, comparison checks = %d", cancelAt, comparisonChecks) + } + + tests := []struct { + name string + cause error + wantCode aicrerrors.ErrorCode + }{ + {name: "canceled", cause: context.Canceled, wantCode: aicrerrors.ErrCodeCanceled}, + {name: "deadline", cause: context.DeadlineExceeded, wantCode: aicrerrors.ErrCodeTimeout}, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + ctx := newSnapshotDiffCheckpointContext(t.Context(), cancelAt, tt.cause) + result, err := client.DiffSnapshots(ctx, baseline, target, aicr.SnapshotDiffOptions{}) + if result != nil { + t.Fatalf("DiffSnapshots() result = %#v, want nil after cancellation", result) + } + if !stderrors.Is(err, aicrerrors.New(tt.wantCode, "")) { + t.Errorf("DiffSnapshots() error = %v, want code %s", err, tt.wantCode) + } + if !stderrors.Is(err, tt.cause) { + t.Errorf("DiffSnapshots() error = %v, want cause %v", err, tt.cause) + } + if ctx.checks < cancelAt { + t.Errorf("context checks = %d, want at least %d to prove traversal began", ctx.checks, cancelAt) + } + }) + } +} + +func TestSnapshotDiff_HasDrift(t *testing.T) { + tests := []struct { + name string + result *aicr.SnapshotDiff + want bool + }{ + {name: "nil", result: nil, want: false}, + {name: "no changes", result: &aicr.SnapshotDiff{Changes: []aicr.SnapshotChange{}}, want: false}, + {name: "change present", result: &aicr.SnapshotDiff{Changes: []aicr.SnapshotChange{{Path: "K8s.server.version"}}}, want: true}, + {name: "summary alone does not imply drift", result: &aicr.SnapshotDiff{Summary: aicr.SnapshotDiffSummary{Total: 1}}, want: false}, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + if got := tt.result.HasDrift(); got != tt.want { + t.Errorf("HasDrift() = %v, want %v", got, tt.want) + } + }) + } +} + +func TestWriteSnapshotDiffTable(t *testing.T) { + empty := "" + result := &aicr.SnapshotDiff{ + Changes: []aicr.SnapshotChange{{ + Kind: aicr.SnapshotChangeAdded, + Severity: aicr.SnapshotChangeSeverityInfo, + Path: "K8s.server.version", + Target: &empty, + }}, + Summary: aicr.SnapshotDiffSummary{Added: 1, Total: 1}, + } + + t.Run("renders absent and explicit empty distinctly", func(t *testing.T) { + var buf bytes.Buffer + if err := aicr.WriteSnapshotDiffTable(&buf, result); err != nil { + t.Fatalf("WriteSnapshotDiffTable() error = %v", err) + } + output := buf.String() + missingExpectedContent := !strings.Contains(output, "CHANGES (1 added, 0 removed, 0 modified)") || + !strings.Contains(output, `K8s.server.version - ""`) || + !strings.Contains(output, "DRIFT DETECTED") + if missingExpectedContent { + t.Errorf("table output missing expected content:\n%s", output) + } + }) + + t.Run("no changes", func(t *testing.T) { + var buf bytes.Buffer + if err := aicr.WriteSnapshotDiffTable(&buf, &aicr.SnapshotDiff{Changes: []aicr.SnapshotChange{}}); err != nil { + t.Fatalf("WriteSnapshotDiffTable() error = %v", err) + } + if got := strings.TrimSpace(buf.String()); got != "NO CHANGES" { + t.Errorf("output = %q, want NO CHANGES", got) + } + }) + + t.Run("nil inputs", func(t *testing.T) { + tests := []struct { + name string + writer io.Writer + result *aicr.SnapshotDiff + }{ + {name: "nil writer", writer: nil, result: result}, + {name: "nil result", writer: io.Discard, result: nil}, + } + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + err := aicr.WriteSnapshotDiffTable(tt.writer, tt.result) + if !stderrors.Is(err, aicrerrors.New(aicrerrors.ErrCodeInvalidRequest, "")) { + t.Errorf("error = %v, want ErrCodeInvalidRequest", err) + } + }) + } + }) + + t.Run("writer failure", func(t *testing.T) { + writeErr := stderrors.New("write failed") + err := aicr.WriteSnapshotDiffTable(snapshotDiffFailWriter{err: writeErr}, result) + if err == nil { + t.Fatal("WriteSnapshotDiffTable() error = nil, want writer failure") + } + if !stderrors.Is(err, writeErr) { + t.Errorf("error = %v, want wrapped writer failure", err) + } + }) +} + +type snapshotDiffFailWriter struct { + err error +} + +func (w snapshotDiffFailWriter) Write(_ []byte) (int, error) { + return 0, w.err +} + +type snapshotDiffCheckpointContext struct { + context.Context + cancelAt int + cause error + done chan struct{} + checks int + closed bool +} + +type snapshotDiffCountingContext struct { + context.Context + checks int +} + +func (c *snapshotDiffCountingContext) Err() error { + c.checks++ + return c.Context.Err() +} + +func newSnapshotDiffCheckpointContext( + parent context.Context, + cancelAt int, + cause error, +) *snapshotDiffCheckpointContext { + + return &snapshotDiffCheckpointContext{ + Context: parent, + cancelAt: cancelAt, + cause: cause, + done: make(chan struct{}), + } +} + +func (c *snapshotDiffCheckpointContext) Done() <-chan struct{} { + return c.done +} + +func (c *snapshotDiffCheckpointContext) Err() error { + c.checks++ + if c.checks < c.cancelAt { + return nil + } + if !c.closed { + close(c.done) + c.closed = true + } + return c.cause +} + +func diffTestSnapshot( + data map[string]measurement.Reading, + contextValues map[string]string, + items []measurement.ItemEntry, +) *aicr.Snapshot { + + return aicr.WrapSnapshot(&snapshotter.Snapshot{ + Measurements: []*measurement.Measurement{{ + Type: measurement.TypeK8s, + Subtypes: []measurement.Subtype{{ + Name: "server", + Data: data, + Context: contextValues, + Items: items, + }}, + }}, + }) +} + +func diffTestString(value string) *string { + return &value +} diff --git a/pkg/client/v1/example_test.go b/pkg/client/v1/example_test.go index 8fc30b84c..7f71d8d51 100644 --- a/pkg/client/v1/example_test.go +++ b/pkg/client/v1/example_test.go @@ -224,6 +224,43 @@ func Example_resolveFromSnapshot() { } } +// ExampleClient_DiffSnapshots compares two previously captured snapshots in +// memory. Loading local files needs no cluster access; cm:// sources use the +// kubeconfig argument passed to LoadSnapshot. +func ExampleClient_DiffSnapshots() { + ctx := context.Background() + + client, err := aicr.NewClient(aicr.WithRecipeSource(aicr.EmbeddedSource())) + if err != nil { + log.Print(err) + return + } + defer func() { _ = client.Close() }() + + baseline, err := client.LoadSnapshot(ctx, "before.yaml", "") + if err != nil { + log.Print(err) + return + } + target, err := client.LoadSnapshot(ctx, "after.yaml", "") + if err != nil { + log.Print(err) + return + } + + result, err := client.DiffSnapshots(ctx, baseline, target, aicr.SnapshotDiffOptions{ + BaselineSource: "before.yaml", + TargetSource: "after.yaml", + }) + if err != nil { + log.Print(err) + return + } + if result.HasDrift() { + fmt.Printf("detected %d change(s)\n", result.Summary.Total) + } +} + // Example_bundleAndVerify is the integrator path end to end: resolve a recipe, // render its deployment bundle, then check what was written. // diff --git a/pkg/client/v1/stability_test.go b/pkg/client/v1/stability_test.go index 576a7ced2..c262ec74c 100644 --- a/pkg/client/v1/stability_test.go +++ b/pkg/client/v1/stability_test.go @@ -29,6 +29,7 @@ package aicr_test import ( "context" + "io" "testing" "time" @@ -105,6 +106,47 @@ func TestStability_RecipeResolution(t *testing.T) { } } +// TestStability_SnapshotDiff pins the facade-owned drift-detection surface. +func TestStability_SnapshotDiff(t *testing.T) { + t.Parallel() + + requireSignature[func(*aicr.Client, context.Context, *aicr.Snapshot, *aicr.Snapshot, aicr.SnapshotDiffOptions) (*aicr.SnapshotDiff, error)]((*aicr.Client).DiffSnapshots) + requireSignature[func(io.Writer, *aicr.SnapshotDiff) error](aicr.WriteSnapshotDiffTable) + requireSignature[func(*aicr.SnapshotDiff) bool]((*aicr.SnapshotDiff).HasDrift) + + var opts aicr.SnapshotDiffOptions + _ = opts.BaselineSource + _ = opts.TargetSource + + var result aicr.SnapshotDiff + _ = result.BaselineSource + _ = result.TargetSource + _ = result.Changes + _ = result.Summary + + var change aicr.SnapshotChange + _ = change.Kind + _ = change.Severity + _ = change.Path + _ = change.Baseline + _ = change.Target + + var summary aicr.SnapshotDiffSummary + _ = summary.Added + _ = summary.Removed + _ = summary.Modified + _ = summary.Total + + _ = []aicr.SnapshotChangeKind{ + aicr.SnapshotChangeAdded, + aicr.SnapshotChangeRemoved, + aicr.SnapshotChangeModified, + } + _ = []aicr.SnapshotChangeSeverity{ + aicr.SnapshotChangeSeverityInfo, + } +} + func requireSignature[T any](_ T) {} // TestStability_RecipeResult pins the consumer-visible fields and methods diff --git a/pkg/diff/diff.go b/pkg/diff/diff.go index f4537872f..ef06a4266 100644 --- a/pkg/diff/diff.go +++ b/pkg/diff/diff.go @@ -18,12 +18,15 @@ package diff import ( + "context" + stderrors "errors" "fmt" "reflect" "sort" "strconv" "strings" + "github.com/NVIDIA/aicr/pkg/errors" "github.com/NVIDIA/aicr/pkg/measurement" "github.com/NVIDIA/aicr/pkg/snapshotter" ) @@ -110,41 +113,108 @@ func (r *Result) HasDrift() bool { // The baseline is the reference state; the target is the current state. // If either baseline or target is nil, returns an empty Result (no drift). func Snapshots(baseline, target *snapshotter.Snapshot) *Result { - if baseline == nil || target == nil { + result, err := snapshots(context.Background(), baseline, target) + if err != nil { + // context.Background cannot be canceled. Keep the legacy no-error API + // total if that invariant ever changes. return &Result{Changes: make([]Change, 0)} } + return result +} + +// SnapshotsWithContext compares two snapshots while honoring cancellation +// throughout the in-memory traversal. It returns no partial result when the +// context is canceled or its deadline expires. +func SnapshotsWithContext(ctx context.Context, baseline, target *snapshotter.Snapshot) (*Result, error) { + if ctx == nil { + return nil, errors.New(errors.ErrCodeInvalidRequest, "snapshot diff context is required (got nil)") + } + + result, err := snapshots(ctx, baseline, target) + if err != nil { + return nil, snapshotContextError(err) + } + return result, nil +} + +func snapshots(ctx context.Context, baseline, target *snapshotter.Snapshot) (*Result, error) { + if err := ctx.Err(); err != nil { + return nil, err + } + if baseline == nil || target == nil { + return &Result{Changes: make([]Change, 0)}, nil + } result := &Result{ Changes: make([]Change, 0), } - baseByType := indexMeasurements(baseline.Measurements) - targetByType := indexMeasurements(target.Measurements) + baseByType, err := indexMeasurements(ctx, baseline.Measurements) + if err != nil { + return nil, err + } + targetByType, err := indexMeasurements(ctx, target.Measurements) + if err != nil { + return nil, err + } - allTypes := mergeKeys(baseByType, targetByType) + allTypes, err := mergeKeys(ctx, baseByType, targetByType) + if err != nil { + return nil, err + } + if err := ctx.Err(); err != nil { + return nil, err + } sort.Strings(allTypes) + if err := ctx.Err(); err != nil { + return nil, err + } for _, typeName := range allTypes { + if err := ctx.Err(); err != nil { + return nil, err + } baseMeasurement, baseExists := baseByType[typeName] targetMeasurement, targetExists := targetByType[typeName] if !baseExists { - result.Changes = append(result.Changes, addedMeasurement(targetMeasurement)...) + changes, err := addedMeasurement(ctx, targetMeasurement) + if err != nil { + return nil, err + } + result.Changes = append(result.Changes, changes...) continue } if !targetExists { - result.Changes = append(result.Changes, removedMeasurement(baseMeasurement)...) + changes, err := removedMeasurement(ctx, baseMeasurement) + if err != nil { + return nil, err + } + result.Changes = append(result.Changes, changes...) continue } - result.Changes = append(result.Changes, compareMeasurements(baseMeasurement, targetMeasurement)...) + changes, err := compareMeasurements(ctx, baseMeasurement, targetMeasurement) + if err != nil { + return nil, err + } + result.Changes = append(result.Changes, changes...) } + if err := ctx.Err(); err != nil { + return nil, err + } sort.Slice(result.Changes, func(i, j int) bool { return result.Changes[i].Path < result.Changes[j].Path }) + if err := ctx.Err(); err != nil { + return nil, err + } for _, c := range result.Changes { + if err := ctx.Err(); err != nil { + return nil, err + } switch c.Kind { case Added: result.Summary.Added++ @@ -155,8 +225,11 @@ func Snapshots(baseline, target *snapshotter.Snapshot) *Result { } } result.Summary.Total = len(result.Changes) + if err := ctx.Err(); err != nil { + return nil, err + } - return result + return result, nil } // --- helpers --- @@ -180,58 +253,115 @@ func safeReadingString(r measurement.Reading) string { return r.String() } -func indexMeasurements(measurements []*measurement.Measurement) map[string]*measurement.Measurement { +func indexMeasurements(ctx context.Context, measurements []*measurement.Measurement) (map[string]*measurement.Measurement, error) { idx := make(map[string]*measurement.Measurement, len(measurements)) for _, m := range measurements { + if err := ctx.Err(); err != nil { + return nil, err + } if m == nil { continue } idx[string(m.Type)] = m } - return idx + return idx, nil } -func compareMeasurements(base, target *measurement.Measurement) []Change { +func compareMeasurements(ctx context.Context, base, target *measurement.Measurement) ([]Change, error) { var changes []Change - base, target = alignTopologyEncoding(base, target) + var err error + base, target, err = alignTopologyEncoding(ctx, base, target) + if err != nil { + return nil, err + } - baseByName := indexSubtypes(base.Subtypes) - targetByName := indexSubtypes(target.Subtypes) + baseByName, err := indexSubtypes(ctx, base.Subtypes) + if err != nil { + return nil, err + } + targetByName, err := indexSubtypes(ctx, target.Subtypes) + if err != nil { + return nil, err + } - allNames := mergeKeys(baseByName, targetByName) + allNames, err := mergeKeys(ctx, baseByName, targetByName) + if err != nil { + return nil, err + } + if err := ctx.Err(); err != nil { + return nil, err + } sort.Strings(allNames) + if err := ctx.Err(); err != nil { + return nil, err + } for _, name := range allNames { + if err := ctx.Err(); err != nil { + return nil, err + } baseSt, baseExists := baseByName[name] targetSt, targetExists := targetByName[name] prefix := string(base.Type) + "." + name if !baseExists { - changes = append(changes, addedSubtype(prefix, targetSt)...) + added, err := addedSubtype(ctx, prefix, targetSt) + if err != nil { + return nil, err + } + changes = append(changes, added...) continue } if !targetExists { - changes = append(changes, removedSubtype(prefix, baseSt)...) + removed, err := removedSubtype(ctx, prefix, baseSt) + if err != nil { + return nil, err + } + changes = append(changes, removed...) continue } - changes = append(changes, compareReadings(prefix, baseSt.Data, targetSt.Data)...) - changes = append(changes, compareStrings(prefix+".context", baseSt.Context, targetSt.Context)...) - changes = append(changes, compareItems(prefix, baseSt.Items, targetSt.Items)...) + readingChanges, err := compareReadings(ctx, prefix, baseSt.Data, targetSt.Data) + if err != nil { + return nil, err + } + changes = append(changes, readingChanges...) + stringChanges, err := compareStrings(ctx, prefix+".context", baseSt.Context, targetSt.Context) + if err != nil { + return nil, err + } + changes = append(changes, stringChanges...) + itemChanges, err := compareItems(ctx, prefix, baseSt.Items, targetSt.Items) + if err != nil { + return nil, err + } + changes = append(changes, itemChanges...) } - return changes + return changes, nil } -func compareReadings(prefix string, base, target map[string]measurement.Reading) []Change { +func compareReadings(ctx context.Context, prefix string, base, target map[string]measurement.Reading) ([]Change, error) { var changes []Change - allKeys := mergeKeys(base, target) + allKeys, err := mergeKeys(ctx, base, target) + if err != nil { + return nil, err + } + if err := ctx.Err(); err != nil { + return nil, err + } sort.Strings(allKeys) + if err := ctx.Err(); err != nil { + return nil, err + } for _, key := range allKeys { + if err := ctx.Err(); err != nil { + return nil, err + } path := dataPath(prefix, key) baseReading, baseExists := base[key] targetReading, targetExists := target[key] @@ -247,12 +377,15 @@ func compareReadings(prefix string, base, target map[string]measurement.Reading) baseVal := safeReadingString(baseReading) targetVal := safeReadingString(targetReading) + if err := ctx.Err(); err != nil { + return nil, err + } if baseVal != targetVal { changes = append(changes, Change{Kind: Modified, Severity: SeverityInfo, Path: path, Baseline: strPtr(baseVal), Target: strPtr(targetVal)}) } } - return changes + return changes, nil } func dataPath(prefix, key string) string { @@ -264,20 +397,32 @@ func dataPath(prefix, key string) string { return prefix + "." + key } -func addedReadings(prefix string, values map[string]measurement.Reading) []Change { - return compareReadings(prefix, nil, values) +func addedReadings(ctx context.Context, prefix string, values map[string]measurement.Reading) ([]Change, error) { + return compareReadings(ctx, prefix, nil, values) } -func removedReadings(prefix string, values map[string]measurement.Reading) []Change { - return compareReadings(prefix, values, nil) +func removedReadings(ctx context.Context, prefix string, values map[string]measurement.Reading) ([]Change, error) { + return compareReadings(ctx, prefix, values, nil) } -func compareStrings(prefix string, base, target map[string]string) []Change { +func compareStrings(ctx context.Context, prefix string, base, target map[string]string) ([]Change, error) { changes := make([]Change, 0) - keys := mergeKeys(base, target) + keys, err := mergeKeys(ctx, base, target) + if err != nil { + return nil, err + } + if err := ctx.Err(); err != nil { + return nil, err + } sort.Strings(keys) + if err := ctx.Err(); err != nil { + return nil, err + } for _, key := range keys { + if err := ctx.Err(); err != nil { + return nil, err + } path := prefix + "." + key baseValue, baseExists := base[key] targetValue, targetExists := target[key] @@ -292,15 +437,15 @@ func compareStrings(prefix string, base, target map[string]string) []Change { } } - return changes + return changes, nil } -func addedStrings(prefix string, values map[string]string) []Change { - return compareStrings(prefix, nil, values) +func addedStrings(ctx context.Context, prefix string, values map[string]string) ([]Change, error) { + return compareStrings(ctx, prefix, nil, values) } -func removedStrings(prefix string, values map[string]string) []Change { - return compareStrings(prefix, values, nil) +func removedStrings(ctx context.Context, prefix string, values map[string]string) ([]Change, error) { + return compareStrings(ctx, prefix, values, nil) } func itemPrefix(prefix string, index int) string { @@ -317,8 +462,11 @@ func lengthChange(prefix string, kind ChangeKind, baseline, target *string) Chan } } -func compareItems(prefix string, base, target []measurement.ItemEntry) []Change { +func compareItems(ctx context.Context, prefix string, base, target []measurement.ItemEntry) ([]Change, error) { changes := make([]Change, 0) + if err := ctx.Err(); err != nil { + return nil, err + } if len(base) != len(target) { changes = append(changes, lengthChange( prefix, @@ -330,107 +478,213 @@ func compareItems(prefix string, base, target []measurement.ItemEntry) []Change shared := min(len(base), len(target)) for i := 0; i < shared; i++ { + if err := ctx.Err(); err != nil { + return nil, err + } path := itemPrefix(prefix, i) - changes = append(changes, compareStrings(path+".context", base[i].Context, target[i].Context)...) - changes = append(changes, compareReadings(path+".data", base[i].Data, target[i].Data)...) + stringChanges, err := compareStrings(ctx, path+".context", base[i].Context, target[i].Context) + if err != nil { + return nil, err + } + changes = append(changes, stringChanges...) + readingChanges, err := compareReadings(ctx, path+".data", base[i].Data, target[i].Data) + if err != nil { + return nil, err + } + changes = append(changes, readingChanges...) } for i := shared; i < len(target); i++ { - changes = append(changes, addedItem(itemPrefix(prefix, i), &target[i])...) + if err := ctx.Err(); err != nil { + return nil, err + } + added, err := addedItem(ctx, itemPrefix(prefix, i), &target[i]) + if err != nil { + return nil, err + } + changes = append(changes, added...) } for i := shared; i < len(base); i++ { - changes = append(changes, removedItem(itemPrefix(prefix, i), &base[i])...) + if err := ctx.Err(); err != nil { + return nil, err + } + removed, err := removedItem(ctx, itemPrefix(prefix, i), &base[i]) + if err != nil { + return nil, err + } + changes = append(changes, removed...) } - return changes + return changes, nil } -func addedItem(prefix string, item *measurement.ItemEntry) []Change { - changes := addedStrings(prefix+".context", item.Context) - return append(changes, addedReadings(prefix+".data", item.Data)...) +func addedItem(ctx context.Context, prefix string, item *measurement.ItemEntry) ([]Change, error) { + changes, err := addedStrings(ctx, prefix+".context", item.Context) + if err != nil { + return nil, err + } + readings, err := addedReadings(ctx, prefix+".data", item.Data) + if err != nil { + return nil, err + } + return append(changes, readings...), nil } -func removedItem(prefix string, item *measurement.ItemEntry) []Change { - changes := removedStrings(prefix+".context", item.Context) - return append(changes, removedReadings(prefix+".data", item.Data)...) +func removedItem(ctx context.Context, prefix string, item *measurement.ItemEntry) ([]Change, error) { + changes, err := removedStrings(ctx, prefix+".context", item.Context) + if err != nil { + return nil, err + } + readings, err := removedReadings(ctx, prefix+".data", item.Data) + if err != nil { + return nil, err + } + return append(changes, readings...), nil } -func addedItems(prefix string, items []measurement.ItemEntry) []Change { +func addedItems(ctx context.Context, prefix string, items []measurement.ItemEntry) ([]Change, error) { if len(items) == 0 { - return nil + return nil, nil } changes := []Change{ lengthChange(prefix, Added, nil, strPtr(strconv.Itoa(len(items)))), } for i := range items { - changes = append(changes, addedItem(itemPrefix(prefix, i), &items[i])...) + if err := ctx.Err(); err != nil { + return nil, err + } + added, err := addedItem(ctx, itemPrefix(prefix, i), &items[i]) + if err != nil { + return nil, err + } + changes = append(changes, added...) } - return changes + return changes, nil } -func removedItems(prefix string, items []measurement.ItemEntry) []Change { +func removedItems(ctx context.Context, prefix string, items []measurement.ItemEntry) ([]Change, error) { if len(items) == 0 { - return nil + return nil, nil } changes := []Change{ lengthChange(prefix, Removed, strPtr(strconv.Itoa(len(items))), nil), } for i := range items { - changes = append(changes, removedItem(itemPrefix(prefix, i), &items[i])...) + if err := ctx.Err(); err != nil { + return nil, err + } + removed, err := removedItem(ctx, itemPrefix(prefix, i), &items[i]) + if err != nil { + return nil, err + } + changes = append(changes, removed...) } - return changes + return changes, nil } -func addedMeasurement(m *measurement.Measurement) []Change { +func addedMeasurement(ctx context.Context, m *measurement.Measurement) ([]Change, error) { changes := make([]Change, 0, len(m.Subtypes)) for i := range m.Subtypes { - changes = append(changes, addedSubtype(string(m.Type)+"."+m.Subtypes[i].Name, &m.Subtypes[i])...) + if err := ctx.Err(); err != nil { + return nil, err + } + added, err := addedSubtype(ctx, string(m.Type)+"."+m.Subtypes[i].Name, &m.Subtypes[i]) + if err != nil { + return nil, err + } + changes = append(changes, added...) } - return changes + return changes, nil } -func removedMeasurement(m *measurement.Measurement) []Change { +func removedMeasurement(ctx context.Context, m *measurement.Measurement) ([]Change, error) { changes := make([]Change, 0, len(m.Subtypes)) for i := range m.Subtypes { - changes = append(changes, removedSubtype(string(m.Type)+"."+m.Subtypes[i].Name, &m.Subtypes[i])...) + if err := ctx.Err(); err != nil { + return nil, err + } + removed, err := removedSubtype(ctx, string(m.Type)+"."+m.Subtypes[i].Name, &m.Subtypes[i]) + if err != nil { + return nil, err + } + changes = append(changes, removed...) } - return changes + return changes, nil } -func addedSubtype(prefix string, st *measurement.Subtype) []Change { - changes := addedReadings(prefix, st.Data) - changes = append(changes, addedStrings(prefix+".context", st.Context)...) - changes = append(changes, addedItems(prefix, st.Items)...) - return changes +func addedSubtype(ctx context.Context, prefix string, st *measurement.Subtype) ([]Change, error) { + changes, err := addedReadings(ctx, prefix, st.Data) + if err != nil { + return nil, err + } + stringChanges, err := addedStrings(ctx, prefix+".context", st.Context) + if err != nil { + return nil, err + } + changes = append(changes, stringChanges...) + items, err := addedItems(ctx, prefix, st.Items) + if err != nil { + return nil, err + } + return append(changes, items...), nil } -func removedSubtype(prefix string, st *measurement.Subtype) []Change { - changes := removedReadings(prefix, st.Data) - changes = append(changes, removedStrings(prefix+".context", st.Context)...) - changes = append(changes, removedItems(prefix, st.Items)...) - return changes +func removedSubtype(ctx context.Context, prefix string, st *measurement.Subtype) ([]Change, error) { + changes, err := removedReadings(ctx, prefix, st.Data) + if err != nil { + return nil, err + } + stringChanges, err := removedStrings(ctx, prefix+".context", st.Context) + if err != nil { + return nil, err + } + changes = append(changes, stringChanges...) + items, err := removedItems(ctx, prefix, st.Items) + if err != nil { + return nil, err + } + return append(changes, items...), nil } -func indexSubtypes(subtypes []measurement.Subtype) map[string]*measurement.Subtype { +func indexSubtypes(ctx context.Context, subtypes []measurement.Subtype) (map[string]*measurement.Subtype, error) { idx := make(map[string]*measurement.Subtype, len(subtypes)) for i := range subtypes { + if err := ctx.Err(); err != nil { + return nil, err + } idx[subtypes[i].Name] = &subtypes[i] } - return idx + return idx, nil } -func mergeKeys[V any](a, b map[string]V) []string { +func mergeKeys[V any](ctx context.Context, a, b map[string]V) ([]string, error) { seen := make(map[string]struct{}, len(a)+len(b)) for k := range a { + if err := ctx.Err(); err != nil { + return nil, err + } seen[k] = struct{}{} } for k := range b { + if err := ctx.Err(); err != nil { + return nil, err + } seen[k] = struct{}{} } keys := make([]string, 0, len(seen)) for k := range seen { + if err := ctx.Err(); err != nil { + return nil, err + } keys = append(keys, k) } - return keys + return keys, nil +} + +func snapshotContextError(cause error) error { + if stderrors.Is(cause, context.Canceled) { + return errors.Wrap(errors.ErrCodeCanceled, "snapshot diff canceled", cause) + } + return errors.Wrap(errors.ErrCodeTimeout, "snapshot diff deadline exceeded", cause) } diff --git a/pkg/diff/diff_test.go b/pkg/diff/diff_test.go index c07d89068..bad7a1b83 100644 --- a/pkg/diff/diff_test.go +++ b/pkg/diff/diff_test.go @@ -16,11 +16,14 @@ package diff import ( "bytes" + "context" + stderrors "errors" "fmt" "reflect" "strings" "testing" + aicrerrors "github.com/NVIDIA/aicr/pkg/errors" "github.com/NVIDIA/aicr/pkg/header" "github.com/NVIDIA/aicr/pkg/measurement" "github.com/NVIDIA/aicr/pkg/snapshotter" @@ -687,6 +690,130 @@ func TestSnapshots_EmptySnapshots(t *testing.T) { } } +func TestSnapshotsWithContext_MidTraversalCancellation(t *testing.T) { + baselineData := make(map[string]measurement.Reading, 64) + targetData := make(map[string]measurement.Reading, 64) + for i := 0; i < 64; i++ { + key := fmt.Sprintf("reading-%02d", i) + baselineData[key] = measurement.Int(i) + targetData[key] = measurement.Int(i + 1) + } + baseline := makeSnapshot(makeMeasurement(measurement.TypeK8s, makeSubtype("server", baselineData))) + target := makeSnapshot(makeMeasurement(measurement.TypeK8s, makeSubtype("server", targetData))) + probeCtx := &countingContext{Context: t.Context()} + probeResult, err := SnapshotsWithContext(probeCtx, baseline, target) + if err != nil { + t.Fatalf("SnapshotsWithContext() probe error = %v", err) + } + // The final summary traversal checks the context once per accumulated + // change. Cancel halfway through it without depending on the number of + // checkpoints used by earlier comparison stages. + cancelAt := probeCtx.checks - probeResult.Summary.Total/2 + if cancelAt <= 0 || cancelAt >= probeCtx.checks { + t.Fatalf("derived cancellation checkpoint = %d, probe checks = %d", cancelAt, probeCtx.checks) + } + + tests := []struct { + name string + cause error + wantCode aicrerrors.ErrorCode + }{ + {name: "canceled", cause: context.Canceled, wantCode: aicrerrors.ErrCodeCanceled}, + {name: "deadline", cause: context.DeadlineExceeded, wantCode: aicrerrors.ErrCodeTimeout}, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + ctx := newCheckpointContext(t.Context(), cancelAt, tt.cause) + result, err := SnapshotsWithContext(ctx, baseline, target) + if result != nil { + t.Fatalf("SnapshotsWithContext() result = %#v, want nil after cancellation", result) + } + if !stderrors.Is(err, aicrerrors.New(tt.wantCode, "")) { + t.Errorf("SnapshotsWithContext() error = %v, want code %s", err, tt.wantCode) + } + if !stderrors.Is(err, tt.cause) { + t.Errorf("SnapshotsWithContext() error = %v, want cause %v", err, tt.cause) + } + if ctx.checks < cancelAt { + t.Errorf("context checks = %d, want at least %d to prove traversal began", ctx.checks, cancelAt) + } + }) + } +} + +func TestSnapshotsWithContext_LegacyOutputUnchanged(t *testing.T) { + baseline := makeSnapshot(makeMeasurement(measurement.Type("Network"), makePFSubtype( + map[string]measurement.Reading{"mtu": measurement.Int(1500)}, + map[string]string{"node": "n1"}, + []measurement.ItemEntry{{ + Context: map[string]string{"name": "pf0"}, + Data: map[string]measurement.Reading{"speed": measurement.Int(100)}, + }}, + ))) + target := makeSnapshot(makeMeasurement(measurement.Type("Network"), makePFSubtype( + map[string]measurement.Reading{"mtu": measurement.Int(9000)}, + map[string]string{"node": "n2"}, + []measurement.ItemEntry{{ + Context: map[string]string{"name": "pf1"}, + Data: map[string]measurement.Reading{"speed": measurement.Int(200)}, + }}, + ))) + + legacy := Snapshots(baseline, target) + withContext, err := SnapshotsWithContext(t.Context(), baseline, target) + if err != nil { + t.Fatalf("SnapshotsWithContext() error = %v", err) + } + if !reflect.DeepEqual(withContext, legacy) { + t.Errorf("SnapshotsWithContext() = %#v, want legacy output %#v", withContext, legacy) + } +} + +type checkpointContext struct { + context.Context + cancelAt int + cause error + done chan struct{} + checks int + closed bool +} + +type countingContext struct { + context.Context + checks int +} + +func (c *countingContext) Err() error { + c.checks++ + return c.Context.Err() +} + +func newCheckpointContext(parent context.Context, cancelAt int, cause error) *checkpointContext { + return &checkpointContext{ + Context: parent, + cancelAt: cancelAt, + cause: cause, + done: make(chan struct{}), + } +} + +func (c *checkpointContext) Done() <-chan struct{} { + return c.done +} + +func (c *checkpointContext) Err() error { + c.checks++ + if c.checks < c.cancelAt { + return nil + } + if !c.closed { + close(c.done) + c.closed = true + } + return c.cause +} + // TestHasDrift_DerivedFromChanges verifies HasDrift derives from len(Changes) // rather than Summary.Total, so a caller-constructed Result whose Summary // hasn't been populated still reports drift correctly. Also verifies a nil diff --git a/pkg/diff/topology.go b/pkg/diff/topology.go index 906817b4a..99f61e646 100644 --- a/pkg/diff/topology.go +++ b/pkg/diff/topology.go @@ -15,6 +15,8 @@ package diff import ( + "context" + "github.com/NVIDIA/aicr/pkg/collector/topology" "github.com/NVIDIA/aicr/pkg/measurement" ) @@ -45,22 +47,38 @@ const ( // The cost is that a mixed-vintage comparison sees only what the folded // encoding could express, which is the accuracy the older snapshot was // captured with. Once both sides carry items the comparison is exact. -func alignTopologyEncoding(base, target *measurement.Measurement) (*measurement.Measurement, *measurement.Measurement) { +func alignTopologyEncoding( + ctx context.Context, + base, target *measurement.Measurement, +) (*measurement.Measurement, *measurement.Measurement, error) { + + if err := ctx.Err(); err != nil { + return nil, nil, err + } if base == nil || target == nil { - return base, target + return base, target, nil } if base.Type != measurement.TypeNodeTopology || target.Type != measurement.TypeNodeTopology { - return base, target + return base, target, nil } - baseIdx := indexSubtypes(base.Subtypes) - targetIdx := indexSubtypes(target.Subtypes) + baseIdx, indexErr := indexSubtypes(ctx, base.Subtypes) + if indexErr != nil { + return nil, nil, indexErr + } + targetIdx, indexErr := indexSubtypes(ctx, target.Subtypes) + if indexErr != nil { + return nil, nil, indexErr + } plan := map[string]subtypePlan{} for subtype, countKey := range map[string]string{ topologyLabelSubtype: topologyLabelCountKey, topologyTaintSubtype: topologyTaintCountKey, } { + if err := ctx.Err(); err != nil { + return nil, nil, err + } b, t := baseIdx[subtype], targetIdx[subtype] if b == nil || t == nil { continue @@ -78,8 +96,14 @@ func alignTopologyEncoding(base, target *measurement.Measurement) (*measurement. // Exclude collision-ambiguous keys: Go map iteration decides the // winner, so an unchanged cluster can write different values on each // side across an upgrade/rollback. - _, ambiguous, err := topology.HydrateItems(itemSide(b, t)) - if err != nil { + if err := ctx.Err(); err != nil { + return nil, nil, err + } + _, ambiguous, hydrateErr := topology.HydrateItems(itemSide(b, t)) + if contextErr := ctx.Err(); contextErr != nil { + return nil, nil, contextErr + } + if hydrateErr != nil { plan[subtype] = subtypePlan{dropItems: true, countKey: countKey} continue } @@ -89,10 +113,18 @@ func alignTopologyEncoding(base, target *measurement.Measurement) (*measurement. // the count here would mask a corrupted one rather than report it. } if len(plan) == 0 { - return base, target + return base, target, nil } - return alignMeasurement(base, plan), alignMeasurement(target, plan) + alignedBase, alignErr := alignMeasurement(ctx, base, plan) + if alignErr != nil { + return nil, nil, alignErr + } + alignedTarget, alignErr := alignMeasurement(ctx, target, plan) + if alignErr != nil { + return nil, nil, alignErr + } + return alignedBase, alignedTarget, nil } // subtypePlan is how one subtype is reduced before comparison. @@ -113,24 +145,41 @@ func itemSide(a, b *measurement.Subtype) *measurement.Subtype { } // alignMeasurement copies m with the topology subtypes reduced as directed. -func alignMeasurement(m *measurement.Measurement, plan map[string]subtypePlan) *measurement.Measurement { +func alignMeasurement( + ctx context.Context, + m *measurement.Measurement, + plan map[string]subtypePlan, +) (*measurement.Measurement, error) { + + if err := ctx.Err(); err != nil { + return nil, err + } out := *m out.Subtypes = make([]measurement.Subtype, len(m.Subtypes)) copy(out.Subtypes, m.Subtypes) folded := map[string]int{} for i := range out.Subtypes { + if err := ctx.Err(); err != nil { + return nil, err + } st := &out.Subtypes[i] p, ok := plan[st.Name] if !ok { continue } if p.hydrate { + if err := ctx.Err(); err != nil { + return nil, err + } if items, _, err := topology.HydrateItems(st); err == nil { st.Items = items } else { p.dropData = false // keep data when hydration fails } + if err := ctx.Err(); err != nil { + return nil, err + } } if p.dropData { st.Data = nil @@ -144,6 +193,9 @@ func alignMeasurement(m *measurement.Measurement, plan map[string]subtypePlan) * if len(p.skipKeys) > 0 && st.Data != nil { data := make(map[string]measurement.Reading, len(st.Data)) for k, v := range st.Data { + if err := ctx.Err(); err != nil { + return nil, err + } if !p.skipKeys[k] { data[k] = v } @@ -152,24 +204,33 @@ func alignMeasurement(m *measurement.Measurement, plan map[string]subtypePlan) * } } if len(folded) == 0 { - return &out + return &out, nil } for i := range out.Subtypes { + if err := ctx.Err(); err != nil { + return nil, err + } st := &out.Subtypes[i] if st.Name != topologySummarySubtype || st.Data == nil { continue } data := make(map[string]measurement.Reading, len(st.Data)) for k, v := range st.Data { + if err := ctx.Err(); err != nil { + return nil, err + } data[k] = v } for key, count := range folded { + if err := ctx.Err(); err != nil { + return nil, err + } if _, ok := data[key]; ok { data[key] = measurement.Int(count) } } st.Data = data } - return &out + return &out, nil }