Fixed early termination of bootstrap fetching

snow/engine/avalanche/bootstrapper.go

@@ -119,7 +119,7 @@ func (b *bootstrapper) fetch(vtxID ids.ID) {
 		b.sendRequest(vtxID)
 		return
 	}
-	b.addVertex(vtx)
+	b.storeVertex(vtx)
 }
 
 func (b *bootstrapper) sendRequest(vtxID ids.ID) {
@@ -138,6 +138,14 @@ func (b *bootstrapper) sendRequest(vtxID ids.ID) {
 }
 
 func (b *bootstrapper) addVertex(vtx avalanche.Vertex) {
+	b.storeVertex(vtx)
+
+	if numPending := b.pending.Len(); numPending == 0 {
+		b.finish()
+	}
+}
+
+func (b *bootstrapper) storeVertex(vtx avalanche.Vertex) {
 	vts := []avalanche.Vertex{vtx}
 
 	for len(vts) > 0 {
@@ -181,9 +189,6 @@ func (b *bootstrapper) addVertex(vtx avalanche.Vertex) {
 
 	numPending := b.pending.Len()
 	b.numPendingRequests.Set(float64(numPending))
-	if numPending == 0 {
-		b.finish()
-	}
 }
 
 func (b *bootstrapper) finish() {

snow/engine/avalanche/bootstrapper_test.go

@@ -957,3 +957,53 @@ func TestBootstrapperFilterAccepted(t *testing.T) {
 		t.Fatalf("Vtx shouldn't be accepted")
 	}
 }
+
+func TestBootstrapperPartialFetch(t *testing.T) {
+	config, _, sender, state, _ := newConfig(t)
+
+	vtxID0 := ids.Empty.Prefix(0)
+	vtxID1 := ids.Empty.Prefix(1)
+
+	vtxBytes0 := []byte{0}
+
+	vtx0 := &Vtx{
+		id:     vtxID0,
+		height: 0,
+		status: choices.Processing,
+		bytes:  vtxBytes0,
+	}
+
+	bs := bootstrapper{}
+	bs.metrics.Initialize(config.Context.Log, fmt.Sprintf("gecko_%s", config.Context.ChainID), prometheus.NewRegistry())
+	bs.Initialize(config)
+
+	acceptedIDs := ids.Set{}
+	acceptedIDs.Add(
+		vtxID0,
+		vtxID1,
+	)
+
+	state.getVertex = func(vtxID ids.ID) (avalanche.Vertex, error) {
+		switch {
+		case vtxID.Equals(vtxID0):
+			return vtx0, nil
+		case vtxID.Equals(vtxID1):
+			return nil, errUnknownVertex
+		default:
+			t.Fatal(errUnknownVertex)
+			panic(errUnknownVertex)
+		}
+	}
+
+	sender.CantGet = false
+
+	bs.ForceAccepted(acceptedIDs)
+
+	if bs.finished {
+		t.Fatalf("should have requested a vertex")
+	}
+
+	if bs.pending.Len() != 1 {
+		t.Fatalf("wrong number pending")
+	}
+}

snow/engine/snowman/bootstrapper.go

@@ -113,7 +113,7 @@ func (b *bootstrapper) fetch(blkID ids.ID) {
 		b.sendRequest(blkID)
 		return
 	}
-	b.addBlock(blk)
+	b.storeBlock(blk)
 }
 
 func (b *bootstrapper) sendRequest(blkID ids.ID) {
@@ -132,6 +132,14 @@ func (b *bootstrapper) sendRequest(blkID ids.ID) {
 }
 
 func (b *bootstrapper) addBlock(blk snowman.Block) {
+	b.storeBlock(blk)
+
+	if numPending := b.pending.Len(); numPending == 0 {
+		b.finish()
+	}
+}
+
+func (b *bootstrapper) storeBlock(blk snowman.Block) {
 	status := blk.Status()
 	blkID := blk.ID()
 	for status == choices.Processing {
@@ -161,9 +169,6 @@ func (b *bootstrapper) addBlock(blk snowman.Block) {
 
 	numPending := b.pending.Len()
 	b.numPendingRequests.Set(float64(numPending))
-	if numPending == 0 {
-		b.finish()
-	}
 }
 
 func (b *bootstrapper) finish() {

snow/engine/snowman/bootstrapper_test.go

@@ -425,3 +425,54 @@ func TestBootstrapperFilterAccepted(t *testing.T) {
 		t.Fatalf("Blk shouldn't be accepted")
 	}
 }
+
+func TestBootstrapperPartialFetch(t *testing.T) {
+	config, _, sender, vm := newConfig(t)
+
+	blkID0 := ids.Empty.Prefix(0)
+	blkID1 := ids.Empty.Prefix(1)
+
+	blkBytes0 := []byte{0}
+
+	blk0 := &Blk{
+		id:     blkID0,
+		height: 0,
+		status: choices.Accepted,
+		bytes:  blkBytes0,
+	}
+
+	bs := bootstrapper{}
+	bs.metrics.Initialize(config.Context.Log, fmt.Sprintf("gecko_%s", config.Context.ChainID), prometheus.NewRegistry())
+	bs.Initialize(config)
+
+	acceptedIDs := ids.Set{}
+	acceptedIDs.Add(
+		blkID0,
+		blkID1,
+	)
+
+	vm.GetBlockF = func(blkID ids.ID) (snowman.Block, error) {
+		switch {
+		case blkID.Equals(blkID0):
+			return blk0, nil
+		case blkID.Equals(blkID1):
+			return nil, errUnknownBlock
+		default:
+			t.Fatal(errUnknownBlock)
+			panic(errUnknownBlock)
+		}
+	}
+
+	sender.CantGet = false
+	bs.onFinished = func() {}
+
+	bs.ForceAccepted(acceptedIDs)
+
+	if bs.finished {
+		t.Fatalf("should have requested a block")
+	}
+
+	if bs.pending.Len() != 1 {
+		t.Fatalf("wrong number pending")
+	}
+}