github.com/ydb-platform/ydb-go-sdk/v3@v3.89.2/internal/table/client_test.go (about)

     1  package table
     2  
     3  import (
     4  	"context"
     5  	"fmt"
     6  	"math/rand"
     7  	"sync"
     8  	"testing"
     9  	"time"
    10  
    11  	"github.com/ydb-platform/ydb-go-genproto/protos/Ydb_Table"
    12  	"google.golang.org/grpc"
    13  	"google.golang.org/protobuf/proto"
    14  	"google.golang.org/protobuf/types/known/emptypb"
    15  
    16  	"github.com/ydb-platform/ydb-go-sdk/v3/internal/table/config"
    17  	"github.com/ydb-platform/ydb-go-sdk/v3/internal/xcontext"
    18  	"github.com/ydb-platform/ydb-go-sdk/v3/internal/xerrors"
    19  	"github.com/ydb-platform/ydb-go-sdk/v3/internal/xsync"
    20  	"github.com/ydb-platform/ydb-go-sdk/v3/internal/xtest"
    21  	"github.com/ydb-platform/ydb-go-sdk/v3/table"
    22  	"github.com/ydb-platform/ydb-go-sdk/v3/testutil"
    23  )
    24  
    25  func TestRaceWgClosed(t *testing.T) {
    26  	defer func() {
    27  		if e := recover(); e != nil {
    28  			t.Fatal(e)
    29  		}
    30  	}()
    31  
    32  	var (
    33  		limit   = 100
    34  		start   = time.Now()
    35  		counter int
    36  	)
    37  
    38  	xtest.TestManyTimes(t, func(t testing.TB) {
    39  		counter++
    40  		defer func() {
    41  			if counter%1000 == 0 {
    42  				t.Logf("%0.1fs: %d times test passed", time.Since(start).Seconds(), counter)
    43  			}
    44  		}()
    45  		ctx, cancel := xcontext.WithTimeout(context.Background(),
    46  			//nolint:gosec
    47  			time.Duration(rand.Int31n(int32(100*time.Millisecond))),
    48  		)
    49  		defer cancel()
    50  
    51  		wg := sync.WaitGroup{}
    52  		p := New(ctx,
    53  			testutil.NewBalancer(testutil.WithInvokeHandlers(testutil.InvokeHandlers{
    54  				testutil.TableCreateSession: func(interface{}) (proto.Message, error) {
    55  					return &Ydb_Table.CreateSessionResult{
    56  						SessionId: testutil.SessionID(),
    57  					}, nil
    58  				},
    59  			})),
    60  			config.New(config.WithSizeLimit(limit)),
    61  		)
    62  		for j := 0; j < limit*10; j++ {
    63  			wg.Add(1)
    64  			go func() {
    65  				defer wg.Done()
    66  				for {
    67  					err := p.Do(ctx,
    68  						func(ctx context.Context, s table.Session) error {
    69  							return nil
    70  						},
    71  					)
    72  					if err != nil && xerrors.Is(err, errClosedClient) {
    73  						return
    74  					}
    75  				}
    76  			}()
    77  		}
    78  		_ = p.Close(context.Background())
    79  		wg.Wait()
    80  	}, xtest.StopAfter(27*time.Second))
    81  }
    82  
    83  var okHandler = func(interface{}) (proto.Message, error) {
    84  	return &emptypb.Empty{}, nil
    85  }
    86  
    87  var simpleCluster = testutil.NewBalancer(
    88  	testutil.WithInvokeHandlers(
    89  		testutil.InvokeHandlers{
    90  			testutil.TableExecuteDataQuery: func(interface{}) (proto.Message, error) {
    91  				return &Ydb_Table.ExecuteQueryResult{
    92  					TxMeta: &Ydb_Table.TransactionMeta{
    93  						Id: "",
    94  					},
    95  				}, nil
    96  			},
    97  			testutil.TableBeginTransaction: func(interface{}) (proto.Message, error) {
    98  				return &Ydb_Table.BeginTransactionResult{
    99  					TxMeta: &Ydb_Table.TransactionMeta{
   100  						Id: "",
   101  					},
   102  				}, nil
   103  			},
   104  			testutil.TableExplainDataQuery: func(interface{}) (proto.Message, error) {
   105  				return &Ydb_Table.ExecuteQueryResult{}, nil
   106  			},
   107  			testutil.TablePrepareDataQuery: func(interface{}) (proto.Message, error) {
   108  				return &Ydb_Table.PrepareQueryResult{}, nil
   109  			},
   110  			testutil.TableCreateSession: func(interface{}) (proto.Message, error) {
   111  				return &Ydb_Table.CreateSessionResult{
   112  					SessionId: testutil.SessionID(),
   113  				}, nil
   114  			},
   115  			testutil.TableDeleteSession: func(interface{}) (proto.Message, error) {
   116  				return &Ydb_Table.DeleteSessionResponse{}, nil
   117  			},
   118  			testutil.TableCommitTransaction: func(interface{}) (proto.Message, error) {
   119  				return &Ydb_Table.CommitTransactionResponse{}, nil
   120  			},
   121  			testutil.TableRollbackTransaction: func(interface{}) (proto.Message, error) {
   122  				return &Ydb_Table.RollbackTransactionResponse{}, nil
   123  			},
   124  			testutil.TableKeepAlive: func(interface{}) (proto.Message, error) {
   125  				return &Ydb_Table.KeepAliveResult{}, nil
   126  			},
   127  		},
   128  	),
   129  )
   130  
   131  func simpleSession(t testing.TB) *session {
   132  	s, err := newSession(context.Background(), simpleCluster, config.New())
   133  	if err != nil {
   134  		t.Fatalf("newSession unexpected error: %v", err)
   135  	}
   136  
   137  	return s
   138  }
   139  
   140  type StubBuilder struct {
   141  	OnCreateSession func(ctx context.Context) (*session, error)
   142  
   143  	cc    grpc.ClientConnInterface
   144  	Limit int
   145  	T     testing.TB
   146  
   147  	mu     xsync.Mutex
   148  	actual int
   149  }
   150  
   151  func (s *StubBuilder) createSession(ctx context.Context) (session *session, err error) {
   152  	defer s.mu.WithLock(func() {
   153  		if session != nil {
   154  			s.actual++
   155  		}
   156  	})
   157  
   158  	s.mu.WithLock(func() {
   159  		if s.Limit > 0 && s.actual == s.Limit {
   160  			err = fmt.Errorf("stub session: limit overflow")
   161  		}
   162  	})
   163  	if err != nil {
   164  		return nil, err
   165  	}
   166  
   167  	if f := s.OnCreateSession; f != nil {
   168  		return f(ctx)
   169  	}
   170  
   171  	return newSession(ctx, s.cc, config.New())
   172  }