github.com/matrixorigin/matrixone@v1.2.0/pkg/proxy/label_info_test.go (about)

     1  // Copyright 2021 - 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 proxy
    16  
    17  import (
    18  	"context"
    19  	"sync"
    20  	"testing"
    21  	"time"
    22  
    23  	"github.com/lni/goutils/leaktest"
    24  	"github.com/matrixorigin/matrixone/pkg/clusterservice"
    25  	"github.com/matrixorigin/matrixone/pkg/common/runtime"
    26  	logpb "github.com/matrixorigin/matrixone/pkg/pb/logservice"
    27  	"github.com/matrixorigin/matrixone/pkg/pb/metadata"
    28  	"github.com/stretchr/testify/require"
    29  )
    30  
    31  type mockHAKeeperClient struct {
    32  	sync.RWMutex
    33  	value logpb.ClusterDetails
    34  }
    35  
    36  func (c *mockHAKeeperClient) updateCN(uuid string, addr string, labels map[string]metadata.LabelList) {
    37  	c.Lock()
    38  	defer c.Unlock()
    39  	var cs *logpb.CNStore
    40  	for i := range c.value.CNStores {
    41  		if c.value.CNStores[i].UUID == uuid {
    42  			cs = &c.value.CNStores[i]
    43  			break
    44  		}
    45  	}
    46  	if cs != nil {
    47  		cs.Labels = labels
    48  		cs.SQLAddress = addr
    49  		return
    50  	}
    51  	cs = &logpb.CNStore{
    52  		UUID:       uuid,
    53  		SQLAddress: addr,
    54  		Labels:     labels,
    55  		WorkState:  metadata.WorkState_Working,
    56  	}
    57  	c.value.CNStores = append(c.value.CNStores, *cs)
    58  }
    59  
    60  func (c *mockHAKeeperClient) Close() error                                 { return nil }
    61  func (c *mockHAKeeperClient) AllocateID(_ context.Context) (uint64, error) { return 0, nil }
    62  func (c *mockHAKeeperClient) AllocateIDByKey(_ context.Context, _ string) (uint64, error) {
    63  	return uint64(nextClientConnID()), nil
    64  }
    65  func (c *mockHAKeeperClient) AllocateIDByKeyWithBatch(_ context.Context, _ string, _ uint64) (uint64, error) {
    66  	return uint64(nextClientConnID()), nil
    67  }
    68  func (c *mockHAKeeperClient) GetClusterDetails(_ context.Context) (logpb.ClusterDetails, error) {
    69  	c.RLock()
    70  	defer c.RUnlock()
    71  
    72  	// return a deepcopy of inner state to avoid data race
    73  	r := logpb.ClusterDetails{}
    74  	bytes, err := c.value.Marshal()
    75  	if err != nil {
    76  		return r, err
    77  	}
    78  	if err := r.Unmarshal(bytes); err != nil {
    79  		return r, err
    80  	}
    81  	return r, nil
    82  }
    83  
    84  func (c *mockHAKeeperClient) GetClusterState(_ context.Context) (logpb.CheckerState, error) {
    85  	return logpb.CheckerState{
    86  		TaskTableUser: logpb.TaskTableUser{
    87  			Username: "u1",
    88  			Password: "p1",
    89  		},
    90  	}, nil
    91  }
    92  
    93  func (c *mockHAKeeperClient) GetCNState(_ context.Context) (logpb.CNState, error) {
    94  	return logpb.CNState{}, nil
    95  }
    96  
    97  func (c *mockHAKeeperClient) UpdateCNLabel(_ context.Context, _ logpb.CNStoreLabel) error {
    98  	return nil
    99  }
   100  
   101  func (c *mockHAKeeperClient) UpdateCNWorkState(_ context.Context, _ logpb.CNWorkState) error {
   102  	return nil
   103  }
   104  
   105  func (c *mockHAKeeperClient) PatchCNStore(_ context.Context, _ logpb.CNStateLabel) error {
   106  	return nil
   107  }
   108  
   109  func (c *mockHAKeeperClient) DeleteCNStore(_ context.Context, _ logpb.DeleteCNStore) error {
   110  	return nil
   111  }
   112  
   113  func (c *mockHAKeeperClient) SendProxyHeartbeat(_ context.Context, _ logpb.ProxyHeartbeat) (logpb.CommandBatch, error) {
   114  	return logpb.CommandBatch{}, nil
   115  }
   116  
   117  func (c *mockHAKeeperClient) updateCNWorkState(_ context.Context, state logpb.CNWorkState) error {
   118  	c.Lock()
   119  	defer c.Unlock()
   120  	for i := range c.value.CNStores {
   121  		if c.value.CNStores[i].UUID == state.UUID {
   122  			c.value.CNStores[i].WorkState = state.State
   123  		}
   124  	}
   125  	return nil
   126  }
   127  
   128  func TestLabelInfoReserved(t *testing.T) {
   129  	defer leaktest.AfterTest(t)()
   130  
   131  	tenant := Tenant("t1")
   132  	labels := map[string]string{
   133  		"k1":           "v1",
   134  		"k2":           "v2",
   135  		"os_user":      "u1",
   136  		"os_sudouser":  "u2",
   137  		"program_name": "p1",
   138  		"_r1":          "_v1",
   139  		"_r2":          "_v2",
   140  	}
   141  	info := newLabelInfo(tenant, labels)
   142  	require.Equal(t, 2, len(info.Labels))
   143  	_, ok := info.Labels["os_user"]
   144  	require.False(t, ok)
   145  	_, ok = info.Labels["os_sudouser"]
   146  	require.False(t, ok)
   147  	_, ok = info.Labels["program_name"]
   148  	require.False(t, ok)
   149  	_, ok = info.Labels["_r1"]
   150  	require.False(t, ok)
   151  	_, ok = info.Labels["_r2"]
   152  	require.False(t, ok)
   153  	v1, ok := info.Labels["k1"]
   154  	require.True(t, ok)
   155  	require.Equal(t, "v1", v1)
   156  	v2, ok := info.Labels["k2"]
   157  	require.True(t, ok)
   158  	require.Equal(t, "v2", v2)
   159  }
   160  
   161  func TestLabelInfoAll(t *testing.T) {
   162  	defer leaktest.AfterTest(t)()
   163  	tenant := Tenant("t1")
   164  	labels := map[string]string{
   165  		"k1": "v1",
   166  		"k2": "v2",
   167  	}
   168  	info := newLabelInfo(tenant, labels)
   169  
   170  	t.Run("all", func(t *testing.T) {
   171  		all := info.allLabels()
   172  		require.Equal(t, 3, len(all))
   173  
   174  		v0, ok := all[tenantLabelKey]
   175  		require.True(t, ok)
   176  		require.Equal(t, "t1", v0)
   177  
   178  		v1, ok := all["k1"]
   179  		require.True(t, ok)
   180  		require.Equal(t, "v1", v1)
   181  
   182  		v2, ok := all["k2"]
   183  		require.True(t, ok)
   184  		require.Equal(t, "v2", v2)
   185  
   186  		_, ok = all["no"]
   187  		require.False(t, ok)
   188  	})
   189  
   190  	t.Run("common", func(t *testing.T) {
   191  		common := info.commonLabels()
   192  		require.Equal(t, 2, len(common))
   193  		v1, ok := common["k1"]
   194  		require.True(t, ok)
   195  		require.Equal(t, "v1", v1)
   196  		v2, ok := common["k2"]
   197  		require.True(t, ok)
   198  		require.Equal(t, "v2", v2)
   199  	})
   200  
   201  	t.Run("tenant", func(t *testing.T) {
   202  		tenantLabel := info.tenantLabel()
   203  		require.Equal(t, 1, len(tenantLabel))
   204  		v1, ok := tenantLabel[tenantLabelKey]
   205  		require.True(t, ok)
   206  		require.Equal(t, "t1", v1)
   207  	})
   208  }
   209  
   210  func TestLabelSuperTenant(t *testing.T) {
   211  	defer leaktest.AfterTest(t)()
   212  	tenant := Tenant("t1")
   213  	labels := map[string]string{
   214  		"k1": "v1",
   215  		"k2": "v2",
   216  	}
   217  	info := newLabelInfo(tenant, labels)
   218  	require.Equal(t, false, info.isSuperTenant())
   219  
   220  	tenant = ""
   221  	info = newLabelInfo(tenant, labels)
   222  	require.Equal(t, true, info.isSuperTenant())
   223  
   224  	tenant = "sys"
   225  	info = newLabelInfo(tenant, labels)
   226  	require.Equal(t, true, info.isSuperTenant())
   227  }
   228  
   229  func TestLabelSelector(t *testing.T) {
   230  	defer leaktest.AfterTest(t)()
   231  	tenant := Tenant("t1")
   232  	labels := map[string]string{
   233  		tenantLabelKey: "t1",
   234  		"k1":           "v1",
   235  		"k2":           "v2",
   236  		"k3":           "v3",
   237  	}
   238  	info := newLabelInfo(tenant, labels)
   239  	se := info.genSelector(clusterservice.EQ_Globbing)
   240  
   241  	hc := &mockHAKeeperClient{}
   242  	cnLabels1 := map[string]metadata.LabelList{
   243  		tenantLabelKey: {
   244  			Labels: []string{"t1"},
   245  		},
   246  		"k1": {
   247  			Labels: []string{"v1", "multi1"},
   248  		},
   249  		"k2": {
   250  			Labels: []string{"v2", "multi2"},
   251  		},
   252  		"k3": {
   253  			Labels: []string{"v3", "multi3"},
   254  		},
   255  	}
   256  	cnLabels2 := map[string]metadata.LabelList{
   257  		tenantLabelKey: {
   258  			Labels: []string{"t1"},
   259  		},
   260  		"k1": {
   261  			Labels: []string{"no1", "no2"},
   262  		},
   263  		"k2": {
   264  			Labels: []string{"v2", "multi2"},
   265  		},
   266  		"k3": {
   267  			Labels: []string{"v3", "multi3"},
   268  		},
   269  	}
   270  	hc.updateCN("cn1", "", cnLabels1)
   271  	hc.updateCN("cn2", "", cnLabels2)
   272  
   273  	runtime.SetupProcessLevelRuntime(runtime.DefaultRuntime())
   274  	mc := clusterservice.NewMOCluster(hc, 3*time.Second)
   275  	defer func() { mc.Close() }()
   276  	mc.ForceRefresh(true)
   277  
   278  	var servers []CNServer
   279  	mc.GetCNService(se, func(s metadata.CNService) bool {
   280  		servers = append(servers, CNServer{
   281  			uuid: s.ServiceID,
   282  		})
   283  		return true
   284  	})
   285  	require.Equal(t, 1, len(servers))
   286  }
   287  
   288  func TestLabelHash(t *testing.T) {
   289  	defer leaktest.AfterTest(t)()
   290  	label1 := labelInfo{
   291  		Tenant: "t1",
   292  		Labels: map[string]string{
   293  			"k1": "v1",
   294  			"k2": "v2",
   295  		},
   296  	}
   297  	h1, err := label1.getHash()
   298  	require.NoError(t, err)
   299  
   300  	label2 := labelInfo{
   301  		Labels: map[string]string{
   302  			"k2": "v2",
   303  			"k1": "v1",
   304  		},
   305  		Tenant: "t1",
   306  	}
   307  	h2, err := label2.getHash()
   308  	require.NoError(t, err)
   309  
   310  	label3 := labelInfo{
   311  		Tenant: "t3",
   312  		Labels: map[string]string{
   313  			"k1": "v1",
   314  			"k2": "v2",
   315  		},
   316  	}
   317  	h3, err := label3.getHash()
   318  	require.NoError(t, err)
   319  
   320  	label4 := labelInfo{
   321  		Tenant: "t1",
   322  		Labels: map[string]string{
   323  			"k1": "v1",
   324  			"k2": "v3",
   325  		},
   326  	}
   327  	h4, err := label4.getHash()
   328  	require.NoError(t, err)
   329  
   330  	require.Equal(t, h1, h2)
   331  	require.NotEqual(t, h1, h3)
   332  	require.NotEqual(t, h1, h4)
   333  }