notify/bulk_dataid_test.go

package notify

import (
	"context"
	"errors"
	"math"
	"sync"
	"testing"
	"time"
)

func TestBulkRuntimeSeparatesBidirectionalDataIDNamespaces(t *testing.T) {
	clientRuntime := newBulkRuntime("cblk")
	serverRuntime := newBulkRuntime("sblk")
	clientID, err := clientRuntime.reserveDataID("peer", 0)
	if err != nil {
		t.Fatalf("reserve client data id: %v", err)
	}
	serverID, err := serverRuntime.reserveDataID("peer", 0)
	if err != nil {
		t.Fatalf("reserve server data id: %v", err)
	}
	if clientID == serverID || clientID%2 != 1 || serverID%2 != 0 {
		t.Fatalf("client/server data ids = %d/%d, want disjoint odd/even namespaces", clientID, serverID)
	}

	clientBulk := newBulkHandle(context.Background(), clientRuntime, "peer", BulkOpenRequest{BulkID: "client", DataID: clientID}, 0, nil, nil, 0, nil, nil, nil, nil, nil)
	serverBulk := newBulkHandle(context.Background(), serverRuntime, "peer", BulkOpenRequest{BulkID: "server", DataID: serverID}, 0, nil, nil, 0, nil, nil, nil, nil, nil)
	if err := clientRuntime.registerReserved("peer", clientBulk); err != nil {
		t.Fatalf("register client bulk: %v", err)
	}
	if err := serverRuntime.registerReserved("peer", serverBulk); err != nil {
		t.Fatalf("register server bulk: %v", err)
	}
	if got, ok := clientRuntime.lookupInboundFrame("peer", clientID); !ok || got != clientBulk {
		t.Fatalf("client local frame lookup = %p/%v, want client bulk", got, ok)
	}
	if got, ok := clientRuntime.lookupInboundFrame("peer", serverID); ok || got != nil {
		t.Fatalf("client peer frame lookup = %p/%v, want missing inbound bulk", got, ok)
	}
}

func TestBulkRuntimeRoutesLegacyZeroDataIDInboundFrames(t *testing.T) {
	for _, role := range []string{"cblk", "sblk"} {
		t.Run(role, func(t *testing.T) {
			runtime := newBulkRuntime(role)
			bulk := newBulkHandle(context.Background(), runtime, "peer", BulkOpenRequest{
				BulkID: "legacy-inbound",
				// A legacy initiator leaves DataID unset and uses the ID
				// allocated by the receiver's open response.
			}, 0, nil, nil, 0, nil, nil, nil, nil, nil)
			if err := runtime.registerInbound("peer", bulk); err != nil {
				t.Fatalf("register legacy inbound bulk: %v", err)
			}
			if got := bulk.dataIDSnapshot(); got == 0 {
				t.Fatal("legacy inbound registration allocated zero data id")
			}
			got, ok := runtime.lookupInboundFrame("peer", bulk.dataIDSnapshot())
			if !ok || got != bulk {
				t.Fatalf("legacy inbound frame lookup = %p/%v, want %p/true", got, ok, bulk)
			}
		})
	}
}

func TestBulkRuntimeAllocatorSkipsLegacyInboundCollision(t *testing.T) {
	runtime := newBulkRuntime("cblk")
	inbound := newBulkHandle(context.Background(), runtime, "peer", BulkOpenRequest{
		BulkID: "legacy-inbound",
		DataID: 1,
	}, 0, nil, nil, 0, nil, nil, nil, nil, nil)
	if err := runtime.registerInbound("peer", inbound); err != nil {
		t.Fatalf("register legacy inbound bulk: %v", err)
	}
	reserved, err := runtime.reserveDataID("peer", 0)
	if err != nil {
		t.Fatalf("reserve outbound data id: %v", err)
	}
	if reserved == inbound.dataIDSnapshot() {
		t.Fatalf("outbound allocator reused legacy inbound data id %d", reserved)
	}
}

func TestBulkRuntimeInboundExplicitDataIDCannotExhaustOutboundAllocator(t *testing.T) {
	tests := []struct {
		role     string
		poisonID uint64
		wantID   uint64
	}{
		{role: "cblk", poisonID: math.MaxUint64, wantID: 1},
		{role: "sblk", poisonID: math.MaxUint64 - 1, wantID: 2},
	}
	for _, test := range tests {
		t.Run(test.role, func(t *testing.T) {
			runtime := newBulkRuntime(test.role)
			inbound := newBulkHandle(context.Background(), runtime, "peer", BulkOpenRequest{
				BulkID: "peer-controlled",
				DataID: test.poisonID,
			}, 0, nil, nil, 0, nil, nil, nil, nil, nil)
			if err := runtime.registerInbound("peer", inbound); err != nil {
				t.Fatalf("register inbound bulk: %v", err)
			}
			got, err := runtime.reserveDataID("peer", 0)
			if err != nil {
				t.Fatalf("reserve outbound data id after peer-controlled id: %v", err)
			}
			if got != test.wantID {
				t.Fatalf("reserved outbound data id = %d, want %d", got, test.wantID)
			}
		})
	}
}

func TestBulkRuntimeStaleFinalizeDoesNotRemoveReplacement(t *testing.T) {
	runtime := newBulkRuntime("cblk")
	old := newBulkHandle(context.Background(), runtime, "peer", BulkOpenRequest{
		BulkID: "reused",
		DataID: 1,
	}, 0, nil, nil, 0, nil, nil, nil, nil, nil)
	if err := runtime.registerInbound("peer", old); err != nil {
		t.Fatalf("register old bulk: %v", err)
	}
	old.markReset(errors.New("old failed"))

	replacement := newBulkHandle(context.Background(), runtime, "peer", BulkOpenRequest{
		BulkID: "reused",
		DataID: 3,
	}, 0, nil, nil, 0, nil, nil, nil, nil, nil)
	if err := runtime.registerInbound("peer", replacement); err != nil {
		t.Fatalf("register replacement bulk: %v", err)
	}

	old.markReset(errors.New("late duplicate reset"))
	if got, ok := runtime.lookup("peer", "reused"); !ok || got != replacement {
		t.Fatalf("replacement bulk after stale finalize = %p/%v, want %p/true", got, ok, replacement)
	}
}

func TestBulkRuntimeControlLookupRejectsMismatchedIdentity(t *testing.T) {
	runtime := newBulkRuntime("control")
	inbound := newBulkHandle(context.Background(), runtime, "peer", BulkOpenRequest{BulkID: "inbound", DataID: 7}, 0, nil, nil, 0, nil, nil, nil, nil, nil)
	outbound := newBulkHandle(context.Background(), runtime, "peer", BulkOpenRequest{BulkID: "outbound", DataID: 7}, 0, nil, nil, 0, nil, nil, nil, nil, nil)
	if err := runtime.registerInbound("peer", inbound); err != nil {
		t.Fatalf("register inbound bulk: %v", err)
	}
	if err := runtime.registerOutbound("peer", outbound); err != nil {
		t.Fatalf("register outbound bulk: %v", err)
	}
	if got, ok := runtime.lookupControl("peer", "outbound", 7); !ok || got != outbound {
		t.Fatalf("matching control lookup = %p/%v, want outbound", got, ok)
	}
	if got, ok := runtime.lookupControl("peer", "outbound", 0); !ok || got != outbound {
		t.Fatalf("BulkID-only control lookup = %p/%v, want outbound", got, ok)
	}
	if got, ok := runtime.lookupControl("peer", "outbound", 8); ok || got != nil {
		t.Fatalf("mismatched control lookup = %p/%v, want rejection", got, ok)
	}
	if got, ok := runtime.lookupControl("peer", "", 7); ok || got != nil {
		t.Fatalf("ambiguous data-only control lookup = %p/%v, want rejection", got, ok)
	}
}

func TestBulkCloseControlRejectsMismatchedDataID(t *testing.T) {
	t.Run("client", func(t *testing.T) {
		client := NewClient().(*ClientCommon)
		runtime := client.getBulkRuntime()
		bulk := newBulkHandle(context.Background(), runtime, clientFileScope(), BulkOpenRequest{
			BulkID: "close-client",
			DataID: 11,
		}, 0, nil, nil, 0, nil, nil, nil, nil, nil)
		if err := runtime.register(clientFileScope(), bulk); err != nil {
			t.Fatalf("register client bulk: %v", err)
		}
		defer bulk.markReset(errors.New("test cleanup"))
		payload, err := encode(BulkCloseRequest{BulkID: bulk.ID(), DataID: 13, Full: true})
		if err != nil {
			t.Fatalf("encode client close: %v", err)
		}
		client.handleInboundBulkClose(&Message{
			NetType:     NET_CLIENT,
			ServerConn:  client,
			TransferMsg: TransferMsg{Key: BulkCloseSignalKey, Value: payload, Type: MSG_ASYNC},
		})
		bulk.mu.Lock()
		defer bulk.mu.Unlock()
		if bulk.remoteClosed || bulk.peerReadClosed || bulk.resetErr != nil {
			t.Fatalf("mismatched client close mutated bulk: remote=%v peer=%v reset=%v", bulk.remoteClosed, bulk.peerReadClosed, bulk.resetErr)
		}
	})

	t.Run("server", func(t *testing.T) {
		server := NewServer().(*ServerCommon)
		logical := server.bootstrapAcceptedLogical("close-server", nil, nil)
		if logical == nil {
			t.Fatal("bootstrap server logical connection failed")
		}
		runtime := server.getBulkRuntime()
		scope := serverFileScope(logical)
		bulk := newBulkHandle(context.Background(), runtime, scope, BulkOpenRequest{
			BulkID: "close-server",
			DataID: 17,
		}, 0, logical, nil, 0, nil, nil, nil, nil, nil)
		if err := runtime.register(scope, bulk); err != nil {
			t.Fatalf("register server bulk: %v", err)
		}
		defer bulk.markReset(errors.New("test cleanup"))
		payload, err := encode(BulkCloseRequest{BulkID: bulk.ID(), DataID: 19, Full: true})
		if err != nil {
			t.Fatalf("encode server close: %v", err)
		}
		server.handleInboundBulkClose(&Message{
			NetType:       NET_SERVER,
			LogicalConn:   logical,
			TransportConn: nil,
			TransferMsg:   TransferMsg{Key: BulkCloseSignalKey, Value: payload, Type: MSG_ASYNC},
		})
		bulk.mu.Lock()
		defer bulk.mu.Unlock()
		if bulk.remoteClosed || bulk.peerReadClosed || bulk.resetErr != nil {
			t.Fatalf("mismatched server close mutated bulk: remote=%v peer=%v reset=%v", bulk.remoteClosed, bulk.peerReadClosed, bulk.resetErr)
		}
	})
}

func TestBulkOpenDedicatedConcurrentlyFromBothPeers(t *testing.T) {
	server := NewServer().(*ServerCommon)
	if err := UseModernPSKServer(server, integrationSharedSecret, integrationModernPSKOptions()); err != nil {
		t.Fatalf("UseModernPSKServer failed: %v", err)
	}
	serverAccepted := make(chan BulkAcceptInfo, 1)
	server.SetBulkHandler(func(info BulkAcceptInfo) error {
		serverAccepted <- info
		return nil
	})
	if err := server.Listen("tcp", "127.0.0.1:0"); err != nil {
		t.Fatalf("server Listen failed: %v", err)
	}
	defer func() { _ = server.Stop() }()

	client := NewClient().(*ClientCommon)
	if err := UseModernPSKClient(client, integrationSharedSecret, integrationModernPSKOptions()); err != nil {
		t.Fatalf("UseModernPSKClient failed: %v", err)
	}
	clientAccepted := make(chan BulkAcceptInfo, 1)
	client.SetBulkHandler(func(info BulkAcceptInfo) error {
		clientAccepted <- info
		return nil
	})
	if err := client.Connect("tcp", server.listener.Addr().String()); err != nil {
		t.Fatalf("client Connect failed: %v", err)
	}
	defer func() { _ = client.Stop() }()

	var logical *LogicalConn
	deadline := time.Now().Add(2 * time.Second)
	for time.Now().Before(deadline) {
		peers := server.GetLogicalConnList()
		if len(peers) > 0 {
			logical = peers[0]
			break
		}
		time.Sleep(time.Millisecond)
	}
	if logical == nil {
		t.Fatal("timed out waiting for server logical connection")
	}

	type openResult struct {
		bulk Bulk
		err  error
	}
	clientResult := make(chan openResult, 1)
	serverResult := make(chan openResult, 1)
	go func() {
		bulk, err := client.OpenDedicatedBulk(context.Background(), BulkOpenOptions{Range: BulkRange{Length: 1}})
		clientResult <- openResult{bulk: bulk, err: err}
	}()
	go func() {
		bulk, err := server.OpenBulkLogical(context.Background(), logical, BulkOpenOptions{Range: BulkRange{Offset: 1, Length: 1}})
		serverResult <- openResult{bulk: bulk, err: err}
	}()
	clientOpen := <-clientResult
	serverOpen := <-serverResult
	if clientOpen.err != nil || serverOpen.err != nil {
		t.Fatalf("concurrent dedicated opens failed: client=%v server=%v", clientOpen.err, serverOpen.err)
	}
	clientInbound := waitAcceptedBulk(t, serverAccepted, 2*time.Second)
	serverInbound := waitAcceptedBulk(t, clientAccepted, 2*time.Second)
	if clientOpen.bulk.(*bulkHandle).dataIDSnapshot()%2 != 1 || serverOpen.bulk.(*bulkHandle).dataIDSnapshot()%2 != 0 {
		t.Fatalf("local data ids = %d/%d, want odd/even", clientOpen.bulk.(*bulkHandle).dataIDSnapshot(), serverOpen.bulk.(*bulkHandle).dataIDSnapshot())
	}
	if clientInbound.Bulk.(*bulkHandle).dataIDSnapshot() != clientOpen.bulk.(*bulkHandle).dataIDSnapshot() {
		t.Fatalf("client-open data id mismatch across peers")
	}
	if serverInbound.Bulk.(*bulkHandle).dataIDSnapshot() != serverOpen.bulk.(*bulkHandle).dataIDSnapshot() {
		t.Fatalf("server-open data id mismatch across peers")
	}
	if _, err := clientOpen.bulk.Write([]byte("client")); err != nil {
		t.Fatalf("client bulk write failed: %v", err)
	}
	readBulkExactly(t, clientInbound.Bulk, "client", 2*time.Second)
	if _, err := serverOpen.bulk.Write([]byte("server")); err != nil {
		t.Fatalf("server bulk write failed: %v", err)
	}
	readBulkExactly(t, serverInbound.Bulk, "server", 2*time.Second)
	_ = clientOpen.bulk.Close()
	_ = clientInbound.Bulk.Close()
	_ = serverOpen.bulk.Close()
	_ = serverInbound.Bulk.Close()
}

func TestBulkRuntimeDataIDAllocatorObservesExplicitIDs(t *testing.T) {
	runtime := newBulkRuntime("dataid")
	scope := "peer"

	explicit := newBulkHandle(context.Background(), runtime, scope, BulkOpenRequest{
		BulkID: "explicit",
		DataID: 41,
	}, 0, nil, nil, 0, nil, nil, nil, nil, nil)
	if err := runtime.register(scope, explicit); err != nil {
		t.Fatalf("register explicit bulk: %v", err)
	}
	if got := explicit.dataIDSnapshot(); got != 41 {
		t.Fatalf("explicit data id = %d, want 41", got)
	}

	auto := newBulkHandle(context.Background(), runtime, scope, BulkOpenRequest{
		BulkID: "auto",
	}, 0, nil, nil, 0, nil, nil, nil, nil, nil)
	if err := runtime.register(scope, auto); err != nil {
		t.Fatalf("register auto bulk: %v", err)
	}
	if got := auto.dataIDSnapshot(); got != 42 {
		t.Fatalf("auto data id = %d, want 42 after explicit 41", got)
	}

	reserved, err := runtime.reserveDataID(scope, 0)
	if err != nil {
		t.Fatalf("reserve data id: %v", err)
	}
	if reserved != 43 {
		t.Fatalf("reserved data id = %d, want 43", reserved)
	}
	autoWhileReserved := newBulkHandle(context.Background(), runtime, scope, BulkOpenRequest{
		BulkID: "auto-while-reserved",
	}, 0, nil, nil, 0, nil, nil, nil, nil, nil)
	if err := runtime.register(scope, autoWhileReserved); err != nil {
		t.Fatalf("register auto bulk while id is reserved: %v", err)
	}
	if got := autoWhileReserved.dataIDSnapshot(); got != 44 {
		t.Fatalf("auto data id while 43 is reserved = %d, want 44", got)
	}
	reservedBulk := newBulkHandle(context.Background(), runtime, scope, BulkOpenRequest{
		BulkID: "reserved",
		DataID: reserved,
	}, 0, nil, nil, 0, nil, nil, nil, nil, nil)
	if err := runtime.registerReserved(scope, reservedBulk); err != nil {
		t.Fatalf("register reserved bulk: %v", err)
	}

	if _, err := runtime.reserveDataID(scope, reserved); err == nil {
		t.Fatal("re-reserving an active data id should fail")
	}
}

func TestBulkRuntimeDataIDAllocatorConcurrentReservationsAreUnique(t *testing.T) {
	runtime := newBulkRuntime("dataid-concurrent")
	const count = 128
	ids := make(chan uint64, count)
	errs := make(chan error, count)
	var wg sync.WaitGroup
	for i := 0; i < count; i++ {
		wg.Add(1)
		go func() {
			defer wg.Done()
			id, err := runtime.reserveDataID("peer", 0)
			if err != nil {
				errs <- err
				return
			}
			ids <- id
		}()
	}
	wg.Wait()
	close(ids)
	close(errs)
	for err := range errs {
		t.Fatalf("reserve data id: %v", err)
	}

	seen := make(map[uint64]struct{}, count)
	for id := range ids {
		if _, ok := seen[id]; ok {
			t.Fatalf("duplicate reserved data id %d", id)
		}
		seen[id] = struct{}{}
	}
	if len(seen) != count {
		t.Fatalf("reserved id count = %d, want %d", len(seen), count)
	}
}

func TestBulkMixedDedicatedThenSharedUsesFreshDataID(t *testing.T) {
	server := NewServer().(*ServerCommon)
	if err := UseModernPSKServer(server, integrationSharedSecret, integrationModernPSKOptions()); err != nil {
		t.Fatalf("UseModernPSKServer failed: %v", err)
	}
	acceptCh := make(chan BulkAcceptInfo, 2)
	server.SetBulkHandler(func(info BulkAcceptInfo) error {
		acceptCh <- info
		return nil
	})
	if err := server.Listen("tcp", "127.0.0.1:0"); err != nil {
		t.Fatalf("server Listen failed: %v", err)
	}
	defer func() { _ = server.Stop() }()

	client := NewClient().(*ClientCommon)
	if err := UseModernPSKClient(client, integrationSharedSecret, integrationModernPSKOptions()); err != nil {
		t.Fatalf("UseModernPSKClient failed: %v", err)
	}
	if err := client.Connect("tcp", server.listener.Addr().String()); err != nil {
		t.Fatalf("client Connect failed: %v", err)
	}
	defer func() { _ = client.Stop() }()

	first, err := client.OpenDedicatedBulk(context.Background(), BulkOpenOptions{
		Range: BulkRange{Offset: 0, Length: 1},
	})
	if err != nil {
		t.Fatalf("open dedicated bulk: %v", err)
	}
	firstAccepted := waitAcceptedBulk(t, acceptCh, 2*time.Second)
	firstID := first.(*bulkHandle).dataIDSnapshot()
	if firstID == 0 || firstAccepted.Bulk.(*bulkHandle).dataIDSnapshot() != firstID {
		t.Fatalf("dedicated data id mismatch: client=%d server=%d", firstID, firstAccepted.Bulk.(*bulkHandle).dataIDSnapshot())
	}
	_ = first.Close()
	_ = firstAccepted.Bulk.Close()

	second, err := client.OpenSharedBulk(context.Background(), BulkOpenOptions{
		Range: BulkRange{Offset: 1, Length: 1},
	})
	if err != nil {
		t.Fatalf("open shared bulk after dedicated: %v", err)
	}
	secondAccepted := waitAcceptedBulk(t, acceptCh, 2*time.Second)
	secondID := second.(*bulkHandle).dataIDSnapshot()
	if secondID <= firstID {
		t.Fatalf("shared data id = %d, want greater than dedicated id %d", secondID, firstID)
	}
	if got := secondAccepted.Bulk.(*bulkHandle).dataIDSnapshot(); got != secondID {
		t.Fatalf("shared data id mismatch: client=%d server=%d", secondID, got)
	}
	_ = second.Close()
	_ = secondAccepted.Bulk.Close()
}