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 }