routewatch/internal/routewatch/batchwrite_test.go

package routewatch

import (
	"slices"
	"sync"
	"testing"
	"time"

	"git.eeqj.de/sneak/routewatch/internal/database"
	"git.eeqj.de/sneak/routewatch/internal/logger"
	"git.eeqj.de/sneak/routewatch/internal/ristypes"
)

// handleMessageTimeout is how long a test waits for HandleMessage to return,
// or for a write to start, before it fails.
const handleMessageTimeout = 2 * time.Second

// blockingStore is a mockStore whose batch writes for the AS, peer and prefix
// handlers wait until the test closes release. writeStarted is closed when the
// first of them starts. It records what each write was given.
type blockingStore struct {
	*mockStore

	release      chan struct{}
	writeStarted chan struct{}
	startOnce    sync.Once

	mu       sync.Mutex
	prefixes map[string]bool
	asns     map[int]bool
	peers    map[string]bool
}

func newBlockingStore() *blockingStore {
	return &blockingStore{
		mockStore:    newMockStore(),
		release:      make(chan struct{}),
		writeStarted: make(chan struct{}),
		prefixes:     make(map[string]bool),
		asns:         make(map[int]bool),
		peers:        make(map[string]bool),
	}
}

func (s *blockingStore) waitForRelease() {
	s.startOnce.Do(func() { close(s.writeStarted) })
	<-s.release
}

func (s *blockingStore) UpsertLiveRouteBatch(routes []*database.LiveRoute) error {
	s.waitForRelease()
	s.mu.Lock()
	defer s.mu.Unlock()
	for _, route := range routes {
		s.prefixes[route.Prefix] = true
	}

	return nil
}

func (s *blockingStore) GetOrCreateASNBatch(asns map[int]time.Time) error {
	s.waitForRelease()
	s.mu.Lock()
	defer s.mu.Unlock()
	for asn := range asns {
		s.asns[asn] = true
	}

	return nil
}

func (s *blockingStore) UpdatePeerBatch(peers map[string]database.PeerUpdate) error {
	s.waitForRelease()
	s.mu.Lock()
	defer s.mu.Unlock()
	for peerIP := range peers {
		s.peers[peerIP] = true
	}

	return nil
}

// returnsInTime reports whether f returns within handleMessageTimeout.
func returnsInTime(f func()) bool {
	done := make(chan struct{})
	go func() {
		f()
		close(done)
	}()

	select {
	case <-done:
		return true
	case <-time.After(handleMessageTimeout):
		return false
	}
}

// handleWhileAWriteIsBlocked hands h the messages in fill, which fill its
// batch so that the handler writes it to store, waits until the write starts,
// and hands h the message extra while store holds the write. HandleMessage
// must return each time. Then it lets the writes finish and stops h, which
// writes what it holds.
func handleWhileAWriteIsBlocked(
	t *testing.T,
	store *blockingStore,
	h interface {
		HandleMessage(*ristypes.RISMessage)
		Stop()
	},
	fill []*ristypes.RISMessage,
	extra *ristypes.RISMessage,
) {
	t.Helper()
	defer h.Stop()
	defer close(store.release)

	if !returnsInTime(func() {
		for _, msg := range fill {
			h.HandleMessage(msg)
		}
	}) {
		t.Fatal("HandleMessage did not return after filling the batch: it waits for the database write")
	}

	select {
	case <-store.writeStarted:
	case <-time.After(handleMessageTimeout):
		t.Fatal("the handler did not write its full batch")
	}

	if !returnsInTime(func() { h.HandleMessage(extra) }) {
		t.Fatal("HandleMessage did not return while a database write was blocked")
	}
}

func TestPrefixHandlerHandlesMessagesWhileAWriteIsBlocked(t *testing.T) {
	const (
		fillPrefix  = "192.0.2.0/24"
		extraPrefix = "198.51.100.0/24"
	)

	announce := func(prefixes []string) *ristypes.RISMessage {
		return &ristypes.RISMessage{
			Peer:          testPeerIP,
			Path:          ristypes.ASPath{testASNA},
			Announcements: []ristypes.RISAnnouncement{{NextHop: testPeerIP, Prefixes: prefixes}},
		}
	}

	store := newBlockingStore()
	h := NewPrefixHandler(store, logger.New())
	handleWhileAWriteIsBlocked(t, store, h,
		[]*ristypes.RISMessage{announce(slices.Repeat([]string{fillPrefix}, prefixBatchSize))},
		announce([]string{extraPrefix}))

	store.mu.Lock()
	defer store.mu.Unlock()
	if !store.prefixes[extraPrefix] {
		t.Errorf("Stop did not write %s, handled while a write was blocked", extraPrefix)
	}
}

func TestASHandlerHandlesMessagesWhileAWriteIsBlocked(t *testing.T) {
	store := newBlockingStore()
	h := NewASHandler(store, logger.New())
	handleWhileAWriteIsBlocked(t, store, h,
		[]*ristypes.RISMessage{{Path: slices.Repeat(ristypes.ASPath{testASNA}, asnBatchSize)}},
		&ristypes.RISMessage{Path: ristypes.ASPath{testASNB}})

	store.mu.Lock()
	defer store.mu.Unlock()
	if !store.asns[testASNB] {
		t.Errorf("Stop did not write AS%d, handled while a write was blocked", testASNB)
	}
}

func TestPeerHandlerHandlesMessagesWhileAWriteIsBlocked(t *testing.T) {
	const otherPeerIP = "2001:db8:ffff::2"

	store := newBlockingStore()
	h := NewPeerHandler(store, logger.New())
	handleWhileAWriteIsBlocked(t, store, h,
		slices.Repeat([]*ristypes.RISMessage{{Peer: testPeerIP}}, peerBatchSize),
		&ristypes.RISMessage{Peer: otherPeerIP})

	store.mu.Lock()
	defer store.mu.Unlock()
	if !store.peers[otherPeerIP] {
		t.Errorf("Stop did not write peer %s, handled while a write was blocked", otherPeerIP)
	}
}