github.com/matrixorigin/matrixone@v1.2.0/pkg/incrservice/allocator_test.go (about)

     1  // Copyright 2023 Matrix Origin
     2  //
     3  // Licensed under the Apache License, Version 2.0 (the "License");
     4  // you may not use this file except in compliance with the License.
     5  // You may obtain a copy of the License at
     6  //
     7  //      http://www.apache.org/licenses/LICENSE-2.0
     8  //
     9  // Unless required by applicable law or agreed to in writing, software
    10  // distributed under the License is distributed on an "AS IS" BASIS,
    11  // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    12  // See the License for the specific language governing permissions and
    13  // limitations under the License.
    14  
    15  package incrservice
    16  
    17  import (
    18  	"context"
    19  	"sync"
    20  	"testing"
    21  
    22  	"github.com/lni/goutils/leaktest"
    23  	"github.com/matrixorigin/matrixone/pkg/catalog"
    24  	"github.com/matrixorigin/matrixone/pkg/common/runtime"
    25  	"github.com/matrixorigin/matrixone/pkg/defines"
    26  	"github.com/stretchr/testify/require"
    27  )
    28  
    29  func TestAlloc(t *testing.T) {
    30  	runAllocatorTests(
    31  		t,
    32  		func(a valueAllocator) {
    33  			ctx, cancel := context.WithCancel(defines.AttachAccountId(context.Background(), catalog.System_Account))
    34  			defer cancel()
    35  
    36  			cases := []struct {
    37  				key        string
    38  				count      int
    39  				step       int
    40  				expectFrom uint64
    41  				expectTo   uint64
    42  			}{
    43  				{
    44  					key:        "k1",
    45  					count:      1,
    46  					step:       1,
    47  					expectFrom: 1,
    48  					expectTo:   2,
    49  				},
    50  				{
    51  					key:        "k1",
    52  					count:      1,
    53  					step:       1,
    54  					expectFrom: 2,
    55  					expectTo:   3,
    56  				},
    57  				{
    58  					key:        "k1",
    59  					count:      2,
    60  					step:       1,
    61  					expectFrom: 3,
    62  					expectTo:   5,
    63  				},
    64  				{
    65  					key:        "k2",
    66  					count:      2,
    67  					step:       3,
    68  					expectFrom: 3,
    69  					expectTo:   9,
    70  				},
    71  				{
    72  					key:        "k2",
    73  					count:      2,
    74  					step:       3,
    75  					expectFrom: 9,
    76  					expectTo:   15,
    77  				},
    78  			}
    79  
    80  			for _, c := range cases {
    81  				require.NoError(t,
    82  					a.(*allocator).store.Create(
    83  						ctx,
    84  						0,
    85  						[]AutoColumn{
    86  							{
    87  								ColName: c.key,
    88  								Offset:  0,
    89  								Step:    uint64(c.step),
    90  							},
    91  						},
    92  						nil))
    93  				from, to, err := a.allocate(ctx, 0, c.key, c.count, nil)
    94  				require.NoError(t, err)
    95  				require.Equal(t, c.expectFrom, from)
    96  				require.Equal(t, c.expectTo, to)
    97  			}
    98  		})
    99  }
   100  
   101  func TestAsyncAlloc(t *testing.T) {
   102  	runAllocatorTests(
   103  		t,
   104  		func(a valueAllocator) {
   105  			ctx, cancel := context.WithCancel(defines.AttachAccountId(context.Background(), catalog.System_Account))
   106  			defer cancel()
   107  
   108  			cases := []struct {
   109  				key        string
   110  				count      int
   111  				step       int
   112  				expectFrom uint64
   113  				expectTo   uint64
   114  			}{
   115  				{
   116  					key:        "k1",
   117  					count:      1,
   118  					step:       1,
   119  					expectFrom: 1,
   120  					expectTo:   2,
   121  				},
   122  				{
   123  					key:        "k1",
   124  					count:      1,
   125  					step:       1,
   126  					expectFrom: 2,
   127  					expectTo:   3,
   128  				},
   129  				{
   130  					key:        "k1",
   131  					count:      2,
   132  					step:       1,
   133  					expectFrom: 3,
   134  					expectTo:   5,
   135  				},
   136  				{
   137  					key:        "k2",
   138  					count:      2,
   139  					step:       3,
   140  					expectFrom: 3,
   141  					expectTo:   9,
   142  				},
   143  				{
   144  					key:        "k2",
   145  					count:      2,
   146  					step:       3,
   147  					expectFrom: 9,
   148  					expectTo:   15,
   149  				},
   150  			}
   151  
   152  			for _, c := range cases {
   153  				require.NoError(t,
   154  					a.(*allocator).store.Create(
   155  						ctx,
   156  						0,
   157  						[]AutoColumn{
   158  							{
   159  								ColName: c.key,
   160  								Offset:  0,
   161  								Step:    uint64(c.step),
   162  							},
   163  						},
   164  						nil))
   165  				var wg sync.WaitGroup
   166  				wg.Add(1)
   167  				a.asyncAllocate(
   168  					ctx,
   169  					0,
   170  					c.key,
   171  					c.count,
   172  					nil,
   173  					func(from, to uint64, err error) {
   174  						defer wg.Done()
   175  						require.NoError(t, err)
   176  						require.Equal(t, c.expectFrom, from)
   177  						require.Equal(t, c.expectTo, to)
   178  					})
   179  				wg.Wait()
   180  			}
   181  		})
   182  }
   183  
   184  func runAllocatorTests(
   185  	t *testing.T,
   186  	fn func(valueAllocator)) {
   187  	defer leaktest.AfterTest(t)()
   188  	runtime.SetupProcessLevelRuntime(runtime.DefaultRuntime())
   189  	a := newValueAllocator(nil)
   190  	defer a.close()
   191  	fn(a)
   192  }