From 0702e373915c3ec37ce5b226ac66cb2eca3cb3c9 Mon Sep 17 00:00:00 2001
From: Qian Wang <cbeuw.andy@gmail.com>
Date: Sun, 11 Nov 2018 19:30:40 +0000
Subject: [PATCH] Syncing in-memory credit with db regularly

---
 cmd/ck-client/ck-client.go               |  2 +-
 internal/multiplex/qos.go                |  8 +++++++
 internal/server/usermanager/userpanel.go | 29 ++++++++++++++++++++++++
 3 files changed, 38 insertions(+), 1 deletion(-)

diff --git a/cmd/ck-client/ck-client.go b/cmd/ck-client/ck-client.go
index 893895e..2eeb923 100644
--- a/cmd/ck-client/ck-client.go
+++ b/cmd/ck-client/ck-client.go
@@ -146,7 +146,7 @@ func main() {
 		log.Fatal("TicketTimeHint cannot be empty or 0")
 	}
 
-	valve := mux.MakeValve(1e9, 1e9, 1e9, 1e9)
+	valve := mux.MakeValve(1e12, 1e12, 1e12, 1e12)
 	obfs := util.MakeObfs(sta.UID)
 	deobfs := util.MakeDeobfs(sta.UID)
 	sesh := mux.MakeSession(0, valve, obfs, deobfs, util.ReadTLS)
diff --git a/internal/multiplex/qos.go b/internal/multiplex/qos.go
index 581a6ec..ca51b0e 100644
--- a/internal/multiplex/qos.go
+++ b/internal/multiplex/qos.go
@@ -47,6 +47,14 @@ func (v *Valve) txWait(n int) {
 	v.txtb.Load().(*ratelimit.Bucket).Wait(int64(n))
 }
 
+func (v *Valve) GetRxCredit() int64 {
+	return atomic.LoadInt64(&v.rxCredit)
+}
+
+func (v *Valve) GetTxCredit() int64 {
+	return atomic.LoadInt64(&v.txCredit)
+}
+
 // n can be negative
 func (v *Valve) AddRxCredit(n int64) int64 {
 	return atomic.AddInt64(&v.rxCredit, n)
diff --git a/internal/server/usermanager/userpanel.go b/internal/server/usermanager/userpanel.go
index 67378aa..ab48fdd 100644
--- a/internal/server/usermanager/userpanel.go
+++ b/internal/server/usermanager/userpanel.go
@@ -4,6 +4,7 @@ import (
 	"encoding/binary"
 	"errors"
 	"sync"
+	"time"
 
 	"github.com/boltdb/bolt"
 )
@@ -24,6 +25,10 @@ func MakeUserpanel(dbPath string) (*Userpanel, error) {
 		db:          db,
 		activeUsers: make(map[[32]byte]*User),
 	}
+	go func() {
+		time.Sleep(time.Second * 10)
+		up.updateCredits()
+	}()
 	return up, nil
 }
 
@@ -147,3 +152,27 @@ func (up *Userpanel) SetSessionsCap(UID [32]byte, newSessionsCap uint32) error {
 	err := up.updateDBEntryUint32(UID, "sessionsCap", newSessionsCap)
 	return err
 }
+
+func (up *Userpanel) updateCredits() {
+	up.activeUsersM.RLock()
+	defer u.activeUsersM.RUnlock()
+	for _, user := range up.activeUsers {
+		up.db.Update(func(tx *bolt.Tx) error {
+			b := tx.Bucket(u.uid[:])
+			if b == nil {
+				return ErrUserNotFound
+			}
+			oct := make([]byte, 8)
+			binary.BigEndian.PutUint64(oct, uint64(u.valve.GetRxCredit()))
+			if err := b.Put([]byte("rxCredit"), oct); err != nil {
+				return err
+			}
+			binary.BigEndian.PutUint64(oct, uint64(u.valve.GetTxCredit()))
+			if err := b.Put([]byte("txCredit"), oct); err != nil {
+				return err
+			}
+			return nil
+
+		})
+	}
+}