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