notify/record_stream_test.go

package notify

import (
	"context"
	"errors"
	"io"
	"net"
	"strconv"
	"sync"
	"testing"
	"time"
)

func TestRecordStreamBarrierTracksAppliedSeq(t *testing.T) {
	server := NewServer().(*ServerCommon)
	secret := []byte("0123456789abcdef0123456789abcdef")
	server = newRunningPeerAttachServerForTest(t, func(server *ServerCommon) {
		server.SetSecretKey(secret)
	})

	receivedCh := make(chan RecordMessage, 4)
	handlerDone := make(chan error, 1)
	server.SetRecordStreamHandler(func(info RecordAcceptInfo) error {
		for {
			record, err := info.RecordStream.ReadRecord(context.Background())
			if err != nil {
				if errors.Is(err, io.EOF) {
					handlerDone <- nil
					return nil
				}
				handlerDone <- err
				return err
			}
			receivedCh <- record
			if err := info.RecordStream.AckRecord(record.Seq); err != nil {
				handlerDone <- err
				return err
			}
		}
	})

	client := NewClient().(*ClientCommon)
	client.SetSecretKey(secret)
	left, right := net.Pipe()
	defer right.Close()
	bootstrapPeerAttachConnForTest(t, server, right)
	if err := client.ConnectByConn(left); err != nil {
		t.Fatalf("client ConnectByConn failed: %v", err)
	}
	defer func() {
		client.setByeFromServer(true)
		_ = client.Stop()
	}()

	stream, err := client.OpenRecordStream(context.Background(), RecordOpenOptions{})
	if err != nil {
		t.Fatalf("OpenRecordStream failed: %v", err)
	}

	payloads := [][]byte{
		[]byte("alpha"),
		[]byte("beta"),
		[]byte("gamma"),
	}
	for index, payload := range payloads {
		seq, err := stream.WriteRecord(context.Background(), payload)
		if err != nil {
			t.Fatalf("WriteRecord(%d) failed: %v", index, err)
		}
		if got, want := seq, uint64(index+1); got != want {
			t.Fatalf("WriteRecord(%d) seq=%d want=%d", index, got, want)
		}
	}

	ackedSeq, err := stream.Barrier(context.Background())
	if err != nil {
		t.Fatalf("Barrier failed: %v", err)
	}
	if got, want := ackedSeq, uint64(len(payloads)); got != want {
		t.Fatalf("Barrier acked=%d want=%d", got, want)
	}

	for index, want := range payloads {
		select {
		case got := <-receivedCh:
			if got.Seq != uint64(index+1) {
				t.Fatalf("received seq=%d want=%d", got.Seq, index+1)
			}
			if string(got.Payload) != string(want) {
				t.Fatalf("received payload=%q want=%q", string(got.Payload), string(want))
			}
		case <-time.After(2 * time.Second):
			t.Fatalf("timed out waiting received payload %d", index)
		}
	}

	if err := stream.CloseWrite(); err != nil {
		t.Fatalf("CloseWrite failed: %v", err)
	}

	select {
	case err := <-handlerDone:
		if err != nil {
			t.Fatalf("record handler failed: %v", err)
		}
	case <-time.After(2 * time.Second):
		t.Fatal("timed out waiting record handler completion")
	}
}

func TestRecordStreamPropagatesStructuredFailure(t *testing.T) {
	server := NewServer().(*ServerCommon)
	secret := []byte("0123456789abcdef0123456789abcdef")
	server = newRunningPeerAttachServerForTest(t, func(server *ServerCommon) {
		server.SetSecretKey(secret)
	})

	server.SetRecordStreamHandler(func(info RecordAcceptInfo) error {
		record, err := info.RecordStream.ReadRecord(context.Background())
		if err != nil {
			return err
		}
		return info.RecordStream.FailRecord(record.Seq, RecordFailure{
			Code:      "disk_full",
			Retryable: true,
			Message:   "disk full",
		})
	})

	client := NewClient().(*ClientCommon)
	client.SetSecretKey(secret)
	left, right := net.Pipe()
	defer right.Close()
	bootstrapPeerAttachConnForTest(t, server, right)
	if err := client.ConnectByConn(left); err != nil {
		t.Fatalf("client ConnectByConn failed: %v", err)
	}
	defer func() {
		client.setByeFromServer(true)
		_ = client.Stop()
	}()

	stream, err := client.OpenRecordStream(context.Background(), RecordOpenOptions{})
	if err != nil {
		t.Fatalf("OpenRecordStream failed: %v", err)
	}
	if _, err := stream.WriteRecord(context.Background(), []byte("payload")); err != nil {
		t.Fatalf("WriteRecord failed: %v", err)
	}
	if _, err := stream.Barrier(context.Background()); err == nil {
		t.Fatal("Barrier should fail after remote FailRecord")
	} else {
		var failure RecordFailure
		if !errors.As(err, &failure) {
			t.Fatalf("Barrier error=%T %v, want RecordFailure", err, err)
		}
		if got, want := failure.FailedSeq, uint64(1); got != want {
			t.Fatalf("failure seq=%d want=%d", got, want)
		}
		if got, want := failure.Code, RecordErrorCode("disk_full"); got != want {
			t.Fatalf("failure code=%q want=%q", got, want)
		}
		if !failure.Retryable {
			t.Fatal("failure retryable=false, want true")
		}
	}
}

func TestRecordStreamBackpressureUsesUnackedRecords(t *testing.T) {
	server := NewServer().(*ServerCommon)
	secret := []byte("0123456789abcdef0123456789abcdef")
	server = newRunningPeerAttachServerForTest(t, func(server *ServerCommon) {
		server.SetSecretKey(secret)
	})

	firstSeen := make(chan struct{}, 1)
	server.SetRecordStreamHandler(func(info RecordAcceptInfo) error {
		record, err := info.RecordStream.ReadRecord(context.Background())
		if err != nil {
			return err
		}
		if record.Seq != 1 {
			t.Errorf("first record seq=%d want=1", record.Seq)
		}
		firstSeen <- struct{}{}
		time.Sleep(200 * time.Millisecond)
		if err := info.RecordStream.AckRecord(record.Seq); err != nil {
			return err
		}
		for {
			_, err := info.RecordStream.ReadRecord(context.Background())
			if err != nil {
				if errors.Is(err, io.EOF) {
					return nil
				}
				return err
			}
		}
	})

	client := NewClient().(*ClientCommon)
	client.SetSecretKey(secret)
	left, right := net.Pipe()
	defer right.Close()
	bootstrapPeerAttachConnForTest(t, server, right)
	if err := client.ConnectByConn(left); err != nil {
		t.Fatalf("client ConnectByConn failed: %v", err)
	}
	defer func() {
		client.setByeFromServer(true)
		_ = client.Stop()
	}()

	stream, err := client.OpenRecordStream(context.Background(), RecordOpenOptions{
		MaxUnackedRecords: 1,
		MaxUnackedBytes:   1 << 20,
	})
	if err != nil {
		t.Fatalf("OpenRecordStream failed: %v", err)
	}
	if _, err := stream.WriteRecord(context.Background(), []byte("first")); err != nil {
		t.Fatalf("first WriteRecord failed: %v", err)
	}
	select {
	case <-firstSeen:
	case <-time.After(2 * time.Second):
		t.Fatal("timed out waiting server to receive first record")
	}

	writeCtx, cancel := context.WithTimeout(context.Background(), 50*time.Millisecond)
	defer cancel()
	if _, err := stream.WriteRecord(writeCtx, []byte("second")); !errors.Is(err, context.DeadlineExceeded) {
		t.Fatalf("second WriteRecord error=%v want=%v", err, context.DeadlineExceeded)
	}

	if acked, err := stream.Barrier(context.Background()); err != nil {
		t.Fatalf("Barrier failed: %v", err)
	} else if got, want := acked, uint64(1); got != want {
		t.Fatalf("Barrier acked=%d want=%d", got, want)
	}

	if err := stream.CloseWrite(); err != nil {
		t.Fatalf("CloseWrite failed: %v", err)
	}
}

func TestRecordStreamBarrierToWaitsCheckpointSeq(t *testing.T) {
	server := NewServer().(*ServerCommon)
	secret := []byte("0123456789abcdef0123456789abcdef")
	server = newRunningPeerAttachServerForTest(t, func(server *ServerCommon) {
		server.SetSecretKey(secret)
	})

	secondSeen := make(chan struct{}, 1)
	server.SetRecordStreamHandler(func(info RecordAcceptInfo) error {
		first, err := info.RecordStream.ReadRecord(context.Background())
		if err != nil {
			return err
		}
		if err := info.RecordStream.AckRecord(first.Seq); err != nil {
			return err
		}

		second, err := info.RecordStream.ReadRecord(context.Background())
		if err != nil {
			return err
		}
		secondSeen <- struct{}{}
		time.Sleep(200 * time.Millisecond)
		if err := info.RecordStream.AckRecord(second.Seq); err != nil {
			return err
		}

		for {
			_, err := info.RecordStream.ReadRecord(context.Background())
			if err != nil {
				if errors.Is(err, io.EOF) {
					return nil
				}
				return err
			}
		}
	})

	client := NewClient().(*ClientCommon)
	client.SetSecretKey(secret)
	left, right := net.Pipe()
	defer right.Close()
	bootstrapPeerAttachConnForTest(t, server, right)
	if err := client.ConnectByConn(left); err != nil {
		t.Fatalf("client ConnectByConn failed: %v", err)
	}
	defer func() {
		client.setByeFromServer(true)
		_ = client.Stop()
	}()

	stream, err := client.OpenRecordStream(context.Background(), RecordOpenOptions{})
	if err != nil {
		t.Fatalf("OpenRecordStream failed: %v", err)
	}
	firstSeq, err := stream.WriteRecord(context.Background(), []byte("first"))
	if err != nil {
		t.Fatalf("first WriteRecord failed: %v", err)
	}
	if _, err := stream.WriteRecord(context.Background(), []byte("second")); err != nil {
		t.Fatalf("second WriteRecord failed: %v", err)
	}

	select {
	case <-secondSeen:
	case <-time.After(2 * time.Second):
		t.Fatal("timed out waiting server to receive second record")
	}

	start := time.Now()
	acked, err := stream.BarrierTo(context.Background(), firstSeq)
	if err != nil {
		t.Fatalf("BarrierTo failed: %v", err)
	}
	if acked != firstSeq {
		t.Fatalf("BarrierTo acked=%d want=%d", acked, firstSeq)
	}
	if elapsed := time.Since(start); elapsed >= 150*time.Millisecond {
		t.Fatalf("BarrierTo waited too long: %s", elapsed)
	}

	if _, err := stream.Barrier(context.Background()); err != nil {
		t.Fatalf("final Barrier failed: %v", err)
	}
	if err := stream.CloseWrite(); err != nil {
		t.Fatalf("CloseWrite failed: %v", err)
	}
}

func TestRecordStreamConcurrentWritesStayOrdered(t *testing.T) {
	server := NewServer().(*ServerCommon)
	secret := []byte("0123456789abcdef0123456789abcdef")
	server = newRunningPeerAttachServerForTest(t, func(server *ServerCommon) {
		server.SetSecretKey(secret)
	})

	const total = 64
	receivedCh := make(chan RecordMessage, total)
	handlerDone := make(chan error, 1)
	server.SetRecordStreamHandler(func(info RecordAcceptInfo) error {
		for {
			record, err := info.RecordStream.ReadRecord(context.Background())
			if err != nil {
				if errors.Is(err, io.EOF) {
					handlerDone <- nil
					return nil
				}
				handlerDone <- err
				return err
			}
			receivedCh <- record
			if err := info.RecordStream.AckRecord(record.Seq); err != nil {
				handlerDone <- err
				return err
			}
		}
	})

	client := NewClient().(*ClientCommon)
	client.SetSecretKey(secret)
	left, right := net.Pipe()
	defer right.Close()
	bootstrapPeerAttachConnForTest(t, server, right)
	if err := client.ConnectByConn(left); err != nil {
		t.Fatalf("client ConnectByConn failed: %v", err)
	}
	defer func() {
		client.setByeFromServer(true)
		_ = client.Stop()
	}()

	stream, err := client.OpenRecordStream(context.Background(), RecordOpenOptions{})
	if err != nil {
		t.Fatalf("OpenRecordStream failed: %v", err)
	}

	var wg sync.WaitGroup
	for index := 0; index < total; index++ {
		index := index
		wg.Add(1)
		go func() {
			defer wg.Done()
			payload := []byte("item-" + strconv.Itoa(index))
			if _, err := stream.WriteRecord(context.Background(), payload); err != nil {
				t.Errorf("WriteRecord(%d) failed: %v", index, err)
			}
		}()
	}
	wg.Wait()

	if acked, err := stream.Barrier(context.Background()); err != nil {
		t.Fatalf("Barrier failed: %v", err)
	} else if got, want := acked, uint64(total); got != want {
		t.Fatalf("Barrier acked=%d want=%d", got, want)
	}

	seen := make(map[uint64]string, total)
	for index := 0; index < total; index++ {
		select {
		case record := <-receivedCh:
			seen[record.Seq] = string(record.Payload)
		case <-time.After(2 * time.Second):
			t.Fatalf("timed out waiting record %d", index)
		}
	}
	for seq := 1; seq <= total; seq++ {
		payload, ok := seen[uint64(seq)]
		if !ok {
			t.Fatalf("missing seq %d", seq)
		}
		if payload == "" {
			t.Fatalf("empty payload at seq %d", seq)
		}
	}

	if err := stream.CloseWrite(); err != nil {
		t.Fatalf("CloseWrite failed: %v", err)
	}

	select {
	case err := <-handlerDone:
		if err != nil {
			t.Fatalf("record handler failed: %v", err)
		}
	case <-time.After(2 * time.Second):
		t.Fatal("timed out waiting record handler completion")
	}
}