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

     1  //go:build integration
     2  // +build integration
     3  
     4  package integration
     5  
     6  import (
     7  	"os"
     8  	"testing"
     9  	"time"
    10  
    11  	"github.com/ydb-platform/ydb-go-sdk/v3"
    12  	"github.com/ydb-platform/ydb-go-sdk/v3/balancers"
    13  	"github.com/ydb-platform/ydb-go-sdk/v3/config"
    14  	"github.com/ydb-platform/ydb-go-sdk/v3/coordination"
    15  	"github.com/ydb-platform/ydb-go-sdk/v3/internal/xtest"
    16  	"github.com/ydb-platform/ydb-go-sdk/v3/log"
    17  	"github.com/ydb-platform/ydb-go-sdk/v3/ratelimiter"
    18  	"github.com/ydb-platform/ydb-go-sdk/v3/trace"
    19  )
    20  
    21  func TestRatelimiter(sourceTest *testing.T) {
    22  	t := xtest.MakeSyncedTest(sourceTest)
    23  	const (
    24  		testCoordinationNodePath = "/local/ratelimiter_test"
    25  		testResource             = "test_resource"
    26  	)
    27  
    28  	ctx := xtest.Context(t)
    29  
    30  	db, err := ydb.Open(ctx,
    31  		os.Getenv("YDB_CONNECTION_STRING"),
    32  		ydb.WithAccessTokenCredentials(
    33  			os.Getenv("YDB_ACCESS_TOKEN_CREDENTIALS"),
    34  		),
    35  		ydb.With(
    36  			config.WithOperationTimeout(time.Second*2),
    37  			config.WithOperationCancelAfter(time.Second*2),
    38  		),
    39  		ydb.WithBalancer(balancers.SingleConn()),
    40  		ydb.WithLogger(
    41  			newLoggerWithMinLevel(t, log.WARN),
    42  			trace.MatchDetails(`ydb\.(driver|discovery|retry|ratelimiter|coordination).*`),
    43  		),
    44  	)
    45  	if err != nil {
    46  		t.Fatal(err)
    47  	}
    48  	defer func() {
    49  		// cleanup connection
    50  		if e := db.Close(ctx); e != nil {
    51  			t.Fatalf("db close failed: %+v", e)
    52  		}
    53  	}()
    54  	// drop node
    55  	err = db.Coordination().DropNode(ctx, testCoordinationNodePath)
    56  	if err != nil {
    57  		if !ydb.IsOperationErrorSchemeError(err) {
    58  			t.Fatal(err)
    59  		}
    60  	}
    61  	// create node
    62  	err = db.Coordination().CreateNode(ctx, testCoordinationNodePath, coordination.NodeConfig{
    63  		Path:                     "",
    64  		SelfCheckPeriodMillis:    1000,
    65  		SessionGracePeriodMillis: 1000,
    66  		ReadConsistencyMode:      coordination.ConsistencyModeRelaxed,
    67  		AttachConsistencyMode:    coordination.ConsistencyModeRelaxed,
    68  		RatelimiterCountersMode:  coordination.RatelimiterCountersModeDetailed,
    69  	})
    70  	if err != nil {
    71  		t.Fatal(err)
    72  	}
    73  	defer func() {
    74  		// cleanup node
    75  		err = db.Coordination().DropNode(ctx, testCoordinationNodePath)
    76  		if err != nil {
    77  			t.Fatal(err)
    78  		}
    79  	}()
    80  	// create resource
    81  	err = db.Ratelimiter().CreateResource(ctx, testCoordinationNodePath, ratelimiter.Resource{
    82  		ResourcePath: testResource,
    83  		HierarchicalDrr: ratelimiter.HierarchicalDrrSettings{
    84  			MaxUnitsPerSecond:       1,
    85  			MaxBurstSizeCoefficient: 2,
    86  		},
    87  	})
    88  	if err != nil {
    89  		t.Fatal(err)
    90  	}
    91  	defer func() {
    92  		// cleanup resource
    93  		err = db.Ratelimiter().DropResource(ctx, testCoordinationNodePath, testResource)
    94  		if err != nil {
    95  			t.Fatalf("Cannot drop resource: %v", err)
    96  		}
    97  	}()
    98  	// describe resource
    99  	described, err := db.Ratelimiter().DescribeResource(ctx, testCoordinationNodePath, testResource)
   100  	if err != nil {
   101  		t.Fatal(err)
   102  	}
   103  	if described == nil ||
   104  		described.ResourcePath != testResource ||
   105  		described.HierarchicalDrr.MaxUnitsPerSecond != 1.0 ||
   106  		described.HierarchicalDrr.MaxBurstSizeCoefficient != 2.0 {
   107  		t.Fatalf("Resource invalid: %+v", described)
   108  	}
   109  	// alter resource
   110  	err = db.Ratelimiter().AlterResource(ctx, testCoordinationNodePath, ratelimiter.Resource{
   111  		ResourcePath: testResource,
   112  		HierarchicalDrr: ratelimiter.HierarchicalDrrSettings{
   113  			MaxUnitsPerSecond:       3,
   114  			MaxBurstSizeCoefficient: 4,
   115  		},
   116  	})
   117  	if err != nil {
   118  		t.Fatal(err)
   119  	}
   120  	// check altered resource
   121  	described, err = db.Ratelimiter().DescribeResource(ctx, testCoordinationNodePath, testResource)
   122  	if err != nil {
   123  		t.Fatal(err)
   124  	}
   125  	if described == nil ||
   126  		described.ResourcePath != testResource ||
   127  		described.HierarchicalDrr.MaxUnitsPerSecond != 3.0 ||
   128  		described.HierarchicalDrr.MaxBurstSizeCoefficient != 4.0 {
   129  		t.Fatal("Resource invalid")
   130  	}
   131  	// list resource
   132  	list, err := db.Ratelimiter().ListResource(
   133  		ctx,
   134  		testCoordinationNodePath,
   135  		testResource,
   136  		true,
   137  	)
   138  	if err != nil {
   139  		t.Fatal(err)
   140  	}
   141  	if len(list) != 1 || list[0] != testResource {
   142  		t.Fatal("ListResource error")
   143  	}
   144  	// acquire resource amount 1
   145  	time.Sleep(time.Second) // for accumulate
   146  	err = db.Ratelimiter().AcquireResource(
   147  		ctx,
   148  		testCoordinationNodePath,
   149  		testResource,
   150  		1,
   151  		ratelimiter.WithAcquire(),
   152  	)
   153  	if err != nil {
   154  		t.Fatal(err)
   155  	}
   156  	// report resource amount 10000
   157  	err = db.Ratelimiter().AcquireResource(
   158  		ctx,
   159  		testCoordinationNodePath,
   160  		testResource,
   161  		10000,
   162  		ratelimiter.WithReport(),
   163  	)
   164  	if err != nil {
   165  		t.Fatal(err)
   166  	}
   167  	// acquire resource amount 10000
   168  	err = db.Ratelimiter().AcquireResource(
   169  		ctx,
   170  		testCoordinationNodePath,
   171  		testResource,
   172  		10000,
   173  		ratelimiter.WithAcquire(),
   174  	)
   175  	if err == nil {
   176  		t.Fatal("Resource must not be acquired")
   177  	}
   178  }