From 76340cc99c99adb45597cd194346a70d5f610002 Mon Sep 17 00:00:00 2001
From: rookisbusy <rookisbusy@gmail.com>
Date: Sat, 8 Apr 2023 00:55:25 +0800
Subject: [PATCH] feat: core support memory chat

---
 hub/route/server.go         | 48 +++++++++++++++++++++++++++++++++++++
 tunnel/statistic/manager.go |  8 ++++++-
 2 files changed, 55 insertions(+), 1 deletion(-)

diff --git a/hub/route/server.go b/hub/route/server.go
index 848face9..47cb7839 100644
--- a/hub/route/server.go
+++ b/hub/route/server.go
@@ -39,6 +39,11 @@ type Traffic struct {
 	Down int64 `json:"down"`
 }
 
+type Memory struct {
+	Inuse   uint64 `json:"inuse"`
+	OSLimit uint64 `json:"oslimit"` // maybe we need it in the future
+}
+
 func SetUIPath(path string) {
 	uiPath = C.Path.Resolve(path)
 }
@@ -76,6 +81,7 @@ func Start(addr string, tlsAddr string, secret string,
 		r.Get("/", hello)
 		r.Get("/logs", getLogs)
 		r.Get("/traffic", traffic)
+		r.Get("/memory", memory)
 		r.Get("/version", version)
 		r.Mount("/configs", configRouter())
 		r.Mount("/proxies", proxyRouter())
@@ -224,6 +230,48 @@ func traffic(w http.ResponseWriter, r *http.Request) {
 	}
 }
 
+func memory(w http.ResponseWriter, r *http.Request) {
+	var wsConn *websocket.Conn
+	if websocket.IsWebSocketUpgrade(r) {
+		var err error
+		wsConn, err = upgrader.Upgrade(w, r, nil)
+		if err != nil {
+			return
+		}
+	}
+
+	if wsConn == nil {
+		w.Header().Set("Content-Type", "application/json")
+		render.Status(r, http.StatusOK)
+	}
+
+	tick := time.NewTicker(time.Second)
+	defer tick.Stop()
+	t := statistic.DefaultManager
+	buf := &bytes.Buffer{}
+	var err error
+	for range tick.C {
+		buf.Reset()
+
+		if err := json.NewEncoder(buf).Encode(Memory{
+			Inuse:   t.Memory(),
+			OSLimit: 0,
+		}); err != nil {
+			break
+		}
+		if wsConn == nil {
+			_, err = w.Write(buf.Bytes())
+			w.(http.Flusher).Flush()
+		} else {
+			err = wsConn.WriteMessage(websocket.TextMessage, buf.Bytes())
+		}
+
+		if err != nil {
+			break
+		}
+	}
+}
+
 type Log struct {
 	Type    string `json:"type"`
 	Payload string `json:"payload"`
diff --git a/tunnel/statistic/manager.go b/tunnel/statistic/manager.go
index 09bb3410..7d0ee5e4 100644
--- a/tunnel/statistic/manager.go
+++ b/tunnel/statistic/manager.go
@@ -34,6 +34,7 @@ type Manager struct {
 	uploadTotal   *atomic.Int64
 	downloadTotal *atomic.Int64
 	pid           int
+	memory        uint64
 }
 
 func (m *Manager) Join(c tracker) {
@@ -58,6 +59,10 @@ func (m *Manager) Now() (up int64, down int64) {
 	return m.uploadBlip.Load(), m.downloadBlip.Load()
 }
 
+func (m *Manager) Memory() uint64 {
+	return m.memory
+}
+
 func (m *Manager) Snapshot() *Snapshot {
 	connections := []tracker{}
 	m.connections.Range(func(key, value any) bool {
@@ -76,12 +81,13 @@ func (m *Manager) Snapshot() *Snapshot {
 		}
 		return stat.RSS
 	}
+	m.memory = getMem()
 
 	return &Snapshot{
 		UploadTotal:   m.uploadTotal.Load(),
 		DownloadTotal: m.downloadTotal.Load(),
 		Connections:   connections,
-		Memory:        getMem(),
+		Memory:        m.memory,
 	}
 }