From f212144310eef1541a4ac6bd08ab6d4eb4ab4524 Mon Sep 17 00:00:00 2001
From: naiba <hi@nai.ba>
Date: Sat, 21 Dec 2024 10:59:05 +0800
Subject: [PATCH] feat: add online user count

---
 cmd/dashboard/controller/ws.go      |  3 +++
 cmd/dashboard/controller/ws_test.go | 26 ++++++++++++++++++++++++++
 model/server_api.go                 |  1 +
 service/singleton/singleton.go      |  2 ++
 4 files changed, 32 insertions(+)
 create mode 100644 cmd/dashboard/controller/ws_test.go

diff --git a/cmd/dashboard/controller/ws.go b/cmd/dashboard/controller/ws.go
index 6a3ebe2..9fdf92e 100644
--- a/cmd/dashboard/controller/ws.go
+++ b/cmd/dashboard/controller/ws.go
@@ -107,6 +107,8 @@ func serverStream(c *gin.Context) (any, error) {
 		return nil, newWsError("%v", err)
 	}
 	defer conn.Close()
+	singleton.OnlineUsers.Add(1)
+	defer singleton.OnlineUsers.Add(^uint64(0))
 	count := 0
 	for {
 		stat, err := getServerStat(c, count == 0)
@@ -164,6 +166,7 @@ func getServerStat(c *gin.Context, withPublicNote bool) ([]byte, error) {
 
 		return utils.Json.Marshal(model.StreamServerData{
 			Now:     time.Now().Unix() * 1000,
+			Online:  singleton.OnlineUsers.Load(),
 			Servers: servers,
 		})
 	})
diff --git a/cmd/dashboard/controller/ws_test.go b/cmd/dashboard/controller/ws_test.go
new file mode 100644
index 0000000..d1d41e6
--- /dev/null
+++ b/cmd/dashboard/controller/ws_test.go
@@ -0,0 +1,26 @@
+package controller
+
+import (
+	"sync/atomic"
+	"testing"
+)
+
+func TestWs(t *testing.T) {
+	onlineUsers := new(atomic.Uint64)
+	onlineUsers.Add(1)
+	if onlineUsers.Load() != 1 {
+		t.Error("onlineUsers.Add(1) failed")
+	}
+	onlineUsers.Add(1)
+	if onlineUsers.Load() != 2 {
+		t.Error("onlineUsers.Add(1) failed")
+	}
+	onlineUsers.Add(^uint64(0))
+	if onlineUsers.Load() != 1 {
+		t.Error("onlineUsers.Add(^uint64(0)) failed")
+	}
+	onlineUsers.Add(^uint64(0))
+	if onlineUsers.Load() != 0 {
+		t.Error("onlineUsers.Add(^uint64(0)) failed")
+	}
+}
diff --git a/model/server_api.go b/model/server_api.go
index 4edfa8e..5328e73 100644
--- a/model/server_api.go
+++ b/model/server_api.go
@@ -16,6 +16,7 @@ type StreamServer struct {
 
 type StreamServerData struct {
 	Now     int64          `json:"now,omitempty"`
+	Online  uint64         `json:"online,omitempty"`
 	Servers []StreamServer `json:"servers,omitempty"`
 }
 
diff --git a/service/singleton/singleton.go b/service/singleton/singleton.go
index 0d2bedb..ed2c564 100644
--- a/service/singleton/singleton.go
+++ b/service/singleton/singleton.go
@@ -3,6 +3,7 @@ package singleton
 import (
 	_ "embed"
 	"log"
+	"sync/atomic"
 	"time"
 
 	"github.com/patrickmn/go-cache"
@@ -23,6 +24,7 @@ var (
 	Loc               *time.Location
 	FrontendTemplates []model.FrontendTemplate
 	DashboardBootTime = uint64(time.Now().Unix())
+	OnlineUsers       = new(atomic.Uint64)
 )
 
 //go:embed frontend-templates.yaml