vendor/github.com/tendermint/tmlibs/flowrate/io.go

   1 //
   2 // Written by Maxim Khitrov (November 2012)
   3 //
   4
   5 package flowrate
   6
   7 import (
   8         "errors"
   9         "io"
  10 )
  11
  12 // ErrLimit is returned by the Writer when a non-blocking write is short due to
  13 // the transfer rate limit.
  14 var ErrLimit = errors.New("flowrate: flow rate limit exceeded")
  15
  16 // Limiter is implemented by the Reader and Writer to provide a consistent
  17 // interface for monitoring and controlling data transfer.
  18 type Limiter interface {
  19         Done() int64
  20         Status() Status
  21         SetTransferSize(bytes int64)
  22         SetLimit(new int64) (old int64)
  23         SetBlocking(new bool) (old bool)
  24 }
  25
  26 // Reader implements io.ReadCloser with a restriction on the rate of data
  27 // transfer.
  28 type Reader struct {
  29         io.Reader // Data source
  30         *Monitor  // Flow control monitor
  31
  32         limit int64 // Rate limit in bytes per second (unlimited when <= 0)
  33         block bool  // What to do when no new bytes can be read due to the limit
  34 }
  35
  36 // NewReader restricts all Read operations on r to limit bytes per second.
  37 func NewReader(r io.Reader, limit int64) *Reader {
  38         return &Reader{r, New(0, 0), limit, true}
  39 }
  40
  41 // Read reads up to len(p) bytes into p without exceeding the current transfer
  42 // rate limit. It returns (0, nil) immediately if r is non-blocking and no new
  43 // bytes can be read at this time.
  44 func (r *Reader) Read(p []byte) (n int, err error) {
  45         p = p[:r.Limit(len(p), r.limit, r.block)]
  46         if len(p) > 0 {
  47                 n, err = r.IO(r.Reader.Read(p))
  48         }
  49         return
  50 }
  51
  52 // SetLimit changes the transfer rate limit to new bytes per second and returns
  53 // the previous setting.
  54 func (r *Reader) SetLimit(new int64) (old int64) {
  55         old, r.limit = r.limit, new
  56         return
  57 }
  58
  59 // SetBlocking changes the blocking behavior and returns the previous setting. A
  60 // Read call on a non-blocking reader returns immediately if no additional bytes
  61 // may be read at this time due to the rate limit.
  62 func (r *Reader) SetBlocking(new bool) (old bool) {
  63         old, r.block = r.block, new
  64         return
  65 }
  66
  67 // Close closes the underlying reader if it implements the io.Closer interface.
  68 func (r *Reader) Close() error {
  69         defer r.Done()
  70         if c, ok := r.Reader.(io.Closer); ok {
  71                 return c.Close()
  72         }
  73         return nil
  74 }
  75
  76 // Writer implements io.WriteCloser with a restriction on the rate of data
  77 // transfer.
  78 type Writer struct {
  79         io.Writer // Data destination
  80         *Monitor  // Flow control monitor
  81
  82         limit int64 // Rate limit in bytes per second (unlimited when <= 0)
  83         block bool  // What to do when no new bytes can be written due to the limit
  84 }
  85
  86 // NewWriter restricts all Write operations on w to limit bytes per second. The
  87 // transfer rate and the default blocking behavior (true) can be changed
  88 // directly on the returned *Writer.
  89 func NewWriter(w io.Writer, limit int64) *Writer {
  90         return &Writer{w, New(0, 0), limit, true}
  91 }
  92
  93 // Write writes len(p) bytes from p to the underlying data stream without
  94 // exceeding the current transfer rate limit. It returns (n, ErrLimit) if w is
  95 // non-blocking and no additional bytes can be written at this time.
  96 func (w *Writer) Write(p []byte) (n int, err error) {
  97         var c int
  98         for len(p) > 0 && err == nil {
  99                 s := p[:w.Limit(len(p), w.limit, w.block)]
 100                 if len(s) > 0 {
 101                         c, err = w.IO(w.Writer.Write(s))
 102                 } else {
 103                         return n, ErrLimit
 104                 }
 105                 p = p[c:]
 106                 n += c
 107         }
 108         return
 109 }
 110
 111 // SetLimit changes the transfer rate limit to new bytes per second and returns
 112 // the previous setting.
 113 func (w *Writer) SetLimit(new int64) (old int64) {
 114         old, w.limit = w.limit, new
 115         return
 116 }
 117
 118 // SetBlocking changes the blocking behavior and returns the previous setting. A
 119 // Write call on a non-blocking writer returns as soon as no additional bytes
 120 // may be written at this time due to the rate limit.
 121 func (w *Writer) SetBlocking(new bool) (old bool) {
 122         old, w.block = w.block, new
 123         return
 124 }
 125
 126 // Close closes the underlying writer if it implements the io.Closer interface.
 127 func (w *Writer) Close() error {
 128         defer w.Done()
 129         if c, ok := w.Writer.(io.Closer); ok {
 130                 return c.Close()
 131         }
 132         return nil
 133 }