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 }