Merge pull request #11 from parca-dev/add_python_initial_state

Add initial state

.github/workflows/test-integration.yml

@@ -36,11 +36,8 @@ jobs:
         with:
           direnvVersion: 2.32.3
 
-      - name: Download Ruby Runtimes
-        run: ./scripts/download/ruby.sh
-
-      - name: Download Python Runtimes
-        run: ./scripts/download/python.sh
+      - name: Build
+        run: make build
 
       - name: Set up cache for test runtimes
         uses: actions/[email protected]

Makefile

@@ -3,38 +3,37 @@ bootstrap:
 	curl -fsSL https://get.jetpack.io/devbox | bash
 	curl -sfL https://direnv.net/install.sh | bash
 
-.PHONY: dirs
-dirs:
-	mkdir -p pkg/ruby/versions
-	mkdir -p pkg/python/versions
+GO_SRC := $(shell find pkg -type f -name '*.go')
 
-structlayout: cmd/structlayout/structlayout.go
+structlayout: cmd/structlayout/structlayout.go $(filter-out *_test.go,$(GO_SRC))
 	go build -o $@ $<
 
-mergelayout: cmd/mergelayout/mergelayout.go
+mergelayout: cmd/mergelayout/mergelayout.go $(filter-out *_test.go,$(GO_SRC))
 	go build -o $@ $<
 
 .PHONY: build
-build: dirs structlayout mergelayout
+build: structlayout mergelayout
 	go build ./...
 
+.PHONY: generate
+generate: build generate/python generate/ruby
+
+.PHONY: generate/python
+generate/python:
+	./scripts/download/python.sh
+	./scripts/structlayout/python.sh
+	./scripts/mergelayout/python.sh
+
+.PHONY: generate/ruby
+generate/ruby:
+	./scripts/download/ruby.sh
+	./scripts/structlayout/ruby.sh
+	./scripts/mergelayout/ruby.sh
 
 .PHONY: clean
 clean:
 	rm -rf target
 
-.PHONY: test
-test: test
-	go test ./...
-
-.PHONY: test/integration
-test/integration:
-	go test -tags=integration ./tests/integration/...
-
-.PHONY: test/integration/update
-test/integration/update:
-	go test -count=1 -race -tags=integration ./tests/integration/... -update
-
 .PHONY: check
 check: vet lint
 
@@ -58,6 +57,18 @@ format:
 tagalign:
 	go run github.com/4meepo/tagalign/cmd/tagalign@latest -fix -sort ./...
 
+.PHONY: test
+test: generate
+	go test ./...
+
+.PHONY: test/integration
+test/integration: generate
+	go test -tags=integration ./tests/integration/...
+
+.PHONY: test/integration/update
+test/integration/update:
+	go test -count=1 -race -tags=integration ./tests/integration/... -update
+
 TMPDIR := ./tmp
 $(TMPDIR):
 	mkdir -p $(TMPDIR)

README.md

@@ -17,12 +17,12 @@ import (
 )
 
 func main() {
-    versions, err := python.GetVersions()
+    layouts, err := python.GetLayouts()
     if err != nil {
-        return fmt.Errorf("get python versions: %w", err)
+        return fmt.Errorf("get python layouts: %w", err)
     }
 
-    fmt.Println(versions)
+    fmt.Println(layouts)
 }
 ```
 
@@ -38,12 +38,12 @@ import (
 )
 
 func main() {
-    versions, err := ruby.GetVersions()
+    layouts, err := ruby.GetLayouts()
     if err != nil {
-        return fmt.Errorf("get ruby versions: %w", err)
+        return fmt.Errorf("get ruby layouts: %w", err)
     }
 
-    fmt.Println(versions)
+    fmt.Println(layouts)
 }
 ```
 

cmd/mergelayout/mergelayout.go

@@ -79,22 +79,17 @@ func main() {
 	logger.Info("done", "output directory", outputDir)
 }
 
-type layoutWithVersion struct {
-	Version runtimedata.Version `yaml:"version"`
-	Layout  map[string]any      `yaml:"layout"`
-}
-
 func mergeLayoutFiles(logger *slog.Logger, inputFiles []string, output string) error {
 	// Read all the input files and store them in a map with the version as the key.
-	versionedLayouts := map[runtimedata.Version]*layoutWithVersion{}
+	versionedLayouts := map[runtimedata.Version]*runtimedata.DataWithVersion{}
 	for _, file := range inputFiles {
 		logger.Info("reading file", "file", file)
 		data, err := os.ReadFile(file)
 		if err != nil {
 			return err
 		}
 
-		var withVersion layoutWithVersion
+		var withVersion runtimedata.DataWithVersion
 		if err := yaml.Unmarshal(data, &withVersion); err != nil {
 			return err
 		}
@@ -130,23 +125,23 @@ func mergeLayoutFiles(logger *slog.Logger, inputFiles []string, output string) e
 	)
 	for _, v := range versionKeys {
 		currentVersion := convertVersion(v)
-		layout := versionedLayouts[v].Layout
+		data := versionedLayouts[v].Data
 		if minVersion == nil {
 			minVersion = currentVersion
 			maxVersion = currentVersion
-			currentLayout = layout
+			currentLayout = data
 			continue
 		}
 
-		if reflect.DeepEqual(currentLayout, layout) {
+		if reflect.DeepEqual(currentLayout, data) {
 			maxVersion = currentVersion
 			continue
 		}
 
 		addVersionRange()
 		minVersion = currentVersion
 		maxVersion = currentVersion
-		currentLayout = layout
+		currentLayout = data
 	}
 
 	// Add the last version range if it's not there already.

cmd/structlayout/structlayout.go

@@ -1,12 +1,14 @@
 package main
 
 import (
+	"debug/dwarf"
 	"debug/elf"
 	"flag"
 	"fmt"
 	"log/slog"
 	"os"
 	"path/filepath"
+	"reflect"
 	"strings"
 
 	"gopkg.in/yaml.v3"
@@ -53,25 +55,28 @@ func main() {
 	}
 
 	var (
-		layoutMap runtimedata.LayoutMap
-		outputDir = givenOutputDir
+		layoutMap       runtimedata.LayoutMap
+		initialStateMap any
+		outputDir       = givenOutputDir
 	)
 	switch runtime {
 	case "python":
-		layoutMap = python.DataMapForVersion(version)
+		layoutMap = python.DataMapForLayout(version)
+		initialStateMap = python.DataMapForInitialState(version)
 		if outputDir == "" {
-			outputDir = "pkg/python/versions"
+			// Base output directory for python is pkg/python.
+			outputDir = "pkg/python"
 		}
 	case "ruby":
-		layoutMap = ruby.DataMapForVersion(version)
+		layoutMap = ruby.DataMapForLayout(version)
 		if outputDir == "" {
-			outputDir = "pkg/ruby/versions"
+			outputDir = "pkg/ruby"
 		}
 	case "libc":
 		// TODO(kakkoyun): Change depending on the libc implementation. e.g musl, glibc, etc.
-		// layoutMap = libc.DataMapForVersion(version)
+		// layoutMap = libc.DataMapForLayout(version)
 		// if outputDir == "" {
-		// 	outputDir = "pkg/libc/versions"
+		// 	outputDir = "pkg/libc/layout"
 		// }
 	default:
 		logger.Error("invalid offset map module", "mod", runtime)
@@ -84,31 +89,74 @@ func main() {
 	}
 
 	var (
-		input  = fSet.Arg(0)
-		output = filepath.Join(outputDir, fmt.Sprintf("%s_%s.yaml", runtime, sanitizeIdentifier(version)))
+		input = fSet.Arg(0)
 	)
-	if err := processAndWriteLayout(input, output, version, layoutMap); err != nil {
+	dwarfData, err := dwarfDataFromELF(input)
+	if err != nil {
+		logger.Error("failed to read DWARF data", "err", err)
+		os.Exit(1)
+	}
+
+	output := filepath.Join(outputDir, "layout", fmt.Sprintf("%s_%s.yaml", runtime, sanitizeIdentifier(version)))
+	if err := processAndWriteLayout(dwarfData, output, version, layoutMap); err != nil {
 		logger.Error("failed to write layout", "err", err)
 		os.Exit(1)
 	}
+	logger.Info("layout file written", "file", output)
 
-	logger.Info("layout written to file", "file", output)
+	if isNil(initialStateMap) {
+		logger.Info("no initial state map found, skipping initial state generation")
+		os.Exit(0)
+	}
+
+	output = filepath.Join(outputDir, "initialstate", fmt.Sprintf("%s_%s.yaml", runtime, sanitizeIdentifier(version)))
+	if err := processAndWriteInitialState(dwarfData, output, version, initialStateMap); err != nil {
+		logger.Error("failed to write initial state", "err", err)
+		os.Exit(1)
+	}
+	logger.Info("initial state file written", "file", output)
 }
 
 // processAndWriteLayout processes the given ELF file and writes the layout to the given output file.
-func processAndWriteLayout(input, output string, version string, layoutMap runtimedata.LayoutMap) error {
-	ef, err := elf.Open(input)
+func processAndWriteLayout(dwarfData *dwarf.Data, output string, version string, layoutMap runtimedata.LayoutMap) error {
+	dm, err := datamap.New(layoutMap)
 	if err != nil {
-		return fmt.Errorf("failed to open ELF file: %w", err)
+		return fmt.Errorf("failed to create data map: %w", err)
 	}
-	defer ef.Close()
 
-	dwarfData, err := ef.DWARF()
+	if err := dm.ReadFromDWARF(dwarfData); err != nil {
+		return fmt.Errorf("failed to extract struct layout from DWARF data: %w", err)
+	}
+
+	if err := os.MkdirAll(filepath.Dir(output), 0755); err != nil {
+		return fmt.Errorf("failed to create output directory: %w", err)
+	}
+
+	file, err := os.Create(output)
 	if err != nil {
-		return fmt.Errorf("failed to read DWARF info: %w", err)
+		return fmt.Errorf("failed to create output file: %w", err)
 	}
 
-	dm, err := datamap.New(layoutMap)
+	// Extremely in-efficient and hacky but it should work for now.
+	withVersion, err := runtimedata.WithVersion(version, convertToMapOfAny(layoutMap.Layout()))
+	if err != nil {
+		return fmt.Errorf("failed to wrap layout with version: %w", err)
+	}
+
+	encoder := yaml.NewEncoder(file)
+	if err := encoder.Encode(withVersion); err != nil {
+		return fmt.Errorf("failed to encode layout: %w", err)
+	}
+	if err := encoder.Close(); err != nil {
+		return fmt.Errorf("failed to close encoder: %w", err)
+	}
+
+	return nil
+}
+
+// processAndWriteInitialState processes the given ELF file and writes the initial state to the given output file.
+func processAndWriteInitialState(dwarfData *dwarf.Data, output string, version string, initialStateMap any) error {
+	dm, err := datamap.New(initialStateMap)
 	if err != nil {
 		return fmt.Errorf("failed to create data map: %w", err)
 	}
@@ -117,12 +165,17 @@ func processAndWriteLayout(input, output string, version string, layoutMap runti
 		return fmt.Errorf("failed to extract struct layout from DWARF data: %w", err)
 	}
 
+	if err := os.MkdirAll(filepath.Dir(output), 0755); err != nil {
+		return fmt.Errorf("failed to create output directory: %w", err)
+	}
+
 	file, err := os.Create(output)
 	if err != nil {
 		return fmt.Errorf("failed to create output file: %w", err)
 	}
 
-	withVersion, err := runtimedata.WithVersion(version, layoutMap.Layout())
+	// Extremely in-efficient and hacky but it should work for now.
+	withVersion, err := runtimedata.WithVersion(version, convertToMapOfAny(initialStateMap))
 	if err != nil {
 		return fmt.Errorf("failed to wrap layout with version: %w", err)
 	}
@@ -138,7 +191,51 @@ func processAndWriteLayout(input, output string, version string, layoutMap runti
 	return nil
 }
 
+// dwarfDataFromELF returns the DWARF data from the given ELF file.
+func dwarfDataFromELF(input string) (*dwarf.Data, error) {
+	ef, err := elf.Open(input)
+	if err != nil {
+		return nil, fmt.Errorf("failed to open ELF file: %w", err)
+	}
+	defer ef.Close()
+
+	dwarfData, err := ef.DWARF()
+	if err != nil {
+		return nil, fmt.Errorf("failed to read DWARF info: %w", err)
+	}
+
+	return dwarfData, nil
+}
+
 // sanitizeIdentifier sanitizes the identifier to be used as a filename.
 func sanitizeIdentifier(identifier string) string {
 	return strings.TrimPrefix(strings.ReplaceAll(identifier, ".", "_"), "v")
 }
+
+// convertToMapOfAny converts the given struct to a map of string to any.
+func convertToMapOfAny(v interface{}) map[string]any {
+	// Marshal and unmarshal to convert the struct to map[string]any.
+	blob, err := yaml.Marshal(v)
+	if err != nil {
+		panic(err)
+	}
+
+	var anyMap map[string]any
+	if err := yaml.Unmarshal(blob, &anyMap); err != nil {
+		panic(err)
+	}
+
+	return anyMap
+}
+
+func isNil(v any) bool {
+	if v == nil {
+		return true
+	}
+	val := reflect.ValueOf(v)
+	switch val.Kind() {
+	case reflect.Chan, reflect.Func, reflect.Map, reflect.Pointer, reflect.UnsafePointer, reflect.Interface, reflect.Slice:
+		return val.IsNil()
+	}
+	return false
+}

pkg/python/datamap.go

@@ -9,7 +9,7 @@ import (
 
 const doesNotExist = -1
 
-func DataMapForVersion(v string) runtimedata.LayoutMap {
+func DataMapForLayout(v string) runtimedata.LayoutMap {
 	// Keys are version constraints defined in semver format,
 	// check github.com/Masterminds/semver for more details.
 	pythonVersions := map[*semver.Constraints]runtimedata.LayoutMap{