vitess.io/vitess@v0.16.2/go/vt/vttablet/tabletserver/txthrottler/tx_throttler_test.go (about)

     1  /*
     2  Copyright 2019 The Vitess Authors.
     3  
     4  Licensed under the Apache License, Version 2.0 (the "License");
     5  you may not use this file except in compliance with the License.
     6  You may obtain a copy of the License at
     7  
     8      http://www.apache.org/licenses/LICENSE-2.0
     9  
    10  Unless required by applicable law or agreedto in writing, software
    11  distributed under the License is distributed on an "AS IS" BASIS,
    12  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    13  See the License for the specific language governing permissions and
    14  limitations under the License.
    15  */
    16  
    17  package txthrottler
    18  
    19  // Commands to generate the mocks for this test.
    20  //go:generate mockgen -destination mock_healthcheck_test.go -package txthrottler -mock_names "HealthCheck=MockHealthCheck" vitess.io/vitess/go/vt/discovery HealthCheck
    21  //go:generate mockgen -destination mock_throttler_test.go -package txthrottler vitess.io/vitess/go/vt/vttablet/tabletserver/txthrottler ThrottlerInterface
    22  //go:generate mockgen -destination mock_topology_watcher_test.go -package txthrottler vitess.io/vitess/go/vt/vttablet/tabletserver/txthrottler TopologyWatcherInterface
    23  
    24  import (
    25  	"testing"
    26  	"time"
    27  
    28  	"github.com/golang/mock/gomock"
    29  
    30  	"vitess.io/vitess/go/vt/discovery"
    31  	"vitess.io/vitess/go/vt/topo"
    32  	"vitess.io/vitess/go/vt/topo/memorytopo"
    33  	"vitess.io/vitess/go/vt/vttablet/tabletserver/tabletenv"
    34  
    35  	querypb "vitess.io/vitess/go/vt/proto/query"
    36  	topodatapb "vitess.io/vitess/go/vt/proto/topodata"
    37  )
    38  
    39  func TestDisabledThrottler(t *testing.T) {
    40  	config := tabletenv.NewDefaultConfig()
    41  	config.EnableTxThrottler = false
    42  	throttler := NewTxThrottler(config, nil)
    43  	throttler.InitDBConfig(&querypb.Target{
    44  		Keyspace: "keyspace",
    45  		Shard:    "shard",
    46  	})
    47  	if err := throttler.Open(); err != nil {
    48  		t.Fatalf("want: nil, got: %v", err)
    49  	}
    50  	if result := throttler.Throttle(); result != false {
    51  		t.Errorf("want: false, got: %v", result)
    52  	}
    53  	throttler.Close()
    54  }
    55  
    56  func TestEnabledThrottler(t *testing.T) {
    57  	mockCtrl := gomock.NewController(t)
    58  	defer mockCtrl.Finish()
    59  
    60  	defer resetTxThrottlerFactories()
    61  	ts := memorytopo.NewServer("cell1", "cell2")
    62  
    63  	mockHealthCheck := NewMockHealthCheck(mockCtrl)
    64  	hcCall1 := mockHealthCheck.EXPECT().Subscribe()
    65  	hcCall1.Do(func() {})
    66  	hcCall2 := mockHealthCheck.EXPECT().Close()
    67  	hcCall2.After(hcCall1)
    68  	healthCheckFactory = func(topoServer *topo.Server, cell string, cellsToWatch []string) discovery.HealthCheck {
    69  		return mockHealthCheck
    70  	}
    71  
    72  	topologyWatcherFactory = func(topoServer *topo.Server, hc discovery.HealthCheck, cell, keyspace, shard string, refreshInterval time.Duration, topoReadConcurrency int) TopologyWatcherInterface {
    73  		if ts != topoServer {
    74  			t.Errorf("want: %v, got: %v", ts, topoServer)
    75  		}
    76  		if cell != "cell1" && cell != "cell2" {
    77  			t.Errorf("want: cell1 or cell2, got: %v", cell)
    78  		}
    79  		if keyspace != "keyspace" {
    80  			t.Errorf("want: keyspace, got: %v", keyspace)
    81  		}
    82  		if shard != "shard" {
    83  			t.Errorf("want: shard, got: %v", shard)
    84  		}
    85  		result := NewMockTopologyWatcherInterface(mockCtrl)
    86  		result.EXPECT().Stop()
    87  		return result
    88  	}
    89  
    90  	mockThrottler := NewMockThrottlerInterface(mockCtrl)
    91  	throttlerFactory = func(name, unit string, threadCount int, maxRate, maxReplicationLag int64) (ThrottlerInterface, error) {
    92  		if threadCount != 1 {
    93  			t.Errorf("want: 1, got: %v", threadCount)
    94  		}
    95  		return mockThrottler, nil
    96  	}
    97  
    98  	call0 := mockThrottler.EXPECT().UpdateConfiguration(gomock.Any(), true /* copyZeroValues */)
    99  	call1 := mockThrottler.EXPECT().Throttle(0)
   100  	call1.Return(0 * time.Second)
   101  	tabletStats := &discovery.TabletHealth{
   102  		Target: &querypb.Target{
   103  			TabletType: topodatapb.TabletType_REPLICA,
   104  		},
   105  	}
   106  	call2 := mockThrottler.EXPECT().RecordReplicationLag(gomock.Any(), tabletStats)
   107  	call3 := mockThrottler.EXPECT().Throttle(0)
   108  	call3.Return(1 * time.Second)
   109  	call4 := mockThrottler.EXPECT().Close()
   110  	call1.After(call0)
   111  	call2.After(call1)
   112  	call3.After(call2)
   113  	call4.After(call3)
   114  
   115  	config := tabletenv.NewDefaultConfig()
   116  	config.EnableTxThrottler = true
   117  	config.TxThrottlerHealthCheckCells = []string{"cell1", "cell2"}
   118  
   119  	throttler, err := tryCreateTxThrottler(config, ts)
   120  	if err != nil {
   121  		t.Fatalf("want: nil, got: %v", err)
   122  	}
   123  	throttler.InitDBConfig(&querypb.Target{
   124  		Keyspace: "keyspace",
   125  		Shard:    "shard",
   126  	})
   127  	if err := throttler.Open(); err != nil {
   128  		t.Fatalf("want: nil, got: %v", err)
   129  	}
   130  	if result := throttler.Throttle(); result != false {
   131  		t.Errorf("want: false, got: %v", result)
   132  	}
   133  	throttler.state.StatsUpdate(tabletStats)
   134  	rdonlyTabletStats := &discovery.TabletHealth{
   135  		Target: &querypb.Target{
   136  			TabletType: topodatapb.TabletType_RDONLY,
   137  		},
   138  	}
   139  	// This call should not be forwarded to the go/vt/throttler.Throttler object.
   140  	throttler.state.StatsUpdate(rdonlyTabletStats)
   141  	// The second throttle call should reject.
   142  	if result := throttler.Throttle(); result != true {
   143  		t.Errorf("want: true, got: %v", result)
   144  	}
   145  	throttler.Close()
   146  }