github.com/ydb-platform/ydb-go-sdk/v3@v3.57.0/tests/integration/topic_partitions_balanced_test.go (about)

     1  //go:build integration
     2  // +build integration
     3  
     4  package integration
     5  
     6  import (
     7  	"context"
     8  	"sync"
     9  	"sync/atomic"
    10  	"testing"
    11  	"time"
    12  
    13  	"github.com/stretchr/testify/require"
    14  
    15  	"github.com/ydb-platform/ydb-go-sdk/v3"
    16  	"github.com/ydb-platform/ydb-go-sdk/v3/internal/empty"
    17  	"github.com/ydb-platform/ydb-go-sdk/v3/internal/xtest"
    18  	"github.com/ydb-platform/ydb-go-sdk/v3/topic/topicoptions"
    19  	"github.com/ydb-platform/ydb-go-sdk/v3/topic/topictypes"
    20  	"github.com/ydb-platform/ydb-go-sdk/v3/trace"
    21  )
    22  
    23  func TestTopicPartitionsBalanced(t *testing.T) {
    24  	ctx := xtest.Context(t)
    25  	db := connect(t)
    26  	topicPath := db.Name() + "/topic-" + t.Name()
    27  
    28  	err := db.Topic().Drop(ctx, topicPath)
    29  	if err != nil {
    30  		require.True(t, ydb.IsOperationErrorSchemeError(err))
    31  	}
    32  
    33  	consumer := "test-consumer-" + t.Name()
    34  	err = db.Topic().Create(ctx, topicPath,
    35  		topicoptions.CreateWithMinActivePartitions(2),
    36  		topicoptions.CreateWithPartitionCountLimit(2),
    37  		topicoptions.CreateWithConsumer(topictypes.Consumer{Name: consumer}),
    38  	)
    39  	require.NoError(t, err)
    40  
    41  	var connectedPartitions atomic.Int64
    42  	var handled atomic.Int64
    43  
    44  	var sessionsMutex sync.Mutex
    45  	sessions := map[int64]bool{}
    46  
    47  	tracer := trace.Topic{
    48  		OnReaderPartitionReadStartResponse: func(startInfo trace.TopicReaderPartitionReadStartResponseStartInfo) func(doneInfo trace.TopicReaderPartitionReadStartResponseDoneInfo) { //nolint:lll
    49  			handled.Store(1)
    50  
    51  			connectedPartitions.Add(1)
    52  			return nil
    53  		},
    54  		OnReaderPartitionReadStopResponse: func(startInfo trace.TopicReaderPartitionReadStopResponseStartInfo) func(doneInfo trace.TopicReaderPartitionReadStopResponseDoneInfo) { //nolint:lll
    55  			handled.Store(1)
    56  
    57  			sessionsMutex.Lock()
    58  			defer sessionsMutex.Unlock()
    59  			if sessions[startInfo.PartitionSessionID] {
    60  				return nil
    61  			}
    62  			sessions[startInfo.PartitionSessionID] = true
    63  
    64  			connectedPartitions.Add(-1)
    65  			return nil
    66  		},
    67  	}
    68  	firstReader, err := db.Topic().StartReader(consumer, topicoptions.ReadTopic(topicPath),
    69  		topicoptions.WithReaderTrace(tracer),
    70  	)
    71  	require.NoError(t, err)
    72  
    73  	readCtx, firstReaderStopRead := context.WithCancel(ctx)
    74  	firstReaderReadStopped := make(empty.Chan)
    75  	go func() {
    76  		defer close(firstReaderReadStopped)
    77  
    78  		for {
    79  			if readCtx.Err() != nil {
    80  				return
    81  			}
    82  			_, err = firstReader.ReadMessage(readCtx)
    83  			if readCtx.Err() == nil {
    84  				require.NoError(t, err)
    85  			}
    86  		}
    87  	}()
    88  
    89  	xtest.SpinWaitConditionWithTimeout(t, nil, time.Second, func() bool {
    90  		return connectedPartitions.Load() == 2
    91  	})
    92  
    93  	readerSecond, err := db.Topic().StartReader(consumer, topicoptions.ReadTopic(topicPath))
    94  	require.NoError(t, err)
    95  
    96  	xtest.SpinWaitConditionWithTimeout(t, nil, time.Second, func() bool {
    97  		return connectedPartitions.Load() == 1
    98  	})
    99  
   100  	require.NoError(t, readerSecond.Close(ctx))
   101  
   102  	xtest.SpinWaitConditionWithTimeout(t, nil, time.Second, func() bool {
   103  		return connectedPartitions.Load() == 2
   104  	})
   105  
   106  	firstReaderStopRead()
   107  	xtest.WaitChannelClosed(t, firstReaderReadStopped)
   108  	require.NoError(t, firstReader.Close(ctx))
   109  }