Merge pull request #57 from adrianreber/2021-10-29-stats-functions

Move statistics read funtion to stats/

Author: rst0git

Makefile

@@ -40,7 +40,6 @@ phaul-test: $(TEST_BINARIES)
 clean:
 	@rm -f $(TEST_BINARIES)
 	@rm -rf image
-	@rm -f rpc/rpc.proto stats/stats.proto
 
 rpc/rpc.proto:
 	curl -sSL https://raw.githubusercontent.com/checkpoint-restore/criu/master/images/rpc.proto -o $@

magic/types.go

@@ -0,0 +1,12 @@
+package magic
+
+const (
+	ImgCommonMagic  = 0x54564319 /* Sarov (a.k.a. Arzamas-16) */
+	ImgServiceMagic = 0x55105940 /* Zlatoust */
+	StatsMagic      = 0x57093306 /* Ostashkov */
+
+	PrimaryMagicOffset   = 0x0
+	SecondaryMagicOffset = 0x4
+	SizeOffset           = 0x8
+	PayloadOffset        = 0xC
+)

phaul/client.go

@@ -110,7 +110,7 @@ func (pc *Client) Migrate() error {
 			return err
 		}
 
-		st, err := criuGetDumpStats(imgDir)
+		st, err := stats.CriuGetDumpStats(imgDir)
 		if err != nil {
 			return err
 		}

phaul/stats.go

@@ -0,0 +1,6 @@
+package stats
+
+const (
+	StatsDump    = "stats-dump"
+	StatsRestore = "stats-restore"
+)

stats/types.go

@@ -0,0 +1,54 @@
+package stats
+
+import (
+	"encoding/binary"
+	"errors"
+	"io/ioutil"
+	"os"
+	"path/filepath"
+
+	"github.com/checkpoint-restore/go-criu/v5/magic"
+	"google.golang.org/protobuf/proto"
+)
+
+func readStatisticsFile(imgDir *os.File, fileName string) (*StatsEntry, error) {
+	buf, err := ioutil.ReadFile(filepath.Join(imgDir.Name(), fileName))
+	if err != nil {
+		return nil, err
+	}
+
+	if binary.LittleEndian.Uint32(buf[magic.PrimaryMagicOffset:magic.SecondaryMagicOffset]) != magic.ImgServiceMagic {
+		return nil, errors.New("Primary magic not found")
+	}
+
+	if binary.LittleEndian.Uint32(buf[magic.SecondaryMagicOffset:magic.SizeOffset]) != magic.StatsMagic {
+		return nil, errors.New("Secondary magic not found")
+	}
+
+	payloadSize := binary.LittleEndian.Uint32(buf[magic.SizeOffset:magic.PayloadOffset])
+
+	st := &StatsEntry{}
+	if err := proto.Unmarshal(buf[magic.PayloadOffset:magic.PayloadOffset+payloadSize], st); err != nil {
+		return nil, err
+	}
+
+	return st, nil
+}
+
+func CriuGetDumpStats(imgDir *os.File) (*DumpStatsEntry, error) {
+	st, err := readStatisticsFile(imgDir, StatsDump)
+	if err != nil {
+		return nil, err
+	}
+
+	return st.GetDump(), nil
+}
+
+func CriuGetRestoreStats(imgDir *os.File) (*RestoreStatsEntry, error) {
+	st, err := readStatisticsFile(imgDir, StatsRestore)
+	if err != nil {
+		return nil, err
+	}
+
+	return st.GetRestore(), nil
+}