github.com/slspeek/camlistore_namedsearch@v0.0.0-20140519202248-ed6f70f7721a/pkg/throttle/throttle.go (about)

     1  /*
     2  Copyright 2012 Google Inc.
     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 provides a net.Listener that returns
    18  // artificially-delayed connections for testing real-world
    19  // connectivity.
    20  package throttle
    21  
    22  import (
    23  	"fmt"
    24  	"net"
    25  	"sync"
    26  	"time"
    27  )
    28  
    29  const unitSize = 1400 // read/write chunk size. ~MTU size.
    30  
    31  type Rate struct {
    32  	KBps    int // or 0, to not rate-limit bandwidth
    33  	Latency time.Duration
    34  }
    35  
    36  // byteTime returns the time required for n bytes.
    37  func (r Rate) byteTime(n int) time.Duration {
    38  	if r.KBps == 0 {
    39  		return 0
    40  	}
    41  	return time.Duration(float64(n)/1024/float64(r.KBps)) * time.Second
    42  }
    43  
    44  type Listener struct {
    45  	net.Listener
    46  	Down Rate // server Writes to Client
    47  	Up   Rate // server Reads from client
    48  }
    49  
    50  func (ln *Listener) Accept() (net.Conn, error) {
    51  	c, err := ln.Listener.Accept()
    52  	time.Sleep(ln.Up.Latency)
    53  	if err != nil {
    54  		return nil, err
    55  	}
    56  	tc := &conn{Conn: c, Down: ln.Down, Up: ln.Up}
    57  	tc.start()
    58  	return tc, nil
    59  }
    60  
    61  type nErr struct {
    62  	n   int
    63  	err error
    64  }
    65  
    66  type writeReq struct {
    67  	writeAt time.Time
    68  	p       []byte
    69  	resc    chan nErr
    70  }
    71  
    72  type conn struct {
    73  	net.Conn
    74  	Down Rate // for reads
    75  	Up   Rate // for writes
    76  
    77  	wchan     chan writeReq
    78  	closeOnce sync.Once
    79  	closeErr  error
    80  }
    81  
    82  func (c *conn) start() {
    83  	c.wchan = make(chan writeReq, 1024)
    84  	go c.writeLoop()
    85  }
    86  
    87  func (c *conn) writeLoop() {
    88  	for req := range c.wchan {
    89  		time.Sleep(req.writeAt.Sub(time.Now()))
    90  		var res nErr
    91  		for len(req.p) > 0 && res.err == nil {
    92  			writep := req.p
    93  			if len(writep) > unitSize {
    94  				writep = writep[:unitSize]
    95  			}
    96  			n, err := c.Conn.Write(writep)
    97  			time.Sleep(c.Up.byteTime(len(writep)))
    98  			res.n += n
    99  			res.err = err
   100  			req.p = req.p[n:]
   101  		}
   102  		req.resc <- res
   103  	}
   104  }
   105  
   106  func (c *conn) Close() error {
   107  	c.closeOnce.Do(func() {
   108  		err := c.Conn.Close()
   109  		close(c.wchan)
   110  		c.closeErr = err
   111  	})
   112  	return c.closeErr
   113  }
   114  
   115  func (c *conn) Write(p []byte) (n int, err error) {
   116  	defer func() {
   117  		if e := recover(); e != nil {
   118  			n = 0
   119  			err = fmt.Errorf("%v", err)
   120  			return
   121  		}
   122  	}()
   123  	resc := make(chan nErr, 1)
   124  	c.wchan <- writeReq{time.Now().Add(c.Up.Latency), p, resc}
   125  	res := <-resc
   126  	return res.n, res.err
   127  }
   128  
   129  func (c *conn) Read(p []byte) (n int, err error) {
   130  	const max = 1024
   131  	if len(p) > max {
   132  		p = p[:max]
   133  	}
   134  	n, err = c.Conn.Read(p)
   135  	time.Sleep(c.Down.byteTime(n))
   136  	return
   137  }