github.com/matrixorigin/matrixone@v1.2.0/pkg/proxy/scaling_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  	"fmt"
    20  	"net"
    21  	"os"
    22  	"testing"
    23  	"time"
    24  
    25  	"github.com/lni/goutils/leaktest"
    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  func TestDoScalingIn(t *testing.T) {
    32  	defer leaktest.AfterTest(t)()
    33  
    34  	var err error
    35  	tp := newTestProxyHandler(t)
    36  	defer tp.closeFn()
    37  
    38  	temp := os.TempDir()
    39  	// Construct backend CN servers.
    40  	addr1 := fmt.Sprintf("%s/%d.sock", temp, time.Now().Nanosecond())
    41  	require.NoError(t, os.RemoveAll(addr1))
    42  	tenantName := "t1"
    43  	cn11 := testMakeCNServer("cn11", addr1, 0, "",
    44  		newLabelInfo(Tenant(tenantName), map[string]string{}),
    45  	)
    46  	li := labelInfo{
    47  		Tenant: Tenant(tenantName),
    48  	}
    49  	cn11.hash, err = li.getHash()
    50  	require.NoError(t, err)
    51  	tp.hc.updateCN("cn11", cn11.addr, map[string]metadata.LabelList{
    52  		tenantLabelKey: {Labels: []string{tenantName}},
    53  	})
    54  	stopFn11 := startTestCNServer(t, tp.ctx, addr1, nil)
    55  	defer func() {
    56  		require.NoError(t, stopFn11())
    57  	}()
    58  
    59  	addr2 := fmt.Sprintf("%s/%d.sock", temp, time.Now().Nanosecond())
    60  	require.NoError(t, os.RemoveAll(addr2))
    61  	cn12 := testMakeCNServer("cn12", addr2, 0, "",
    62  		newLabelInfo(Tenant(tenantName), map[string]string{}),
    63  	)
    64  	cn12.hash, err = li.getHash()
    65  	require.NoError(t, err)
    66  	tp.hc.updateCN("cn12", cn12.addr, map[string]metadata.LabelList{
    67  		tenantLabelKey: {Labels: []string{tenantName}},
    68  	})
    69  	stopFn12 := startTestCNServer(t, tp.ctx, addr2, nil)
    70  	defer func() {
    71  		require.NoError(t, stopFn12())
    72  	}()
    73  	tp.mc.ForceRefresh(true)
    74  
    75  	ctx, cancel := context.WithTimeout(tp.ctx, 10*time.Second)
    76  	defer cancel()
    77  
    78  	ci := clientInfo{
    79  		labelInfo: li,
    80  		username:  "test",
    81  		originIP:  net.ParseIP("127.0.0.1"),
    82  	}
    83  	cleanup1 := testStartNClients(t, tp, ci, cn11, 2)
    84  	defer cleanup1()
    85  
    86  	cleanup2 := testStartNClients(t, tp, ci, cn12, 2)
    87  	defer cleanup2()
    88  
    89  	err = tp.hc.updateCNWorkState(ctx, logpb.CNWorkState{
    90  		UUID:  cn11.uuid,
    91  		State: metadata.WorkState_Draining,
    92  	})
    93  	require.NoError(t, err)
    94  	tp.mc.ForceRefresh(true)
    95  
    96  	tick := time.NewTicker(time.Millisecond * 200)
    97  	for {
    98  		select {
    99  		case <-ctx.Done():
   100  			require.Fail(t, "scaling in failed")
   101  		case <-tick.C:
   102  			tunnels1 := tp.re.connManager.getTunnelsByCNID(cn11.uuid)
   103  			tunnels2 := tp.re.connManager.getTunnelsByCNID(cn12.uuid)
   104  			tp.re.connManager.Lock()
   105  			if len(tunnels1) == 0 && len(tunnels2) == 4 {
   106  				tp.re.connManager.Unlock()
   107  				return
   108  			}
   109  			tp.re.connManager.Unlock()
   110  		}
   111  	}
   112  }