github.com/uber/kraken@v0.1.4/lib/backend/throttle.go (about)

     1  // Copyright (c) 2016-2019 Uber Technologies, Inc.
     2  //
     3  // Licensed under the Apache License, Version 2.0 (the "License");
     4  // you may not use this file except in compliance with the License.
     5  // You may obtain a copy of the License at
     6  //
     7  //     http://www.apache.org/licenses/LICENSE-2.0
     8  //
     9  // Unless required by applicable law or agreed to in writing, software
    10  // distributed under the License is distributed on an "AS IS" BASIS,
    11  // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    12  // See the License for the specific language governing permissions and
    13  // limitations under the License.
    14  package backend
    15  
    16  import (
    17  	"io"
    18  
    19  	"github.com/uber/kraken/lib/store"
    20  	"github.com/uber/kraken/utils/bandwidth"
    21  	"github.com/uber/kraken/utils/log"
    22  	"github.com/uber/kraken/utils/stringset"
    23  )
    24  
    25  // ThrottledClient is a backend client with speed limit.
    26  type ThrottledClient struct {
    27  	Client
    28  	bandwidth *bandwidth.Limiter
    29  }
    30  
    31  // throttle wraps client with bandwidth limits.
    32  func throttle(client Client, bandwidth *bandwidth.Limiter) *ThrottledClient {
    33  	return &ThrottledClient{client, bandwidth}
    34  }
    35  
    36  type sizer interface {
    37  	Size() int64
    38  }
    39  
    40  // Ensure that we can get size from file store readers.
    41  var _ sizer = (store.FileReader)(nil)
    42  
    43  // Upload uploads src into name.
    44  func (c *ThrottledClient) Upload(namespace, name string, src io.Reader) error {
    45  	if s, ok := src.(sizer); ok {
    46  		// Only throttle if the src implements a Size method.
    47  		if err := c.bandwidth.ReserveEgress(s.Size()); err != nil {
    48  			log.With("name", name).Errorf("Error reserving egress: %s", err)
    49  			// Ignore error.
    50  		}
    51  	}
    52  	return c.Client.Upload(namespace, name, src)
    53  }
    54  
    55  // Download downloads name into dst.
    56  func (c *ThrottledClient) Download(namespace, name string, dst io.Writer) error {
    57  	info, err := c.Client.Stat(namespace, name)
    58  	if err != nil {
    59  		return err
    60  	}
    61  	if err := c.bandwidth.ReserveIngress(info.Size); err != nil {
    62  		log.With("name", name).Errorf("Error reserving ingress: %s", err)
    63  		// Ignore error.
    64  	}
    65  	return c.Client.Download(namespace, name, dst)
    66  }
    67  
    68  func (c *ThrottledClient) adjustBandwidth(denominator int) error {
    69  	return c.bandwidth.Adjust(denominator)
    70  }
    71  
    72  // EgressLimit returns egress limit.
    73  func (c *ThrottledClient) EgressLimit() int64 {
    74  	return c.bandwidth.EgressLimit()
    75  }
    76  
    77  // IngressLimit returns ingress limit.
    78  func (c *ThrottledClient) IngressLimit() int64 {
    79  	return c.bandwidth.IngressLimit()
    80  }
    81  
    82  // BandwidthWatcher is a hashring.Watcher which adjusts bandwidth on throttled
    83  // backends when hashring membership changes.
    84  type BandwidthWatcher struct {
    85  	manager *Manager
    86  }
    87  
    88  // NewBandwidthWatcher creates a new BandwidthWatcher for manager.
    89  func NewBandwidthWatcher(manager *Manager) *BandwidthWatcher {
    90  	return &BandwidthWatcher{manager}
    91  }
    92  
    93  // Notify splits bandwidth across the size of latest.
    94  func (w *BandwidthWatcher) Notify(latest stringset.Set) {
    95  	if err := w.manager.AdjustBandwidth(len(latest)); err != nil {
    96  		log.With("latest", latest.ToSlice()).Errorf("Error adjusting bandwidth: %s", err)
    97  	}
    98  }