allow NewLimiter to pass limits

Author: linkdata

limiter.go

@@ -15,10 +15,12 @@ type Limiter struct {
 	Writes *Operation
 }
 
-func NewLimiter(ctx context.Context) *Limiter {
+// NewLimiter returns a new limiter. If you provide limits, the first will set
+// both read and write limits, the second will set the write limit.
+func NewLimiter(ctx context.Context, limits ...int64) *Limiter {
 	return &Limiter{
-		Reads:  NewOperation(ctx, true),
-		Writes: NewOperation(ctx, false),
+		Reads:  NewOperation(ctx, limits, 0),
+		Writes: NewOperation(ctx, limits, 1),
 	}
 }
 

operation.go

@@ -21,9 +21,17 @@ type Operation struct {
 	reader bool
 }
 
-func NewOperation(ctx context.Context, reader bool) (op *Operation) {
+func NewOperation(ctx context.Context, limits []int64, idx int) (op *Operation) {
 	ch := make(chan struct{}, secparts)
-	op = &Operation{ch: ch, reader: reader}
+	op = &Operation{ch: ch, reader: idx == 0}
+	var limit int64
+	if len(limits) > 0 {
+		limit = limits[0]
+		if len(limits) > idx {
+			limit = limits[idx]
+		}
+	}
+	op.Limit.Store(limit)
 	go op.run(ctx, ch)
 	return
 }

operation_test.go

@@ -92,9 +92,7 @@ func TestOperation_io_read_limit(t *testing.T) {
 func TestOperation_read_rate_low(t *testing.T) {
 	ctx, cancel := context.WithTimeout(context.Background(), time.Second*2)
 	defer cancel()
-	l := NewLimiter(ctx)
-
-	l.Reads.Limit.Store(1000) // 10 bytes @ 1000/sec
+	l := NewLimiter(ctx, 1000)
 
 	now := time.Now()
 	r := bytes.NewReader(make([]byte, 2000))
@@ -109,6 +107,7 @@ func TestOperation_read_rate_low(t *testing.T) {
 	}
 
 	if elapsed := time.Since(now); elapsed < time.Millisecond*900 || elapsed > time.Millisecond*1100 {
+		t.Log(l.Reads.Limit.Load())
 		t.Error(elapsed)
 	}
 	if rate := int(l.Reads.Rate.Load()); rate < 990 || rate > 1000 {
@@ -119,10 +118,8 @@ func TestOperation_read_rate_low(t *testing.T) {
 func TestOperation_read_rate_high(t *testing.T) {
 	ctx, cancel := context.WithTimeout(context.Background(), time.Second)
 	defer cancel()
-	l := NewLimiter(ctx)
-
 	const numbytes = (2 * 1024 * 1024 * 1024) - 1
-	l.Reads.Limit.Store(numbytes)
+	l := NewLimiter(ctx, numbytes)
 
 	now := time.Now()
 	r := &unlimitedReader{}
@@ -158,8 +155,8 @@ func TestOperation_read_rate_high(t *testing.T) {
 func TestOperation_write_rate(t *testing.T) {
 	ctx, cancel := context.WithTimeout(context.Background(), time.Second*2)
 	defer cancel()
-	l := NewLimiter(ctx)
-	l.Writes.Limit.Store(1000000)
+
+	l := NewLimiter(ctx, 1000000)
 
 	buf := make([]byte, 10000)