Flow display and user login and rate limit bug

lib/rate/conn.go

@@ -0,0 +1,38 @@
+package rate
+
+import (
+	"io"
+	"net"
+)
+
+type rateConn struct {
+	conn net.Conn
+	rate *Rate
+}
+
+func NewRateConn(conn net.Conn, rate *Rate) io.ReadWriteCloser {
+	return &rateConn{
+		conn: conn,
+		rate: rate,
+	}
+}
+
+func (s *rateConn) Read(b []byte) (n int, err error) {
+	n, err = s.conn.Read(b)
+	if s.rate != nil {
+		s.rate.Get(int64(n))
+	}
+	return
+}
+
+func (s *rateConn) Write(b []byte) (n int, err error) {
+	n, err = s.conn.Write(b)
+	if s.rate != nil {
+		s.rate.Get(int64(n))
+	}
+	return
+}
+
+func (s *rateConn) Close() error {
+	return s.conn.Close()
+}

lib/rate/rate.go

@@ -10,6 +10,7 @@ type Rate struct {
 	bucketSurplusSize int64     //当前桶中体积
 	bucketAddSize     int64     //每次加水大小
 	stopChan          chan bool //停止
+	NowRate           int64
 }
 
 func NewRate(addSize int64) *Rate {
@@ -26,7 +27,8 @@ func (s *Rate) Start() {
 }
 
 func (s *Rate) add(size int64) {
-	if (s.bucketSize - s.bucketSurplusSize) < s.bucketAddSize {
+	if res := s.bucketSize - s.bucketSurplusSize; res < s.bucketAddSize {
+		atomic.AddInt64(&s.bucketSurplusSize, res)
 		return
 	}
 	atomic.AddInt64(&s.bucketSurplusSize, size)
@@ -65,6 +67,11 @@ func (s *Rate) session() {
 	for {
 		select {
 		case <-ticker.C:
+			if rs := s.bucketAddSize - s.bucketSurplusSize; rs > 0 {
+				s.NowRate = rs
+			} else {
+				s.NowRate = s.bucketSize - s.bucketSurplusSize
+			}
 			s.add(s.bucketAddSize)
 		case <-s.stopChan:
 			ticker.Stop()