Recover parked nodes from mem-cached layers with a small tree

prover/prover.go

@@ -2,7 +2,9 @@ package prover
 
 import (
 	"context"
+	"errors"
 	"fmt"
+	"io"
 	"os"
 	"path/filepath"
 	"regexp"
@@ -144,10 +146,15 @@ func makeRecoveryProofTree(
 		return nil, nil, fmt.Errorf("layer 0 cache file is missing")
 	}
 
+	var topLayer uint
 	parkedNodesMap := make(map[uint][]byte)
 
 	// Validate structure.
 	for layer, file := range layersFiles {
+		if layer > topLayer {
+			topLayer = layer
+		}
+
 		readWriter, err := layerFactory(uint(layer))
 		if err != nil {
 			return nil, nil, err
@@ -211,31 +218,16 @@ func makeRecoveryProofTree(
 			parkedNodes = append(parkedNodes, nil)
 		}
 	}
-
-	// recover parked nodes from mem-cached layers
-	fileLayers := len(parkedNodes)
-	memLayerFactory := NewReadWriterMetaFactory(
-		treeCfg.MinMemoryLayer,
-		treeCfg.Datadir,
-		treeCfg.FileWriterBufSize,
-	).GetFactory()
-	for layer := fileLayers; ; layer++ {
-		numNodes := nextLeafID >> layer
-		if numNodes == 0 {
-			break
-		}
-		if numNodes%2 == 0 {
-			parkedNodes = append(parkedNodes, nil)
-		} else {
-			node, err := GetNode(memLayerFactory, layer, numNodes-1, merkleHashFunc)
-			if err != nil {
-				return nil, nil, fmt.Errorf("recovering parked node in layer %d: %w", layer, err)
-			}
-			logging.FromContext(ctx).
-				Info("recovered (mem) parked node", zap.Uint("layer", uint(layer)), zap.String("node", fmt.Sprintf("%X", node)))
-			parkedNodes = append(parkedNodes, node)
-		}
+	layerReader, err := layerFactory(topLayer)
+	if err != nil {
+		return nil, nil, err
 	}
+	defer layerReader.Close()
+	memCachedParkedNodes, err := recoverMemCachedParkedNodes(layerReader, merkleHashFunc)
+	if err != nil {
+		return nil, nil, fmt.Errorf("Recoveing parked nodes from top layer of disk-cache: %w", err)
+	}
+	parkedNodes = append(parkedNodes, memCachedParkedNodes...)
 
 	logging.FromContext(ctx).
 		Info("all recovered parked nodes", zap.Array("nodes", zapcore.ArrayMarshalerFunc(func(enc zapcore.ArrayEncoder) error {
@@ -403,33 +395,33 @@ func CalcTreeRoot(leaves [][]byte) ([]byte, error) {
 	return tree.Root(), nil
 }
 
-// GetNode returns the node at the given index in the given layer.
-// It will recursively calculate the node if it is not already cached.
-func GetNode(layerFactory mshared.LayerFactory, layer int, index uint64, hash merkle.HashFunc) ([]byte, error) {
-	readWriter, err := layerFactory(uint(layer))
+// build a small tree with the nodes from the top layer of the cache as leafs.
+// this tree will be used to get parked nodes for the merkle tree.
+func recoverMemCachedParkedNodes(layerReader mshared.LayerReader, merkleHashFunc merkle.HashFunc) ([][]byte, error) {
+	tmpDir, err := os.MkdirTemp(os.TempDir(), "poet-recovery-tree")
 	if err != nil {
 		return nil, err
 	}
-	defer readWriter.Close()
-
-	// Try to obtain from cache.
-	if err := readWriter.Seek(index); err == nil {
-		return readWriter.ReadNext()
-	}
+	defer os.RemoveAll(tmpDir)
 
-	if layer == 0 {
-		return nil, fmt.Errorf("cannot recreate leaf at index %d", index)
-	}
-
-	left, err := GetNode(layerFactory, layer-1, index*2, hash)
+	tree, err := merkle.NewTreeBuilder().WithHashFunc(merkleHashFunc).Build()
 	if err != nil {
 		return nil, err
 	}
 
-	right, err := GetNode(layerFactory, layer-1, index*2+1, hash)
-	if err != nil {
-		return nil, err
+	// append nodes as leafs
+	for {
+		node, err := layerReader.ReadNext()
+		if errors.Is(err, io.EOF) {
+			break
+		}
+		if err != nil {
+			return nil, fmt.Errorf("reading node from top layer of disk-cache: %w", err)
+		}
+		if err := tree.AddLeaf(node); err != nil {
+			return nil, fmt.Errorf("adding node to small tree: %w", err)
+		}
 	}
-
-	return hash(nil, left, right), nil
+	// the first parked node is for the leaves from the layerReader.
+	return tree.GetParkedNodes(nil)[1:], nil
 }

service/round.go

@@ -227,7 +227,7 @@ func (r *round) persistExecution(
 	return r.saveState()
 }
 
-func (r *round) recoverExecution(ctx context.Context, end time.Time, minMemoryLayer, fileWriterBufSize uint) error {
+func (r *round) recoverExecution(ctx context.Context, end time.Time, fileWriterBufSize uint) error {
 	logger := logging.FromContext(ctx).With(zap.String("round", r.ID))
 
 	started := time.Now()
@@ -263,7 +263,6 @@ func (r *round) recoverExecution(ctx context.Context, end time.Time, minMemoryLa
 		prover.TreeConfig{
 			Datadir:           r.datadir,
 			FileWriterBufSize: fileWriterBufSize,
-			MinMemoryLayer:    minMemoryLayer,
 		},
 		hash.GenLabelHashFunc(r.execution.Statement),
 		hash.GenMerkleHashFunc(r.execution.Statement),

service/round_test.go

@@ -360,7 +360,7 @@ func TestRound_ExecutionRecovery(t *testing.T) {
 
 		ctx, stop := context.WithTimeout(context.Background(), 100*time.Millisecond)
 		defer stop()
-		req.ErrorIs(round.recoverExecution(ctx, time.Now().Add(time.Hour), 2, 0), context.DeadlineExceeded)
+		req.ErrorIs(round.recoverExecution(ctx, time.Now().Add(time.Hour), 0), context.DeadlineExceeded)
 		req.NoError(round.teardown(context.Background(), false))
 	}
 
@@ -371,7 +371,7 @@ func TestRound_ExecutionRecovery(t *testing.T) {
 		req.Equal(len(challenges), numChallenges(round))
 		req.NoError(round.loadState())
 
-		req.NoError(round.recoverExecution(context.Background(), time.Now().Add(400*time.Millisecond), 2, 0))
+		req.NoError(round.recoverExecution(context.Background(), time.Now().Add(400*time.Millisecond), 0))
 		validateProof(t, round.execution)
 		req.NoError(round.teardown(context.Background(), true))
 	}

service/service.go

@@ -224,7 +224,7 @@ func (s *Service) loop(ctx context.Context, roundToResume *round) error {
 		eg.Go(func() error {
 			unlock := lockOSThread(ctx, roundTidFile)
 			defer unlock()
-			err := round.recoverExecution(ctx, end, s.minMemoryLayer, s.cfg.TreeFileBufferSize)
+			err := round.recoverExecution(ctx, end, s.cfg.TreeFileBufferSize)
 			roundResults <- roundResult{round: round, err: err}
 			return nil
 		})