Merge pull request ipld/go-car#37 from ipld/feat/two-step-user-funcs

Allow user defined block hooks when using two step write for selective cars

This commit was moved from ipld/go-car@8fe0b74

Author: hannahhoward

ipld/car/car_test.go

@@ -117,7 +117,9 @@ func TestRoundtripSelective(t *testing.T) {
 	// write car in one step
 	buf := new(bytes.Buffer)
 	blockCount := 0
+	var oneStepBlocks []Block
 	err := sc.Write(buf, func(block Block) error {
+		oneStepBlocks = append(oneStepBlocks, block)
 		blockCount++
 		return nil
 	})
@@ -128,7 +130,11 @@ func TestRoundtripSelective(t *testing.T) {
 	sc2 := NewSelectiveCar(context.Background(), sourceBs, []Dag{Dag{Root: nd3.Cid(), Selector: selector}})
 
 	// write car in two steps
-	scp, err := sc2.Prepare()
+	var twoStepBlocks []Block
+	scp, err := sc2.Prepare(func(block Block) error {
+		twoStepBlocks = append(twoStepBlocks, block)
+		return nil
+	})
 	require.NoError(t, err)
 	buf2 := new(bytes.Buffer)
 	err = scp.Dump(buf2)
@@ -141,6 +147,9 @@ func TestRoundtripSelective(t *testing.T) {
 	// verify equal data written by both methods
 	require.Equal(t, buf.Bytes(), buf2.Bytes())
 
+	// verify equal blocks were passed to user block hook funcs
+	require.Equal(t, oneStepBlocks, twoStepBlocks)
+
 	// readout car and verify contents
 	bserv := dstest.Bserv()
 	ch, err := LoadCar(bserv.Blockstore(), buf)

ipld/car/selectivecar.go

@@ -50,9 +50,10 @@ type OnNewCarBlockFunc func(Block) error
 // the Car file like size and number of blocks that go into it
 type SelectiveCarPrepared struct {
 	SelectiveCar
-	size   uint64
-	header CarHeader
-	cids   []cid.Cid
+	size               uint64
+	header             CarHeader
+	cids               []cid.Cid
+	userOnNewCarBlocks []OnNewCarBlockFunc
 }
 
 // NewSelectiveCar creates a new SelectiveCar for the given car file based
@@ -72,7 +73,7 @@ func (sc SelectiveCar) traverse(onCarHeader OnCarHeaderFunc, onNewCarBlock OnNew
 
 // Prepare traverse a car file and collects data on what is about to be written, but
 // does not actually write the file
-func (sc SelectiveCar) Prepare() (SelectiveCarPrepared, error) {
+func (sc SelectiveCar) Prepare(userOnNewCarBlocks ...OnNewCarBlockFunc) (SelectiveCarPrepared, error) {
 	var header CarHeader
 	var cids []cid.Cid
 
@@ -88,7 +89,7 @@ func (sc SelectiveCar) Prepare() (SelectiveCarPrepared, error) {
 	if err != nil {
 		return SelectiveCarPrepared{}, err
 	}
-	return SelectiveCarPrepared{sc, size, header, cids}, nil
+	return SelectiveCarPrepared{sc, size, header, cids, userOnNewCarBlocks}, nil
 }
 
 func (sc SelectiveCar) Write(w io.Writer, userOnNewCarBlocks ...OnNewCarBlockFunc) error {
@@ -133,6 +134,10 @@ func (sc SelectiveCarPrepared) Cids() []cid.Cid {
 // Dump writes the car file as quickly as possible based on information already
 // collected
 func (sc SelectiveCarPrepared) Dump(w io.Writer) error {
+	offset, err := HeaderSize(&sc.header)
+	if err != nil {
+		return fmt.Errorf("failed to size car header: %s", err)
+	}
 	if err := WriteHeader(&sc.header, w); err != nil {
 		return fmt.Errorf("failed to write car header: %s", err)
 	}
@@ -142,10 +147,23 @@ func (sc SelectiveCarPrepared) Dump(w io.Writer) error {
 			return err
 		}
 		raw := blk.RawData()
+		size := util.LdSize(c.Bytes(), raw)
 		err = util.LdWrite(w, c.Bytes(), raw)
 		if err != nil {
 			return err
 		}
+		for _, userOnNewCarBlock := range sc.userOnNewCarBlocks {
+			err := userOnNewCarBlock(Block{
+				BlockCID: c,
+				Data:     raw,
+				Offset:   offset,
+				Size:     size,
+			})
+			if err != nil {
+				return err
+			}
+		}
+		offset += size
 	}
 	return nil
 }