Fix hub subscriber race, SOURCE_CHAN_SIZE handling, skipped-block serving, retry count and Range.Equals

forkable/blocks_from_num_test.go

@@ -0,0 +1,66 @@
+package forkable
+
+import (
+	"testing"
+
+	"github.com/streamingfast/bstream"
+	pbbstream "github.com/streamingfast/bstream/pb/sf/bstream/v1"
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+)
+
+// Some chains (Solana, NEAR, ...) can skip block numbers: requesting a start
+// block that was skipped must serve from the first block above it instead of
+// failing forever.
+func TestCallWithBlocksFromNum_SkippedBlockNums(t *testing.T) {
+	p := New(
+		bstream.HandlerFunc(func(blk *pbbstream.Block, obj any) error { return nil }),
+		HoldBlocksUntilLIB(),
+		WithKeptFinalBlocks(100),
+	)
+
+	blocks := []*pbbstream.Block{
+		bstream.TestBlockWithLIBNum("00000003", "00000002", 2),
+		bstream.TestBlockWithLIBNum("00000004", "00000003", 3),
+		bstream.TestBlockWithLIBNum("00000006", "00000004", 3), // number 5 skipped by the chain
+		bstream.TestBlockWithLIBNum("00000007", "00000006", 3),
+	}
+	for _, blk := range blocks {
+		require.NoError(t, p.ProcessBlock(blk, nil))
+	}
+
+	blockNumsFrom := func(num uint64) (out []uint64, err error) {
+		err = p.CallWithBlocksFromNum(num, func(blks []*bstream.PreprocessedBlock) {
+			for _, b := range blks {
+				out = append(out, b.Num())
+			}
+		}, false)
+		return
+	}
+
+	t.Run("requesting a skipped block num starts at the next available block", func(t *testing.T) {
+		nums, err := blockNumsFrom(5)
+		require.NoError(t, err)
+		assert.Equal(t, []uint64{6, 7}, nums)
+	})
+
+	t.Run("requesting an existing block num still starts exactly there", func(t *testing.T) {
+		nums, err := blockNumsFrom(4)
+		require.NoError(t, err)
+		assert.Equal(t, []uint64{4, 6, 7}, nums)
+
+		nums, err = blockNumsFrom(3)
+		require.NoError(t, err)
+		assert.Equal(t, []uint64{3, 4, 6, 7}, nums)
+	})
+
+	t.Run("requesting below the lowest held block fails so caller can use another source", func(t *testing.T) {
+		_, err := blockNumsFrom(2)
+		require.Error(t, err)
+	})
+
+	t.Run("requesting above head fails", func(t *testing.T) {
+		_, err := blockNumsFrom(8)
+		require.Error(t, err)
+	})
+}

forkable/forkable.go

@@ -180,6 +180,10 @@ func (p *Forkable) blocksFromNumWithForks(startNum uint64) ([]*bstream.Preproces
 	return out, nil
 }
 
+// blocksFromNum returns the blocks of the current canonical segment starting
+// at the first block whose number is >= num. The '>=' matters: chains like
+// Solana or NEAR can skip block numbers, so requiring an exact match would
+// make the whole segment unservable when the requested number was skipped.
 func (p *Forkable) blocksFromNum(num uint64) ([]*bstream.PreprocessedBlock, error) {
 	if !p.forkDB.HasLIB() {
 		return nil, fmt.Errorf("no lib")
@@ -196,13 +200,19 @@ func (p *Forkable) blocksFromNum(num uint64) ([]*bstream.PreprocessedBlock, erro
 		return nil, fmt.Errorf("head segment does not reach LIB")
 	}
 
+	if len(seg) > 0 && seg[0].AsRef().Num() > num {
+		// we don't hold the requested block: serving from the lowest block we
+		// have would silently skip blocks, let the caller use another source
+		return nil, fmt.Errorf("lowest block in complete segment %d is above requested block num %d", seg[0].AsRef().Num(), num)
+	}
+
 	libNum := p.forkDB.libRef.Num()
 
 	var out []*bstream.PreprocessedBlock
 	var seenBlock bool
 	for i := range seg {
 		ref := seg[i].AsRef()
-		if !seenBlock && ref.Num() == num {
+		if !seenBlock && ref.Num() >= num {
 			seenBlock = true
 		}
 

hub/hub.go

@@ -22,6 +22,7 @@ import (
 	"os"
 	"strconv"
 	"strings"
+	"sync"
 	"time"
 
 	"github.com/streamingfast/dstore"
@@ -45,6 +46,7 @@ type ForkableHub struct {
 
 	keepFinalBlocks int
 
+	subscribersLock   sync.Mutex // leaf lock: never acquire any other lock while holding it
 	subscribers       []*Subscription
 	sourceChannelSize int
 
@@ -68,14 +70,7 @@ func NewForkableHubWithOptions(liveSourceFactory bstream.SourceFactory, keepFina
 }
 
 func newForkableHub(liveSourceFactory bstream.SourceFactory, keepFinalBlocks int, oneBlocksStore dstore.Store, hubOptions []Option, extraForkableOptions ...forkable.Option) *ForkableHub {
-	sourceChanSize := 100
-	if os.Getenv("SOURCE_CHAN_SIZE") != "" {
-		newSize, err := strconv.Atoi(os.Getenv("SOURCE_CHAN_SIZE"))
-		if err != nil {
-			zlog.Warn("invalid SOURCE_CHAN_SIZE, ignoring", zap.Error(err))
-		}
-		sourceChanSize = newSize
-	}
+	sourceChanSize := sourceChanSizeFromEnv(100)
 
 	hub := &ForkableHub{
 		Shutter:           shutter.New(),
@@ -108,14 +103,35 @@ func newForkableHub(liveSourceFactory bstream.SourceFactory, keepFinalBlocks int
 	)
 
 	hub.OnTerminating(func(err error) {
-		for _, sub := range hub.subscribers {
+		hub.subscribersLock.Lock()
+		subscribers := make([]*Subscription, len(hub.subscribers))
+		copy(subscribers, hub.subscribers)
+		hub.subscribersLock.Unlock()
+
+		for _, sub := range subscribers {
 			sub.Shutdown(err)
 		}
 	})
 
 	return hub
 }
 
+// sourceChanSizeFromEnv returns the subscription channel size from the
+// SOURCE_CHAN_SIZE environment variable, falling back to defaultSize when the
+// variable is unset or invalid.
+func sourceChanSizeFromEnv(defaultSize int) int {
+	value := os.Getenv("SOURCE_CHAN_SIZE")
+	if value == "" {
+		return defaultSize
+	}
+	newSize, err := strconv.Atoi(value)
+	if err != nil {
+		zlog.Warn("invalid SOURCE_CHAN_SIZE, ignoring", zap.Error(err))
+		return defaultSize
+	}
+	return newSize
+}
+
 // Option configures a ForkableHub.
 type Option func(h *ForkableHub)
 
@@ -200,19 +216,26 @@ func (h *ForkableHub) IsReady() bool {
 	}
 }
 
-// subscribe must be called while hub is locked
+// subscribe must be called while the forkable is locked (via forkable.CallWithBlocks*)
+// so that no new block can be broadcast between the snapshot of initialBlocks and the
+// registration of the subscription. The subscribers slice itself is protected by
+// subscribersLock: the forkable lock is not enough because CallWithBlocks* only takes
+// a read lock, so multiple subscribe calls can run concurrently.
 func (h *ForkableHub) subscribe(handler bstream.Handler, initialBlocks []*bstream.PreprocessedBlock, withPartials bool) *Subscription {
 	chanSize := h.sourceChannelSize + len(initialBlocks)
 	sub := NewSubscription(handler, chanSize, withPartials)
 	for _, ppblk := range initialBlocks {
 		_ = sub.push(ppblk)
 	}
+	h.subscribersLock.Lock()
 	h.subscribers = append(h.subscribers, sub)
+	h.subscribersLock.Unlock()
 	return sub
 }
 
-// unsubscribe must be called while hub is locked
 func (h *ForkableHub) unsubscribe(removeSub *Subscription) {
+	h.subscribersLock.Lock()
+	defer h.subscribersLock.Unlock()
 	var newSubscriber []*Subscription
 	for _, sub := range h.subscribers {
 		if sub != removeSub {
@@ -518,7 +541,10 @@ func (h *ForkableHub) broadcastBlock(blk *pbbstream.Block, obj any) error {
 
 	preprocBlock := &bstream.PreprocessedBlock{Block: blk, Obj: obj}
 
-	subscribers := h.subscribers // we may remove some from the original slice during the loop
+	h.subscribersLock.Lock()
+	subscribers := make([]*Subscription, len(h.subscribers))
+	copy(subscribers, h.subscribers) // we may remove some from the original slice during the loop
+	h.subscribersLock.Unlock()
 
 	for _, sub := range subscribers {
 		err := sub.push(preprocBlock)

hub/hub_subscribe_race_test.go

@@ -0,0 +1,107 @@
+package hub
+
+import (
+	"fmt"
+	"sync"
+	"testing"
+	"time"
+
+	"github.com/streamingfast/bstream"
+	"github.com/streamingfast/bstream/forkable"
+	pbbstream "github.com/streamingfast/bstream/pb/sf/bstream/v1"
+	"github.com/streamingfast/shutter"
+	"github.com/stretchr/testify/require"
+)
+
+// TestForkableHub_ConcurrentSubscribe exercises concurrent subscriptions while
+// blocks are being broadcast. All SourceFrom* paths call subscribe under the
+// forkable *read* lock only, so two concurrent subscribers used to race on the
+// h.subscribers slice (caught by -race) and could lose a subscription entirely.
+func TestForkableHub_ConcurrentSubscribe(t *testing.T) {
+	fh := &ForkableHub{
+		Shutter:           shutter.New(),
+		logger:            zlog,
+		sourceChannelSize: 100,
+	}
+	fh.forkable = forkable.New(bstream.HandlerFunc(fh.broadcastBlock),
+		forkable.HoldBlocksUntilLIB(),
+		forkable.WithKeptFinalBlocks(100),
+		forkable.WithFilters(bstream.StepsAllWithPartial),
+	)
+
+	// Seed the forkable so that SourceFromBlockNum can serve block 4.
+	require.NoError(t, fh.forkable.ProcessBlock(bstream.TestBlockWithLIBNum("00000003", "00000002", 2), nil))
+	require.NoError(t, fh.forkable.ProcessBlock(bstream.TestBlockWithLIBNum("00000004", "00000003", 3), nil))
+
+	const subscriberCount = 20
+	const lastBroadcastBlock = uint64(40)
+
+	start := make(chan struct{})
+
+	// Broadcaster: pushes blocks 5..40 through the forkable while subscribers register.
+	broadcastDone := make(chan struct{})
+	go func() {
+		defer close(broadcastDone)
+		<-start
+		prev := "00000004"
+		for num := uint64(5); num <= lastBroadcastBlock; num++ {
+			id := fmt.Sprintf("%08x", num)
+			if err := fh.forkable.ProcessBlock(bstream.TestBlockWithLIBNum(id, prev, 3), nil); err != nil {
+				panic(err)
+			}
+			prev = id
+		}
+	}()
+
+	type subscriber struct {
+		source   bstream.Source
+		received chan uint64
+	}
+	subs := make([]*subscriber, subscriberCount)
+
+	var wg sync.WaitGroup
+	for i := 0; i < subscriberCount; i++ {
+		sub := &subscriber{received: make(chan uint64, 256)}
+		subs[i] = sub
+		wg.Add(1)
+		go func() {
+			defer wg.Done()
+			<-start
+			handler := bstream.HandlerFunc(func(blk *pbbstream.Block, _ any) error {
+				sub.received <- blk.Number
+				return nil
+			})
+			sub.source = fh.SourceFromBlockNum(4, handler)
+		}()
+	}
+
+	close(start)
+	wg.Wait()
+	<-broadcastDone
+
+	// Every subscription must exist and receive blocks up to the last one
+	// broadcast after it registered. A lost subscription (dropped by a racing
+	// append) would never receive it.
+	for i, sub := range subs {
+		require.NotNil(t, sub.source, "subscriber %d did not get a source", i)
+		go sub.source.Run()
+
+		var last uint64
+		timeout := time.After(5 * time.Second)
+	drain:
+		for {
+			select {
+			case num := <-sub.received:
+				if num > last {
+					last = num
+				}
+				if last == lastBroadcastBlock {
+					break drain
+				}
+			case <-timeout:
+				t.Fatalf("subscriber %d timed out waiting for block %d, last received %d (lost subscription?)", i, lastBroadcastBlock, last)
+			}
+		}
+		sub.source.Shutdown(nil)
+	}
+}

hub/source_chan_size_test.go

@@ -0,0 +1,41 @@
+package hub
+
+import (
+	"testing"
+
+	"github.com/streamingfast/bstream"
+	"github.com/stretchr/testify/require"
+)
+
+func TestSourceChanSizeFromEnv(t *testing.T) {
+	tests := []struct {
+		name     string
+		envValue string
+		setEnv   bool
+		expect   int
+	}{
+		{"unset keeps default", "", false, 100},
+		{"empty keeps default", "", true, 100},
+		{"garbage keeps default", "garbage", true, 100},
+		{"valid value is used", "42", true, 42},
+	}
+
+	for _, test := range tests {
+		t.Run(test.name, func(t *testing.T) {
+			if test.setEnv {
+				t.Setenv("SOURCE_CHAN_SIZE", test.envValue)
+			}
+			require.Equal(t, test.expect, sourceChanSizeFromEnv(100))
+		})
+	}
+}
+
+func TestNewForkableHub_InvalidSourceChanSizeEnv(t *testing.T) {
+	t.Setenv("SOURCE_CHAN_SIZE", "not-a-number")
+
+	lsf := bstream.NewTestSourceFactory()
+	fh := NewForkableHub(lsf.NewSource, 0, nil)
+
+	require.Equal(t, 100, fh.sourceChannelSize,
+		"invalid SOURCE_CHAN_SIZE must keep the default channel size instead of using 0 and disconnecting every subscriber")
+}

range.go

@@ -343,10 +343,15 @@ func (r *Range) IsNext(next *Range, size uint64) bool {
 }
 
 func (r *Range) Equals(other *Range) bool {
-	return r.startBlock == other.startBlock &&
-		r.endBlock == other.endBlock &&
-		r.exclusiveStartBlock == other.exclusiveStartBlock &&
-		r.exclusiveEndBlock == other.exclusiveEndBlock
+	if r.startBlock != other.startBlock ||
+		r.exclusiveStartBlock != other.exclusiveStartBlock ||
+		r.exclusiveEndBlock != other.exclusiveEndBlock {
+		return false
+	}
+	if r.endBlock == nil || other.endBlock == nil {
+		return r.endBlock == nil && other.endBlock == nil
+	}
+	return *r.endBlock == *other.endBlock
 }
 
 func (r *Range) Size() (uint64, error) {