rscp

git clone https://orangeshoelaces.net/git/rscp.git

/bwcapio.go

   1 package main
   2 
   3 import (
   4         "io"
   5         "time"
   6 )
   7 
   8 type BwStats struct {
   9         Last   time.Time /* time of last observed event */
  10         Wnd    uint      /* unmetered bytes */
  11         Thresh uint      /* delay after at least this much bytes */
  12         Rate   uint      /* bandwidth limit in bits/second */
  13 }
  14 
  15 func NewBwStats(rate uint) *BwStats {
  16         return &BwStats{Wnd: 0, Thresh: rate, Rate: rate}
  17 }
  18 
  19 func CapReader(r io.Reader, st *BwStats) io.Reader {
  20         if st == nil {
  21                 panic("nil stats")
  22         }
  23         return &BwCapReader{r, st}
  24 }
  25 
  26 func CapWriter(w io.Writer, st *BwStats) io.Writer {
  27         if st == nil {
  28                 panic("nil stats")
  29         }
  30         return &BwCapWriter{w, st}
  31 }
  32 
  33 type BwCapReader struct {
  34         Base  io.Reader
  35         Stats *BwStats
  36 }
  37 
  38 func (r *BwCapReader) Read(p []byte) (int, error) {
  39         n, err := r.Base.Read(p)
  40         bwCap(r.Stats, n)
  41         return n, err
  42 }
  43 
  44 type BwCapWriter struct {
  45         Base  io.Writer
  46         Stats *BwStats
  47 }
  48 
  49 func (w *BwCapWriter) Write(p []byte) (int, error) {
  50         n, err := w.Base.Write(p)
  51         bwCap(w.Stats, n)
  52         return n, err
  53 }
  54 
  55 func bwCap(st *BwStats, transfered int) {
  56         if transfered <= 0 {
  57                 return 
  58         }
  59         if st.Last.IsZero() {
  60                 st.Last = time.Now()
  61                 return
  62         }
  63         st.Wnd += uint(transfered)
  64         if st.Wnd < st.Thresh {
  65                 return
  66         }
  67 
  68         bits := st.Wnd * 8
  69         exp := time.Duration((1e9 * bits) / st.Rate)
  70         ahead := exp - time.Since(st.Last)
  71 
  72         if ahead > 0 {
  73                 if ahead.Seconds() > 1 {
  74                         st.Thresh /= 2
  75                 } else if ahead < 10*time.Millisecond {
  76                         st.Thresh *= 2
  77                 }
  78                 time.Sleep(ahead)
  79         }
  80 
  81         st.Wnd = 0
  82         st.Last = time.Now()
  83 }