Merge pull request #17 from pyroscope-io/feat/symbol_preprocessor

Add cpool preprocessor parse option
grafana · Jun 12, 2022 · 3073e81 · 3073e81
2 parents 3228ecf + 5267e8a
commit 3073e81
Show file tree

Hide file tree

Showing 4 changed files with 22 additions and 8 deletions.
diff --git a/parser/checkpoint.go b/parser/checkpoint.go
@@ -41,7 +41,7 @@ func (c *CheckpointEvent) Parse(r reader.Reader, classes ClassMap, cpools PoolMa
 		}
 		cm, ok := cpools[int(classID)]
 		if !ok {
-			cpools[int(classID)] = &CPool{pool: make(map[int]ParseResolvable)}
+			cpools[int(classID)] = &CPool{Pool: make(map[int]ParseResolvable)}
 			cm = cpools[int(classID)]
 		}
 		m, err := r.VarInt()
@@ -58,7 +58,7 @@ func (c *CheckpointEvent) Parse(r reader.Reader, classes ClassMap, cpools PoolMa
 			if err != nil {
 				return fmt.Errorf("unable to parse constant type %d: %w", classID, err)
 			}
-			cm.pool[int(idx)] = v
+			cm.Pool[int(idx)] = v
 		}
 	}
 	return nil

diff --git a/parser/chunk.go b/parser/chunk.go
@@ -11,7 +11,7 @@ import (
 var magic = []byte{'F', 'L', 'R', 0}
 
 type CPool struct {
-	pool     map[int]ParseResolvable
+	Pool     map[int]ParseResolvable
 	resolved bool
 }
 type ClassMap map[int]ClassMetadata
@@ -24,7 +24,11 @@ type Chunk struct {
 	Events      []Parseable
 }
 
-func (c *Chunk) Parse(r io.Reader) (err error) {
+type ChunkParseOptions struct {
+	CPoolProcessor func(meta ClassMetadata, cpool *CPool)
+}
+
+func (c *Chunk) Parse(r io.Reader, options *ChunkParseOptions) (err error) {
 	buf := make([]byte, len(magic))
 	if _, err = io.ReadFull(r, buf); err != nil {
 		if err == io.EOF {
@@ -103,6 +107,12 @@ func (c *Chunk) Parse(r io.Reader) (err error) {
 		br.Seek(c.Header.ConstantPoolOffset+delta, io.SeekStart)
 	}
 
+	if options.CPoolProcessor != nil {
+		for classID, pool := range cpools {
+			options.CPoolProcessor(classes[classID], pool)
+		}
+	}
+
 	// Second pass over constant pools: resolve constants
 	for classID := range cpools {
 		if err := ResolveConstants(classes, cpools, classID); err != nil {
@@ -153,7 +163,7 @@ func ResolveConstants(classes ClassMap, cpools PoolMap, classID int) (err error)
 		return nil
 	}
 	cpool.resolved = true
-	for _, t := range cpool.pool {
+	for _, t := range cpool.Pool {
 		if err := t.Resolve(classes, cpools); err != nil {
 			return fmt.Errorf("unable to resolve constants: %w", err)
 		}

diff --git a/parser/parser.go b/parser/parser.go
@@ -6,10 +6,14 @@ import (
 )
 
 func Parse(r io.Reader) ([]Chunk, error) {
+	return ParseWithOptions(r, &ChunkParseOptions{})
+}
+
+func ParseWithOptions(r io.Reader, options *ChunkParseOptions) ([]Chunk, error) {
 	var chunks []Chunk
 	for {
 		var chunk Chunk
-		err := chunk.Parse(r)
+		err := chunk.Parse(r, options)
 		if err == io.EOF {
 			return chunks, nil
 		}

diff --git a/parser/types.go b/parser/types.go
@@ -97,7 +97,7 @@ func parseFields(r reader.Reader, classes ClassMap, cpools PoolMap, class ClassM
 				if err != nil {
 					return fmt.Errorf("unable to read constant index")
 				}
-				p, ok := cpool.pool[int(i)]
+				p, ok := cpool.Pool[int(i)]
 				if !ok {
 					continue
 				}
@@ -147,7 +147,7 @@ func resolveConstants(classes ClassMap, cpools PoolMap, constants *[]constant, r
 			// Non-existent constant pool references seem to be used to mark no value
 			continue
 		}
-		it, ok := p.pool[int(c.index)]
+		it, ok := p.Pool[int(c.index)]
 		if !ok {
 			// Non-existent constant pool references seem to be used to mark no value
 			continue