Skip to content

Commit

Permalink
Vendor throttle pkg to avoid a dependency
Browse files Browse the repository at this point in the history
  • Loading branch information
samuel committed Apr 19, 2016
1 parent 218e9c8 commit 47c87b9
Show file tree
Hide file tree
Showing 3 changed files with 142 additions and 7 deletions.
3 changes: 2 additions & 1 deletion .travis.yml
Original file line number Diff line number Diff line change
@@ -1,6 +1,7 @@
language: go
go:
- 1.4.2
- 1.6.1
- tip

sudo: false

Expand Down
136 changes: 136 additions & 0 deletions zk/throttle_test.go
Original file line number Diff line number Diff line change
@@ -0,0 +1,136 @@
/*
Copyright 2012 Google Inc.
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
http://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
*/

// Vendored from go4.org/net/throttle

package zk

import (
"fmt"
"net"
"sync"
"time"
)

const unitSize = 1400 // read/write chunk size. ~MTU size.

type Rate struct {
KBps int // or 0, to not rate-limit bandwidth
Latency time.Duration
}

// byteTime returns the time required for n bytes.
func (r Rate) byteTime(n int) time.Duration {
if r.KBps == 0 {
return 0
}
return time.Duration(float64(n)/1024/float64(r.KBps)) * time.Second
}

type Listener struct {
net.Listener
Down Rate // server Writes to Client
Up Rate // server Reads from client
}

func (ln *Listener) Accept() (net.Conn, error) {
c, err := ln.Listener.Accept()
time.Sleep(ln.Up.Latency)
if err != nil {
return nil, err
}
tc := &conn{Conn: c, Down: ln.Down, Up: ln.Up}
tc.start()
return tc, nil
}

type nErr struct {
n int
err error
}

type writeReq struct {
writeAt time.Time
p []byte
resc chan nErr
}

type conn struct {
net.Conn
Down Rate // for reads
Up Rate // for writes

wchan chan writeReq
closeOnce sync.Once
closeErr error
}

func (c *conn) start() {
c.wchan = make(chan writeReq, 1024)
go c.writeLoop()
}

func (c *conn) writeLoop() {
for req := range c.wchan {
time.Sleep(req.writeAt.Sub(time.Now()))
var res nErr
for len(req.p) > 0 && res.err == nil {
writep := req.p
if len(writep) > unitSize {
writep = writep[:unitSize]
}
n, err := c.Conn.Write(writep)
time.Sleep(c.Up.byteTime(len(writep)))
res.n += n
res.err = err
req.p = req.p[n:]
}
req.resc <- res
}
}

func (c *conn) Close() error {
c.closeOnce.Do(func() {
err := c.Conn.Close()
close(c.wchan)
c.closeErr = err
})
return c.closeErr
}

func (c *conn) Write(p []byte) (n int, err error) {
defer func() {
if e := recover(); e != nil {
n = 0
err = fmt.Errorf("%v", err)
return
}
}()
resc := make(chan nErr, 1)
c.wchan <- writeReq{time.Now().Add(c.Up.Latency), p, resc}
res := <-resc
return res.n, res.err
}

func (c *conn) Read(p []byte) (n int, err error) {
const max = 1024
if len(p) > max {
p = p[:max]
}
n, err = c.Conn.Read(p)
time.Sleep(c.Down.byteTime(n))
return
}
10 changes: 4 additions & 6 deletions zk/zk_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -7,8 +7,6 @@ import (
"strings"
"testing"
"time"

"camlistore.org/pkg/throttle"
)

func TestCreate(t *testing.T) {
Expand Down Expand Up @@ -411,8 +409,8 @@ func TestSlowServer(t *testing.T) {

realAddr := fmt.Sprintf("127.0.0.1:%d", ts.Servers[0].Port)
proxyAddr, stopCh, err := startSlowProxy(t,
throttle.Rate{}, throttle.Rate{},
realAddr, func(ln *throttle.Listener) {
Rate{}, Rate{},
realAddr, func(ln *Listener) {
if ln.Up.Latency == 0 {
ln.Up.Latency = time.Millisecond * 2000
ln.Down.Latency = time.Millisecond * 2000
Expand Down Expand Up @@ -461,12 +459,12 @@ func TestSlowServer(t *testing.T) {
}
}

func startSlowProxy(t *testing.T, up, down throttle.Rate, upstream string, adj func(ln *throttle.Listener)) (string, chan bool, error) {
func startSlowProxy(t *testing.T, up, down Rate, upstream string, adj func(ln *Listener)) (string, chan bool, error) {
ln, err := net.Listen("tcp", "127.0.0.1:0")
if err != nil {
return "", nil, err
}
tln := &throttle.Listener{
tln := &Listener{
Listener: ln,
Up: up,
Down: down,
Expand Down

0 comments on commit 47c87b9

Please sign in to comment.