upstream-seaweedfs/weed/mq/kafka/integration/broker_client_restart_test.go

package integration

import (
	"context"
	"testing"

	"github.com/seaweedfs/seaweedfs/weed/pb/mq_pb"
	"google.golang.org/grpc/metadata"
)

// MockSubscribeStream implements mq_pb.SeaweedMessaging_SubscribeMessageClient for testing
type MockSubscribeStream struct {
	sendCalls []interface{}
	closed    bool
}

func (m *MockSubscribeStream) Send(req *mq_pb.SubscribeMessageRequest) error {
	m.sendCalls = append(m.sendCalls, req)
	return nil
}

func (m *MockSubscribeStream) Recv() (*mq_pb.SubscribeMessageResponse, error) {
	return nil, nil
}

func (m *MockSubscribeStream) CloseSend() error {
	m.closed = true
	return nil
}

func (m *MockSubscribeStream) Header() (metadata.MD, error) { return nil, nil }
func (m *MockSubscribeStream) Trailer() metadata.MD         { return nil }
func (m *MockSubscribeStream) Context() context.Context     { return context.Background() }
func (m *MockSubscribeStream) SendMsg(m2 interface{}) error { return nil }
func (m *MockSubscribeStream) RecvMsg(m2 interface{}) error { return nil }

// TestNeedsRestart tests the NeedsRestart logic
func TestNeedsRestart(t *testing.T) {
	bc := &BrokerClient{}

	tests := []struct {
		name            string
		session         *BrokerSubscriberSession
		requestedOffset int64
		want            bool
		reason          string
	}{
		{
			name: "Stream is nil - needs restart",
			session: &BrokerSubscriberSession{
				Topic:       "test-topic",
				Partition:   0,
				StartOffset: 100,
				Stream:      nil,
			},
			requestedOffset: 100,
			want:            true,
			reason:          "Stream is nil",
		},
		{
			name: "Offset in cache - no restart needed",
			session: &BrokerSubscriberSession{
				Topic:       "test-topic",
				Partition:   0,
				StartOffset: 100,
				Stream:      &MockSubscribeStream{},
				Ctx:         context.Background(),
				consumedRecords: []*SeaweedRecord{
					{Offset: 95},
					{Offset: 96},
					{Offset: 97},
					{Offset: 98},
					{Offset: 99},
				},
			},
			requestedOffset: 97,
			want:            false,
			reason:          "Offset 97 is in cache [95-99]",
		},
		{
			name: "Offset before current - needs restart",
			session: &BrokerSubscriberSession{
				Topic:       "test-topic",
				Partition:   0,
				StartOffset: 100,
				Stream:      &MockSubscribeStream{},
				Ctx:         context.Background(),
			},
			requestedOffset: 50,
			want:            true,
			reason:          "Requested offset 50 < current 100",
		},
		{
			name: "Large gap ahead - needs restart",
			session: &BrokerSubscriberSession{
				Topic:       "test-topic",
				Partition:   0,
				StartOffset: 100,
				Stream:      &MockSubscribeStream{},
				Ctx:         context.Background(),
			},
			requestedOffset: 2000,
			want:            true,
			reason:          "Gap of 1900 is > 1000",
		},
		{
			name: "Small gap ahead - no restart needed",
			session: &BrokerSubscriberSession{
				Topic:       "test-topic",
				Partition:   0,
				StartOffset: 100,
				Stream:      &MockSubscribeStream{},
				Ctx:         context.Background(),
			},
			requestedOffset: 150,
			want:            false,
			reason:          "Gap of 50 is < 1000",
		},
		{
			name: "Exact match - no restart needed",
			session: &BrokerSubscriberSession{
				Topic:       "test-topic",
				Partition:   0,
				StartOffset: 100,
				Stream:      &MockSubscribeStream{},
				Ctx:         context.Background(),
			},
			requestedOffset: 100,
			want:            false,
			reason:          "Exact match with current offset",
		},
		{
			name: "Context is nil - needs restart",
			session: &BrokerSubscriberSession{
				Topic:       "test-topic",
				Partition:   0,
				StartOffset: 100,
				Stream:      &MockSubscribeStream{},
				Ctx:         nil,
			},
			requestedOffset: 100,
			want:            true,
			reason:          "Context is nil",
		},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			got := bc.NeedsRestart(tt.session, tt.requestedOffset)
			if got != tt.want {
				t.Errorf("NeedsRestart() = %v, want %v (reason: %s)", got, tt.want, tt.reason)
			}
		})
	}
}

// TestNeedsRestart_CacheLogic tests cache-based restart decisions
func TestNeedsRestart_CacheLogic(t *testing.T) {
	bc := &BrokerClient{}

	// Create session with cache containing offsets 100-109
	session := &BrokerSubscriberSession{
		Topic:       "test-topic",
		Partition:   0,
		StartOffset: 110,
		Stream:      &MockSubscribeStream{},
		Ctx:         context.Background(),
		consumedRecords: []*SeaweedRecord{
			{Offset: 100}, {Offset: 101}, {Offset: 102}, {Offset: 103}, {Offset: 104},
			{Offset: 105}, {Offset: 106}, {Offset: 107}, {Offset: 108}, {Offset: 109},
		},
	}

	testCases := []struct {
		offset int64
		want   bool
		desc   string
	}{
		{100, false, "First offset in cache"},
		{105, false, "Middle offset in cache"},
		{109, false, "Last offset in cache"},
		{99, true, "Before cache start"},
		{110, false, "Current position"},
		{111, false, "One ahead"},
		{1200, true, "Large gap > 1000"},
	}

	for _, tc := range testCases {
		t.Run(tc.desc, func(t *testing.T) {
			got := bc.NeedsRestart(session, tc.offset)
			if got != tc.want {
				t.Errorf("NeedsRestart(offset=%d) = %v, want %v (%s)", tc.offset, got, tc.want, tc.desc)
			}
		})
	}
}

// TestNeedsRestart_EmptyCache tests behavior with empty cache
func TestNeedsRestart_EmptyCache(t *testing.T) {
	bc := &BrokerClient{}

	session := &BrokerSubscriberSession{
		Topic:           "test-topic",
		Partition:       0,
		StartOffset:     100,
		Stream:          &MockSubscribeStream{},
		Ctx:             context.Background(),
		consumedRecords: nil, // Empty cache
	}

	tests := []struct {
		offset int64
		want   bool
		desc   string
	}{
		{50, true, "Before current"},
		{100, false, "At current"},
		{150, false, "Small gap ahead"},
		{1200, true, "Large gap ahead"},
	}

	for _, tt := range tests {
		t.Run(tt.desc, func(t *testing.T) {
			got := bc.NeedsRestart(session, tt.offset)
			if got != tt.want {
				t.Errorf("NeedsRestart(offset=%d) = %v, want %v (%s)", tt.offset, got, tt.want, tt.desc)
			}
		})
	}
}

// TestNeedsRestart_ThreadSafety tests concurrent access
func TestNeedsRestart_ThreadSafety(t *testing.T) {
	bc := &BrokerClient{}

	session := &BrokerSubscriberSession{
		Topic:       "test-topic",
		Partition:   0,
		StartOffset: 100,
		Stream:      &MockSubscribeStream{},
		Ctx:         context.Background(),
	}

	// Run many concurrent checks
	done := make(chan bool)
	for i := 0; i < 100; i++ {
		go func(offset int64) {
			bc.NeedsRestart(session, offset)
			done <- true
		}(int64(i))
	}

	// Wait for all to complete
	for i := 0; i < 100; i++ {
		<-done
	}

	// Test passes if no panic/race condition
}

// TestRestartSubscriber_StateManagement tests session state management
func TestRestartSubscriber_StateManagement(t *testing.T) {
	oldStream := &MockSubscribeStream{}
	oldCtx, oldCancel := context.WithCancel(context.Background())

	session := &BrokerSubscriberSession{
		Topic:       "test-topic",
		Partition:   0,
		StartOffset: 100,
		Stream:      oldStream,
		Ctx:         oldCtx,
		Cancel:      oldCancel,
		consumedRecords: []*SeaweedRecord{
			{Offset: 100, Key: []byte("key100"), Value: []byte("value100")},
			{Offset: 101, Key: []byte("key101"), Value: []byte("value101")},
			{Offset: 102, Key: []byte("key102"), Value: []byte("value102")},
		},
		nextOffsetToRead: 103,
	}

	// Verify initial state
	if len(session.consumedRecords) != 3 {
		t.Errorf("Initial cache size = %d, want 3", len(session.consumedRecords))
	}
	if session.nextOffsetToRead != 103 {
		t.Errorf("Initial nextOffsetToRead = %d, want 103", session.nextOffsetToRead)
	}
	if session.StartOffset != 100 {
		t.Errorf("Initial StartOffset = %d, want 100", session.StartOffset)
	}

	// Note: Full RestartSubscriber testing requires gRPC mocking
	// These tests verify the core state management and NeedsRestart logic
}

// BenchmarkNeedsRestart_CacheHit benchmarks cache hit performance
func BenchmarkNeedsRestart_CacheHit(b *testing.B) {
	bc := &BrokerClient{}

	session := &BrokerSubscriberSession{
		Topic:           "test-topic",
		Partition:       0,
		StartOffset:     1000,
		Stream:          &MockSubscribeStream{},
		Ctx:             context.Background(),
		consumedRecords: make([]*SeaweedRecord, 100),
	}

	for i := 0; i < 100; i++ {
		session.consumedRecords[i] = &SeaweedRecord{Offset: int64(i)}
	}

	b.ResetTimer()
	for i := 0; i < b.N; i++ {
		bc.NeedsRestart(session, 50) // Hit cache
	}
}

// BenchmarkNeedsRestart_CacheMiss benchmarks cache miss performance
func BenchmarkNeedsRestart_CacheMiss(b *testing.B) {
	bc := &BrokerClient{}

	session := &BrokerSubscriberSession{
		Topic:           "test-topic",
		Partition:       0,
		StartOffset:     1000,
		Stream:          &MockSubscribeStream{},
		Ctx:             context.Background(),
		consumedRecords: make([]*SeaweedRecord, 100),
	}

	for i := 0; i < 100; i++ {
		session.consumedRecords[i] = &SeaweedRecord{Offset: int64(i)}
	}

	b.ResetTimer()
	for i := 0; i < b.N; i++ {
		bc.NeedsRestart(session, 500) // Miss cache (within gap threshold)
	}
}