use ContextDialer isof DialContextFn

Author: linkdata

contextdialer.go

@@ -0,0 +1,10 @@
+package bwlimit
+
+import (
+	"context"
+	"net"
+)
+
+type ContextDialer interface {
+	DialContext(ctx context.Context, network, address string) (conn net.Conn, err error)
+}

dialer.go

@@ -6,12 +6,12 @@ import (
 )
 
 type Dialer struct {
-	DialContextFn // DialContext function we wrap
+	ContextDialer // ContextDialer we wrap
 	*Limiter      // Limiter to use
 }
 
 func (d *Dialer) DialContext(ctx context.Context, network, address string) (conn net.Conn, err error) {
-	if conn, err = d.DialContextFn(ctx, network, address); err == nil {
+	if conn, err = d.ContextDialer.DialContext(ctx, network, address); err == nil {
 		conn = &Conn{
 			Conn:    conn,
 			Limiter: d.Limiter,

dialer_test.go

@@ -13,12 +13,12 @@ func TestDialer_Dial(t *testing.T) {
 	defer l2.Stop()
 
 	d1 := &Dialer{
-		DialContextFn: l1.Wrap(nil),
 		Limiter:       l1,
+		ContextDialer: l1.Wrap(nil),
 	}
 	d2 := &Dialer{
-		DialContextFn: l2.Wrap(d1.DialContext),
 		Limiter:       l2,
+		ContextDialer: l2.Wrap(d1),
 	}
 
 	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {}))

example_test.go

@@ -23,7 +23,7 @@ func ExampleLimiter_NewLimiter() {
 
 	// wrap the default http transport DialContext
 	tp := http.DefaultTransport.(*http.Transport)
-	tp.DialContext = lim.Wrap(tp.DialContext)
+	tp.DialContext = lim.Wrap(nil).DialContext
 
 	// make a request and time it
 	now := time.Now()

limiter.go

@@ -1,12 +1,9 @@
 package bwlimit
 
 import (
-	"context"
 	"net"
 )
 
-type DialContextFn func(ctx context.Context, network string, address string) (net.Conn, error)
-
 var DefaultNetDialer = &net.Dialer{}
 
 type Limiter struct {
@@ -32,12 +29,31 @@ func (l *Limiter) Stop() {
 	l.Writes.Stop()
 }
 
-// Wrap returns a DialContextFn using the given fn that is bandwidth limited by this Limiter.
-// If fn is nil we use DefaultNetDialer.DialContext.
-func (l *Limiter) Wrap(fn DialContextFn) DialContextFn {
-	if fn == nil {
-		fn = DefaultNetDialer.DialContext
+// alreadyLimits returns true if cd is already limited by this Limiter.
+// This lets us help the user avoiding double-accounting bandwidth.
+func (l *Limiter) alreadyLimits(cd ContextDialer) bool {
+	for {
+		if d, ok := cd.(*Dialer); ok {
+			if d.Limiter == l {
+				return true
+			}
+			cd = d.ContextDialer
+		} else {
+			return false
+		}
+	}
+}
+
+// Wrap returns a ContextDialer wrapping cd that is bandwidth limited by this Limiter.
+//
+// If cd is nil we use DefaultNetDialer. If cd is already limited by this Limiter, cd
+// is returned unchanged.
+func (l *Limiter) Wrap(cd ContextDialer) ContextDialer {
+	if cd == nil {
+		cd = DefaultNetDialer
+	}
+	if l.alreadyLimits(cd) {
+		return cd
 	}
-	d := &Dialer{DialContextFn: fn, Limiter: l}
-	return d.DialContext
+	return &Dialer{ContextDialer: cd, Limiter: l}
 }

limiter_test.go

@@ -31,3 +31,14 @@ func TestLimiter_Stop(t *testing.T) {
 		t.Error(n)
 	}
 }
+
+func TestLimiter_double_Wrap(t *testing.T) {
+	l := NewLimiter()
+	defer l.Stop()
+
+	d1 := l.Wrap(nil)
+	d2 := l.Wrap(d1)
+	if d1 != d2 {
+		t.Error(d1, d2)
+	}
+}