vitess.io/vitess@v0.16.2/go/test/endtoend/throttler/util.go (about)

     1  /*
     2  Copyright 2023 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 agreed to 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 throttler
    18  
    19  import (
    20  	"context"
    21  	"fmt"
    22  	"io"
    23  	"net/http"
    24  	"testing"
    25  	"time"
    26  
    27  	"github.com/buger/jsonparser"
    28  	"github.com/stretchr/testify/require"
    29  
    30  	"vitess.io/vitess/go/test/endtoend/cluster"
    31  	"vitess.io/vitess/go/vt/log"
    32  )
    33  
    34  type Config struct {
    35  	Query     string
    36  	Threshold float64
    37  }
    38  
    39  const (
    40  	DefaultQuery     = "select unix_timestamp(now(6))-max(ts/1000000000) as replication_lag from _vt.heartbeat"
    41  	DefaultThreshold = 1 * time.Second
    42  	ConfigTimeout    = 60 * time.Second
    43  )
    44  
    45  var DefaultConfig = &Config{
    46  	Query:     DefaultQuery,
    47  	Threshold: DefaultThreshold.Seconds(),
    48  }
    49  
    50  // UpdateThrottlerTopoConfig runs vtctlclient UpdateThrottlerConfig.
    51  // This retries the command until it succeeds or times out as the
    52  // SrvKeyspace record may not yet exist for a newly created
    53  // Keyspace that is still initializing before it becomes serving.
    54  func UpdateThrottlerTopoConfig(clusterInstance *cluster.LocalProcessCluster, enable bool, disable bool, threshold float64, metricsQuery string, viaVtctldClient bool) (result string, err error) {
    55  	args := []string{}
    56  	clientfunc := clusterInstance.VtctldClientProcess.ExecuteCommandWithOutput
    57  	if !viaVtctldClient {
    58  		args = append(args, "--")
    59  		clientfunc = clusterInstance.VtctlclientProcess.ExecuteCommandWithOutput
    60  	}
    61  	args = append(args, "UpdateThrottlerConfig")
    62  	if enable {
    63  		args = append(args, "--enable")
    64  	}
    65  	if disable {
    66  		args = append(args, "--disable")
    67  	}
    68  	if threshold > 0 {
    69  		args = append(args, "--threshold", fmt.Sprintf("%f", threshold))
    70  	}
    71  	args = append(args, "--custom-query", metricsQuery)
    72  	if metricsQuery != "" {
    73  		args = append(args, "--check-as-check-self")
    74  	} else {
    75  		args = append(args, "--check-as-check-shard")
    76  	}
    77  	args = append(args, clusterInstance.Keyspaces[0].Name)
    78  
    79  	ctx, cancel := context.WithTimeout(context.Background(), ConfigTimeout)
    80  	defer cancel()
    81  
    82  	ticker := time.NewTicker(time.Second)
    83  	defer ticker.Stop()
    84  
    85  	for {
    86  		result, err = clientfunc(args...)
    87  		if err == nil {
    88  			return result, nil
    89  		}
    90  		select {
    91  		case <-ctx.Done():
    92  			return "", fmt.Errorf("timed out waiting for UpdateThrottlerConfig to succeed after %v; last seen value: %+v, error: %v", ConfigTimeout, result, err)
    93  		case <-ticker.C:
    94  		}
    95  	}
    96  }
    97  
    98  // WaitForThrottlerStatusEnabled waits for a tablet to report its throttler status as
    99  // enabled/disabled and have the provided config (if any) until the specified timeout.
   100  func WaitForThrottlerStatusEnabled(t *testing.T, tablet *cluster.Vttablet, enabled bool, config *Config, timeout time.Duration) {
   101  	enabledJSONPath := "IsEnabled"
   102  	queryJSONPath := "Query"
   103  	thresholdJSONPath := "Threshold"
   104  	url := fmt.Sprintf("http://localhost:%d/throttler/status", tablet.HTTPPort)
   105  	ctx, cancel := context.WithTimeout(context.Background(), timeout)
   106  	defer cancel()
   107  	ticker := time.NewTicker(time.Second)
   108  	defer ticker.Stop()
   109  
   110  	for {
   111  		body := getHTTPBody(url)
   112  		isEnabled, err := jsonparser.GetBoolean([]byte(body), enabledJSONPath)
   113  		require.NoError(t, err)
   114  		if isEnabled == enabled {
   115  			if config == nil {
   116  				return
   117  			}
   118  			query, err := jsonparser.GetString([]byte(body), queryJSONPath)
   119  			require.NoError(t, err)
   120  			threshold, err := jsonparser.GetFloat([]byte(body), thresholdJSONPath)
   121  			require.NoError(t, err)
   122  			if query == config.Query && threshold == config.Threshold {
   123  				return
   124  			}
   125  		}
   126  		select {
   127  		case <-ctx.Done():
   128  			t.Errorf("timed out waiting for the %s tablet's throttler status enabled to be %t with the correct config after %v; last seen value: %s",
   129  				tablet.Alias, enabled, timeout, body)
   130  			return
   131  		case <-ticker.C:
   132  		}
   133  	}
   134  }
   135  
   136  func getHTTPBody(url string) string {
   137  	resp, err := http.Get(url)
   138  	if err != nil {
   139  		log.Infof("http Get returns %+v", err)
   140  		return ""
   141  	}
   142  	defer resp.Body.Close()
   143  	if resp.StatusCode != 200 {
   144  		log.Infof("http Get returns status %d", resp.StatusCode)
   145  		return ""
   146  	}
   147  	respByte, _ := io.ReadAll(resp.Body)
   148  	body := string(respByte)
   149  	return body
   150  }
   151  
   152  // WaitForQueryResult waits for a tablet to return the given result for the given
   153  // query until the specified timeout.
   154  // This is for simple queries that return 1 column in 1 row. It compares the result
   155  // for that column as a string with the provided result.
   156  func WaitForQueryResult(t *testing.T, tablet *cluster.Vttablet, query, result string, timeout time.Duration) {
   157  	ctx, cancel := context.WithTimeout(context.Background(), timeout)
   158  	defer cancel()
   159  	ticker := time.NewTicker(time.Second)
   160  	defer ticker.Stop()
   161  
   162  	for {
   163  		res, err := tablet.VttabletProcess.QueryTablet(query, "", false)
   164  		require.NoError(t, err)
   165  		if res != nil && len(res.Rows) == 1 && res.Rows[0][0].ToString() == result {
   166  			return
   167  		}
   168  		select {
   169  		case <-ctx.Done():
   170  			t.Errorf("timed out waiting for the %q query to produce a result of %q on tablet %s after %v; last seen value: %s",
   171  				query, result, tablet.Alias, timeout, res.Rows[0][0].ToString())
   172  			return
   173  		case <-ticker.C:
   174  		}
   175  	}
   176  }
   177  
   178  // WaitForValidData waits for a tablet's checks to return a non 500 http response
   179  // which indicates that it's not able to provide valid results. This is most
   180  // commonly caused by the throttler still gathering the initial results for
   181  // the given configuration.
   182  func WaitForValidData(t *testing.T, tablet *cluster.Vttablet, timeout time.Duration) {
   183  	checkURL := fmt.Sprintf("http://localhost:%d/throttler/check", tablet.HTTPPort)
   184  	selfCheckURL := fmt.Sprintf("http://localhost:%d/throttler/check-self", tablet.HTTPPort)
   185  	ctx, cancel := context.WithTimeout(context.Background(), timeout)
   186  	defer cancel()
   187  	ticker := time.NewTicker(500 * time.Millisecond)
   188  	defer ticker.Stop()
   189  
   190  	for {
   191  		checkResp, checkErr := http.Get(checkURL)
   192  		if checkErr != nil {
   193  			defer checkResp.Body.Close()
   194  		}
   195  		selfCheckResp, selfCheckErr := http.Get(selfCheckURL)
   196  		if selfCheckErr != nil {
   197  			defer selfCheckResp.Body.Close()
   198  		}
   199  		if checkErr == nil && selfCheckErr == nil &&
   200  			checkResp.StatusCode != http.StatusInternalServerError &&
   201  			selfCheckResp.StatusCode != http.StatusInternalServerError {
   202  			return
   203  		}
   204  		select {
   205  		case <-ctx.Done():
   206  			t.Errorf("timed out waiting for %s tablet's throttler to return a valid result after %v; last seen value: %+v",
   207  				tablet.Alias, timeout, checkResp)
   208  			return
   209  		case <-ticker.C:
   210  		}
   211  	}
   212  }