feat(crit): add utility functions for stats

This adds the commonly used functions to fetch dump and restore
statistics for a checkpoint.

Signed-off-by: Prajwal S N <[email protected]>

Author: snprajwal

crit/stats.go

@@ -0,0 +1,46 @@
+package crit
+
+import (
+	"errors"
+	"path/filepath"
+
+	"github.com/checkpoint-restore/go-criu/v6/crit/images"
+)
+
+// Helper function to load stats file into Go struct
+func getStats(path string) (*images.StatsEntry, error) {
+	c := New(path, "", "", false, false)
+	statsImg, err := c.Decode()
+	if err != nil {
+		return nil, err
+	}
+
+	stats, ok := statsImg.Entries[0].Message.(*images.StatsEntry)
+	if !ok {
+		return nil, errors.New("Failed to type assert stats image")
+	}
+
+	return stats, nil
+}
+
+// GetDumpStats returns the dump statistics of a checkpoint.
+// dir is the path to the directory with the checkpoint images.
+func GetDumpStats(dir string) (*images.DumpStatsEntry, error) {
+	stats, err := getStats(filepath.Join(dir, "stats-dump"))
+	if err != nil {
+		return nil, err
+	}
+
+	return stats.GetDump(), nil
+}
+
+// GetRestoreStats returns the restore statistics of a checkpoint.
+// dir is the path to the directory with the checkpoint images.
+func GetRestoreStats(dir string) (*images.RestoreStatsEntry, error) {
+	stats, err := getStats(filepath.Join(dir, "stats-restore"))
+	if err != nil {
+		return nil, err
+	}
+
+	return stats.GetRestore(), nil
+}

test/crit/Makefile

@@ -2,7 +2,10 @@ CC ?= gcc
 GO ?= go
 CRIU ?= criu
 
-all: integration-test e2e-test clean
+all: unit-test integration-test e2e-test clean
+
+unit-test: test-imgs
+	go test -v ./...
 
 integration-test: test-imgs crit-test
 	@echo "Running integration test"
@@ -16,6 +19,8 @@ test-imgs: ../loop/loop
 	$(eval PID := $(shell ../loop/loop))
 	mkdir -p $@
 	$(CRIU) dump -v4 -o dump.log -D $@ -t $(PID)
+	$(CRIU) restore -v4 -o restore.log -D $@ -d
+	pkill -9 loop
 
 ../../crit/bin/crit:
 	$(MAKE) -C ../../crit bin/crit
@@ -29,4 +34,4 @@ crit-test: main.go
 clean:
 	@rm -rf test-imgs
 
-.PHONY: all test integration-test e2e-test clean
+.PHONY: all test unit-test integration-test e2e-test clean

test/crit/stats_test.go

@@ -0,0 +1,27 @@
+package main
+
+import (
+	"testing"
+
+	"github.com/checkpoint-restore/go-criu/v6/crit"
+)
+
+func TestGetDumpStats(t *testing.T) {
+	dumpStats, err := crit.GetDumpStats("test-imgs")
+	if err != nil {
+		t.Error("Failed to get stats")
+	}
+	if dumpStats.GetPagesWritten() == 0 {
+		t.Error("PagesWritten is 0")
+	}
+}
+
+func TestGetRestoreStats(t *testing.T) {
+	restoreStats, err := crit.GetRestoreStats("test-imgs")
+	if err != nil {
+		t.Error("Failed to get stats")
+	}
+	if restoreStats.GetForkingTime() == 0 {
+		t.Error("ForkingTime is 0")
+	}
+}