routewatch/internal/routewatch/batchwrite_internal_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"
)

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

	// stillWaitingAfter is how long a test watches a HandleMessage that must
	// wait. One that does not wait returns well within it.
	stillWaitingAfter = 100 * time.Millisecond
)

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

	release      chan struct{}
	releaseOnce  sync.Once
	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) 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
}

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

// releaseWrites lets the held writes, and all later ones, finish. A second call
// does nothing.
func (s *blockingStore) releaseWrites() {
	s.releaseOnce.Do(func() { close(s.release) })
}

// 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
	}
}

// startBlockedWrite hands h the messages in fill, which fill its batch so that
// the handler writes it to store, and waits until that write starts. store
// holds it until the test calls releaseWrites.
func startBlockedWrite(
	t *testing.T,
	store *blockingStore,
	h interface {
		HandleMessage(msg *ristypes.RISMessage)
	},
	fill []*ristypes.RISMessage,
) {
	t.Helper()

	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")
	}
}

// handleWhileAWriteIsBlocked starts a blocked write with the messages in fill
// 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(msg *ristypes.RISMessage)
		Stop()
	},
	fill []*ristypes.RISMessage,
	extra *ristypes.RISMessage,
) {
	t.Helper()

	defer h.Stop()
	defer store.releaseWrites()

	startBlockedWrite(t, store, h, fill)

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

// handleAtMaxBatchSize starts a blocked write with the messages in fill, then,
// while store holds the write, hands h the messages in toMax, which bring its
// next batch to its maximum size, and the message extra. HandleMessage for
// extra must wait until the test lets the write finish, and return after. Then
// it stops h, which writes what it holds.
func handleAtMaxBatchSize(
	t *testing.T,
	store *blockingStore,
	h interface {
		HandleMessage(msg *ristypes.RISMessage)
		Stop()
	},
	fill, toMax []*ristypes.RISMessage,
	extra *ristypes.RISMessage,
) {
	t.Helper()

	defer h.Stop()
	defer store.releaseWrites()

	startBlockedWrite(t, store, h, fill)

	for _, msg := range toMax {
		h.HandleMessage(msg)
	}

	returned := make(chan struct{})

	go func() {
		h.HandleMessage(extra)
		close(returned)
	}()

	select {
	case <-returned:
		t.Fatal("HandleMessage returned while the batch was at its maximum " +
			"size and a write was blocked")
	case <-time.After(stillWaitingAfter):
	}

	store.releaseWrites()

	select {
	case <-returned:
	case <-time.After(handleMessageTimeout):
		t.Fatal("HandleMessage did not return after the database write finished")
	}
}

// prefixAnnouncement returns a message in which testPeerIP announces prefixes
// with origin AS testASNA.
func prefixAnnouncement(prefixes []string) *ristypes.RISMessage {
	return &ristypes.RISMessage{
		Peer: testPeerIP,
		Path: ristypes.ASPath{Sequence: []int{testASNA}},
		Announcements: []ristypes.RISAnnouncement{
			{NextHop: testPeerIP, Prefixes: prefixes},
		},
	}
}

func TestPrefixHandlerHandlesMessagesWhileAWriteIsBlocked(t *testing.T) {
	t.Parallel()

	const extraPrefix = "198.51.100.0/24"

	store := newBlockingStore()
	h := NewPrefixHandler(store, logger.New())
	fill := slices.Repeat([]string{testV4Prefix}, prefixBatchSize)
	handleWhileAWriteIsBlocked(t, store, h,
		[]*ristypes.RISMessage{prefixAnnouncement(fill)},
		prefixAnnouncement([]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) {
	t.Parallel()

	store := newBlockingStore()
	h := NewASHandler(store, logger.New())
	fill := ristypes.ASPath{Sequence: slices.Repeat([]int{testASNA}, asnBatchSize)}
	handleWhileAWriteIsBlocked(t, store, h,
		[]*ristypes.RISMessage{{Path: fill}},
		&ristypes.RISMessage{Path: ristypes.ASPath{Sequence: []int{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) {
	t.Parallel()

	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)
	}
}

func TestPrefixHandlerWaitsAtTheMaximumBatchSize(t *testing.T) {
	t.Parallel()

	const extraPrefix = "198.51.100.0/24"

	store := newBlockingStore()
	h := NewPrefixHandler(store, logger.New())
	fill := slices.Repeat([]string{testV4Prefix}, prefixBatchSize)
	toMax := slices.Repeat([]string{testV4Prefix}, prefixBatchMaxSize)
	handleAtMaxBatchSize(t, store, h,
		[]*ristypes.RISMessage{prefixAnnouncement(fill)},
		[]*ristypes.RISMessage{prefixAnnouncement(toMax)},
		prefixAnnouncement([]string{extraPrefix}))

	store.mu.Lock()
	defer store.mu.Unlock()

	if !store.prefixes[extraPrefix] {
		t.Errorf("Stop did not write %s, handled after the batch reached its "+
			"maximum size", extraPrefix)
	}
}

func TestASHandlerWaitsAtTheMaximumBatchSize(t *testing.T) {
	t.Parallel()

	store := newBlockingStore()
	h := NewASHandler(store, logger.New())
	fill := ristypes.ASPath{Sequence: slices.Repeat([]int{testASNA}, asnBatchSize)}
	toMax := ristypes.ASPath{Sequence: slices.Repeat([]int{testASNA}, asnBatchMaxSize)}
	handleAtMaxBatchSize(t, store, h,
		[]*ristypes.RISMessage{{Path: fill}},
		[]*ristypes.RISMessage{{Path: toMax}},
		&ristypes.RISMessage{Path: ristypes.ASPath{Sequence: []int{testASNB}}})

	store.mu.Lock()
	defer store.mu.Unlock()

	if !store.asns[testASNB] {
		t.Errorf("Stop did not write AS%d, handled after the batch reached its "+
			"maximum size", testASNB)
	}
}

func TestPeerHandlerWaitsAtTheMaximumBatchSize(t *testing.T) {
	t.Parallel()

	const otherPeerIP = "2001:db8:ffff::2"

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

	store.mu.Lock()
	defer store.mu.Unlock()

	if !store.peers[otherPeerIP] {
		t.Errorf("Stop did not write peer %s, handled after the batch reached "+
			"its maximum size", otherPeerIP)
	}
}

// TestHandlersWakeTheFlushGoroutineWhenTheBatchIsFull hands each handler, built
// without its flush goroutine so that nothing takes a wake-up, one update short
// of a full batch, and a second one exactly a full batch. Only the full batch
// leaves a wake-up in batchFullCh.
func TestHandlersWakeTheFlushGoroutineWhenTheBatchIsFull(t *testing.T) {
	t.Parallel()

	tests := []struct {
		name      string
		batchSize int
		// handle hands a new handler n updates and returns its batchFullCh.
		handle func(n int) chan struct{}
	}{
		{
			name:      "prefix",
			batchSize: prefixBatchSize,
			handle: func(n int) chan struct{} {
				h := &PrefixHandler{batchFullCh: make(chan struct{}, 1)}
				h.HandleMessage(prefixAnnouncement(
					slices.Repeat([]string{testV4Prefix}, n)))

				return h.batchFullCh
			},
		},
		{
			name:      "AS",
			batchSize: asnBatchSize,
			handle: func(n int) chan struct{} {
				h := &ASHandler{batchFullCh: make(chan struct{}, 1)}
				// The origin AS and each AS in the path are one update each.
				h.HandleMessage(&ristypes.RISMessage{
					Path: ristypes.ASPath{Sequence: slices.Repeat([]int{testASNA}, n-1)},
				})

				return h.batchFullCh
			},
		},
		{
			name:      "peer",
			batchSize: peerBatchSize,
			handle: func(n int) chan struct{} {
				h := &PeerHandler{batchFullCh: make(chan struct{}, 1)}
				for range n {
					h.HandleMessage(&ristypes.RISMessage{Peer: testPeerIP})
				}

				return h.batchFullCh
			},
		},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			t.Parallel()

			got := len(tt.handle(tt.batchSize - 1))
			if got != 0 {
				t.Errorf("one update short of a full batch left %d wake-ups, "+
					"want none", got)
			}

			got = len(tt.handle(tt.batchSize))
			if got != 1 {
				t.Errorf("a full batch left %d wake-ups, want 1", got)
			}
		})
	}
}