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  }