Fix fork compression preemption metric

Author: sjmiller609

lib/instances/fork.go

@@ -245,7 +245,7 @@ func (m *manager) forkInstanceFromStoppedOrStandby(ctx context.Context, id strin
 	defer cu.Clean()
 
 	if source.State == StateStandby {
-		if err := m.ensureSnapshotMemoryReady(ctx, m.paths.InstanceSnapshotLatest(id), m.snapshotJobKeyForInstance(id), stored.HypervisorType); err != nil {
+		if err := m.ensureSnapshotMemoryReady(ctx, m.paths.InstanceSnapshotLatest(id), m.snapshotJobKeyForInstance(id), stored.HypervisorType, snapshotCompressionPreemptionForkInstance); err != nil {
 			return nil, fmt.Errorf("prepare standby snapshot for fork: %w", err)
 		}
 	}

lib/instances/metrics.go

@@ -39,6 +39,7 @@ type snapshotCompressionPreemptionOperation string
 const (
 	snapshotCompressionPreemptionRestoreInstance snapshotCompressionPreemptionOperation = "restore_instance"
 	snapshotCompressionPreemptionRestoreSnapshot snapshotCompressionPreemptionOperation = "restore_snapshot"
+	snapshotCompressionPreemptionForkInstance    snapshotCompressionPreemptionOperation = "fork_instance"
 	snapshotCompressionPreemptionForkSnapshot    snapshotCompressionPreemptionOperation = "fork_snapshot"
 	snapshotCompressionPreemptionCreateSnapshot  snapshotCompressionPreemptionOperation = "create_snapshot"
 	snapshotCompressionPreemptionDeleteInstance  snapshotCompressionPreemptionOperation = "delete_instance"

lib/instances/metrics_test.go

@@ -1,6 +1,7 @@
 package instances
 
 import (
+	"os"
 	"testing"
 	"time"
 
@@ -113,6 +114,59 @@ func TestSnapshotCompressionMetrics_RecordAndObserve(t *testing.T) {
 	assert.Equal(t, "standby", metricLabel(t, active.DataPoints[0].Attributes, "source"))
 }
 
+func TestEnsureSnapshotMemoryReadyRecordsProvidedPreemptionOperation(t *testing.T) {
+	t.Parallel()
+
+	reader := otelmetric.NewManualReader()
+	provider := otelmetric.NewMeterProvider(otelmetric.WithReader(reader))
+
+	rawDir := t.TempDir()
+	rawPath := rawDir + "/memory-ranges"
+	require.NoError(t, os.WriteFile(rawPath, []byte("raw snapshot"), 0644))
+
+	jobDone := make(chan struct{})
+	m := &manager{
+		paths: paths.New(t.TempDir()),
+		compressionJobs: map[string]*compressionJob{
+			"job-1": {
+				done: jobDone,
+				target: compressionTarget{
+					Key:            "job-1",
+					HypervisorType: hypervisor.TypeCloudHypervisor,
+					Source:         snapshotCompressionSourceStandby,
+					Policy: snapshotstore.SnapshotCompressionConfig{
+						Enabled:   true,
+						Algorithm: snapshotstore.SnapshotCompressionAlgorithmLz4,
+					},
+				},
+			},
+		},
+	}
+	m.compressionJobs["job-1"].cancel = func() {
+		select {
+		case <-jobDone:
+		default:
+			close(jobDone)
+		}
+	}
+
+	metrics, err := newInstanceMetrics(provider.Meter("test"), nil, m)
+	require.NoError(t, err)
+	m.metrics = metrics
+
+	err = m.ensureSnapshotMemoryReady(t.Context(), rawDir, "job-1", hypervisor.TypeCloudHypervisor, snapshotCompressionPreemptionForkInstance)
+	require.NoError(t, err)
+
+	var rm metricdata.ResourceMetrics
+	require.NoError(t, reader.Collect(t.Context(), &rm))
+
+	preemptionsMetric := findMetric(t, rm, "hypeman_snapshot_compression_preemptions_total")
+	preemptions, ok := preemptionsMetric.Data.(metricdata.Sum[int64])
+	require.True(t, ok)
+	require.Len(t, preemptions.DataPoints, 1)
+	assert.Equal(t, "fork_instance", metricLabel(t, preemptions.DataPoints[0].Attributes, "operation"))
+}
+
 func assertMetricNames(t *testing.T, rm metricdata.ResourceMetrics, expected []string) {
 	t.Helper()
 

lib/instances/restore.go

@@ -73,7 +73,7 @@ func (m *manager) restoreInstance(
 	if m.metrics != nil && m.metrics.tracer != nil {
 		ctx, prepareSnapshotSpan = m.metrics.tracer.Start(ctx, "PrepareSnapshotMemory")
 	}
-	err = m.ensureSnapshotMemoryReady(ctx, snapshotDir, m.snapshotJobKeyForInstance(id), stored.HypervisorType)
+	err = m.ensureSnapshotMemoryReady(ctx, snapshotDir, m.snapshotJobKeyForInstance(id), stored.HypervisorType, snapshotCompressionPreemptionRestoreInstance)
 	if prepareSnapshotSpan != nil {
 		prepareSnapshotSpan.End()
 	}

lib/instances/snapshot.go

@@ -103,7 +103,7 @@ func (m *manager) createSnapshot(ctx context.Context, id string, req CreateSnaps
 			if target != nil {
 				m.recordSnapshotCompressionPreemption(ctx, snapshotCompressionPreemptionCreateSnapshot, *target)
 			}
-			if err := m.ensureSnapshotMemoryReady(ctx, m.paths.InstanceSnapshotLatest(id), "", stored.HypervisorType); err != nil {
+			if err := m.ensureSnapshotMemoryReady(ctx, m.paths.InstanceSnapshotLatest(id), "", stored.HypervisorType, snapshotCompressionPreemptionCreateSnapshot); err != nil {
 				return nil, fmt.Errorf("prepare source snapshot memory for copy: %w", err)
 			}
 		default:
@@ -393,7 +393,7 @@ func (m *manager) forkSnapshot(ctx context.Context, snapshotID string, req ForkS
 	if target != nil {
 		m.recordSnapshotCompressionPreemption(ctx, snapshotCompressionPreemptionForkSnapshot, *target)
 	}
-	if err := m.ensureSnapshotMemoryReady(ctx, m.paths.SnapshotGuestDir(snapshotID), "", rec.StoredMetadata.HypervisorType); err != nil {
+	if err := m.ensureSnapshotMemoryReady(ctx, m.paths.SnapshotGuestDir(snapshotID), "", rec.StoredMetadata.HypervisorType, snapshotCompressionPreemptionForkSnapshot); err != nil {
 		return nil, fmt.Errorf("prepare snapshot memory for fork: %w", err)
 	}
 

lib/instances/snapshot_compression.go

@@ -417,7 +417,7 @@ func (m *manager) cancelAndWaitCompressionJob(ctx context.Context, key string) (
 	}
 }
 
-func (m *manager) ensureSnapshotMemoryReady(ctx context.Context, snapshotDir, jobKey string, hvType hypervisor.Type) error {
+func (m *manager) ensureSnapshotMemoryReady(ctx context.Context, snapshotDir, jobKey string, hvType hypervisor.Type, preemptionOp snapshotCompressionPreemptionOperation) error {
 	start := time.Now()
 
 	if jobKey != "" {
@@ -426,7 +426,7 @@ func (m *manager) ensureSnapshotMemoryReady(ctx context.Context, snapshotDir, jo
 			return err
 		}
 		if target != nil {
-			m.recordSnapshotCompressionPreemption(ctx, snapshotCompressionPreemptionRestoreInstance, *target)
+			m.recordSnapshotCompressionPreemption(ctx, preemptionOp, *target)
 		}
 	}