vitess.io/vitess@v0.16.2/go/vt/vttablet/tabletserver/throttle/client.go (about) 1 /* 2 Copyright 2021 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 throttle 18 19 import ( 20 "context" 21 "net/http" 22 "sync" 23 "sync/atomic" 24 "time" 25 ) 26 27 const ( 28 throttleCheckDuration = 250 * time.Millisecond 29 ) 30 31 var throttleTicks int64 32 var throttleInit sync.Once 33 34 func initThrottleTicker() { 35 throttleInit.Do(func() { 36 go func() { 37 tick := time.NewTicker(throttleCheckDuration) 38 defer tick.Stop() 39 for range tick.C { 40 atomic.AddInt64(&throttleTicks, 1) 41 } 42 }() 43 }) 44 } 45 46 // Client construct is used by apps who wish to consult with a throttler. It encapsulates the check/throttling/backoff logic 47 type Client struct { 48 throttler *Throttler 49 appName string 50 checkType ThrottleCheckType 51 flags CheckFlags 52 53 lastSuccessfulThrottle int64 54 } 55 56 // NewProductionClient creates a client suitable for foreground/production jobs, which have normal priority. 57 func NewProductionClient(throttler *Throttler, appName string, checkType ThrottleCheckType) *Client { 58 initThrottleTicker() 59 return &Client{ 60 throttler: throttler, 61 appName: appName, 62 checkType: checkType, 63 flags: CheckFlags{ 64 LowPriority: false, 65 }, 66 } 67 } 68 69 // NewBackgroundClient creates a client suitable for background jobs, which have low priority over productio ntraffic, 70 // e.g. migration, table pruning, vreplication 71 func NewBackgroundClient(throttler *Throttler, appName string, checkType ThrottleCheckType) *Client { 72 initThrottleTicker() 73 return &Client{ 74 throttler: throttler, 75 appName: appName, 76 checkType: checkType, 77 flags: CheckFlags{ 78 LowPriority: true, 79 }, 80 } 81 } 82 83 // ThrottleCheckOK checks the throttler, and returns 'true' when the throttler is satisfied. 84 // It does not sleep. 85 // The function caches results for a brief amount of time, hence it's safe and efficient to 86 // be called very frequenty. 87 // The function is not thread safe. 88 func (c *Client) ThrottleCheckOK(ctx context.Context, overrideAppName string) (throttleCheckOK bool) { 89 if c == nil { 90 // no client 91 return true 92 } 93 if c.throttler == nil { 94 // no throttler 95 return true 96 } 97 if c.lastSuccessfulThrottle >= atomic.LoadInt64(&throttleTicks) { 98 // if last check was OK just very recently there is no need to check again 99 return true 100 } 101 // It's time to run a throttler check 102 checkApp := c.appName 103 if overrideAppName != "" { 104 checkApp = overrideAppName 105 } 106 checkResult := c.throttler.CheckByType(ctx, checkApp, "", &c.flags, c.checkType) 107 if checkResult.StatusCode != http.StatusOK { 108 return false 109 } 110 c.lastSuccessfulThrottle = atomic.LoadInt64(&throttleTicks) 111 return true 112 113 } 114 115 // ThrottleCheckOKOrWait checks the throttler; if throttler is satisfied, the function returns 'true' mmediately, 116 // otherwise it briefly sleeps and returns 'false'. 117 // Non-empty appName overrides the default appName. 118 // The function is not thread safe. 119 func (c *Client) ThrottleCheckOKOrWaitAppName(ctx context.Context, appName string) bool { 120 ok := c.ThrottleCheckOK(ctx, appName) 121 if !ok { 122 time.Sleep(throttleCheckDuration) 123 } 124 return ok 125 } 126 127 // ThrottleCheckOKOrWait checks the throttler; if throttler is satisfied, the function returns 'true' mmediately, 128 // otherwise it briefly sleeps and returns 'false'. 129 // The function is not thread safe. 130 func (c *Client) ThrottleCheckOKOrWait(ctx context.Context) bool { 131 return c.ThrottleCheckOKOrWaitAppName(ctx, "") 132 } 133 134 // Throttle throttles until the throttler is satisfied, or until context is cancelled. 135 // The function sleeps between throttle checks. 136 // The function is not thread safe. 137 func (c *Client) Throttle(ctx context.Context) { 138 for { 139 if ctx.Err() != nil { 140 return 141 } 142 if c.ThrottleCheckOKOrWait(ctx) { 143 break 144 } 145 } 146 }