builder.go

@@ -168,6 +168,7 @@ func (sdb *builder) WriteStateDiff(
 
 // WriteStateDiff writes a statediff object to output sinks
 func (sdb *builder) WriteStateSnapshot(
+	ctx context.Context,
 	stateRoot common.Hash, params Params,
 	nodeSink sdtypes.StateNodeSink,
 	ipldSink sdtypes.IPLDSink,
@@ -200,7 +201,7 @@ func (sdb *builder) WriteStateSnapshot(
 		}
 	}
 	// errgroup will cancel if any group fails
-	g, ctx := errgroup.WithContext(context.Background())
+	g, ctx := errgroup.WithContext(ctx)
 	for i := range subiters {
 		func(subdiv uint) {
 			g.Go(func() error {

test_helpers/builder.go

@@ -2,6 +2,7 @@ package test_helpers
 
 import (
 	"bytes"
+	"context"
 	"fmt"
 	"math/big"
 	"math/rand"
@@ -93,6 +94,7 @@ func RunStateSnapshot(
 				tr := tracker.New(recoveryFile, subtries)
 				defer tr.CloseAndSave()
 				return builder.WriteStateSnapshot(
+					context.Background(),
 					test.StateRoot, params, stateAppender, ipldAppender, tr,
 				)
 			}