Rewrite rate

Author: mohanson

cmd/daze/main.go

@@ -11,6 +11,7 @@ import (
 	"os"
 	"path/filepath"
 	"strings"
+	"time"
 
 	"github.com/mohanson/daze"
 	"github.com/mohanson/daze/lib/doa"
@@ -95,7 +96,7 @@ func main() {
 			if *flLimits != "" {
 				n := daze.SizeParser(*flLimits)
 				doa.Doa(n <= math.MaxInt)
-				server.Limits = rate.NewLimiter(rate.Limit(n), int(n))
+				server.Limits = rate.NewLimits(n, time.Second)
 			}
 			defer server.Close()
 			doa.Nil(server.Run())
@@ -107,7 +108,7 @@ func main() {
 			if *flLimits != "" {
 				n := daze.SizeParser(*flLimits)
 				doa.Doa(n <= math.MaxInt)
-				server.Limits = rate.NewLimiter(rate.Limit(n), int(n))
+				server.Limits = rate.NewLimits(n, time.Second)
 			}
 			defer server.Close()
 			doa.Nil(server.Run())
@@ -116,7 +117,7 @@ func main() {
 			if *flLimits != "" {
 				n := daze.SizeParser(*flLimits)
 				doa.Doa(n <= math.MaxInt)
-				server.Limits = rate.NewLimiter(rate.Limit(n), int(n))
+				server.Limits = rate.NewLimits(n, time.Second)
 			}
 			defer server.Close()
 			doa.Nil(server.Run())
@@ -125,7 +126,7 @@ func main() {
 			if *flLimits != "" {
 				n := daze.SizeParser(*flLimits)
 				doa.Doa(n <= math.MaxInt)
-				server.Limits = rate.NewLimiter(rate.Limit(n), int(n))
+				server.Limits = rate.NewLimits(n, time.Second)
 			}
 			defer server.Close()
 			doa.Nil(server.Run())
@@ -177,7 +178,7 @@ func main() {
 			if *flLimits != "" {
 				n := daze.SizeParser(*flLimits)
 				doa.Doa(n <= math.MaxInt)
-				locale.Limits = rate.NewLimiter(rate.Limit(n), int(n))
+				locale.Limits = rate.NewLimits(n, time.Second)
 			}
 			defer locale.Close()
 			doa.Nil(locale.Run())
@@ -191,7 +192,7 @@ func main() {
 			if *flLimits != "" {
 				n := daze.SizeParser(*flLimits)
 				doa.Doa(n <= math.MaxInt)
-				locale.Limits = rate.NewLimiter(rate.Limit(n), int(n))
+				locale.Limits = rate.NewLimits(n, time.Second)
 			}
 			defer locale.Close()
 			doa.Nil(locale.Run())
@@ -206,7 +207,7 @@ func main() {
 			if *flLimits != "" {
 				n := daze.SizeParser(*flLimits)
 				doa.Doa(n <= math.MaxInt)
-				locale.Limits = rate.NewLimiter(rate.Limit(n), int(n))
+				locale.Limits = rate.NewLimits(n, time.Second)
 			}
 			defer locale.Close()
 			doa.Nil(locale.Run())
@@ -215,7 +216,7 @@ func main() {
 			if *flLimits != "" {
 				n := daze.SizeParser(*flLimits)
 				doa.Doa(n <= math.MaxInt)
-				client.Limits = rate.NewLimiter(rate.Limit(n), int(n))
+				client.Limits = rate.NewLimits(n, time.Second)
 			}
 			defer client.Close()
 			doa.Nil(client.Run())

daze.go

@@ -192,7 +192,7 @@ func (d *Direct) Dial(ctx *Context, network string, address string) (io.ReadWrit
 type Locale struct {
 	Closer io.Closer
 	Dialer Dialer
-	Limits *rate.Limiter
+	Limits *rate.Limits
 	Listen string
 }
 
@@ -627,7 +627,7 @@ func (l *Locale) Run() error {
 func NewLocale(listen string, dialer Dialer) *Locale {
 	return &Locale{
 		Dialer: dialer,
-		Limits: rate.NewLimiter(rate.Inf, 0),
+		Limits: rate.NewLimits(math.MaxUint64, time.Second),
 		Listen: listen,
 	}
 }
@@ -1082,7 +1082,7 @@ func (r *RandomReader) Read(p []byte) (int, error) {
 // RateConn wraps a net.Conn with a per-conn and a rate limiter.
 type RateConn struct {
 	Conn io.ReadWriteCloser
-	Rate *rate.Limiter
+	Rate *rate.Limits
 }
 
 // Close closes the connection.
@@ -1093,14 +1093,14 @@ func (r *RateConn) Close() error {
 // Read reads up to len(p) bytes into p.
 func (r *RateConn) Read(p []byte) (int, error) {
 	n, err := r.Conn.Read(p)
-	r.Rate.WaitN(context.Background(), n)
+	r.Rate.Wait(uint64(n))
 	return n, err
 }
 
 // Write writes len(p) bytes from p to the underlying data stream.
 func (r *RateConn) Write(p []byte) (int, error) {
 	n, err := r.Conn.Write(p)
-	r.Rate.WaitN(context.Background(), n)
+	r.Rate.Wait(uint64(n))
 	return n, err
 }
 

lib/rate/README.md

@@ -0,0 +1,3 @@
+# Rate
+
+Package rate provides a rate limiter.