1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
|
package rules
import (
"os"
"path/filepath"
"testing"
"time"
)
func TestConnectionStoreConnectAndRecommend(t *testing.T) {
tmpDir := t.TempDir()
store := NewConnectionStore(tmpDir)
if err := store.Load(); err != nil {
t.Fatal(err)
}
availableIPs := []string{"5.180.97.198", "5.180.97.199", "5.180.97.197"}
healthyIPs := map[string]bool{"5.180.97.198": true, "5.180.97.199": true, "5.180.97.197": true}
// Studio 1 connects to 198
store.Connect("1.2.3.4", "5.180.97.198", "nl-198", "windows", "2.0.11")
// Studio 1 asks — should get 198 (least loaded: 1 client vs 0/0)
// Actually: 198=1, 199=0, 197=0 → should get 199 or 197 (least loaded)
rec1 := store.GetRecommendation("1.2.3.4", availableIPs, healthyIPs)
if rec1.RecommendedServerIP == "5.180.97.198" {
// This is OK if load balancing picks a different server
t.Logf("studio 1 recommended: %s (reason: %s)", rec1.RecommendedServerIP, rec1.Reason)
} else {
t.Logf("studio 1 recommended different server: %s (load-based)", rec1.RecommendedServerIP)
}
// Studio 2 is new — should also get least loaded
rec2 := store.GetRecommendation("9.9.9.9", availableIPs, healthyIPs)
if rec2.RecommendedServerIP == "" {
t.Fatal("expected recommendation for new studio")
}
t.Logf("studio 2 recommended: %s (reason: %s)", rec2.RecommendedServerIP, rec2.Reason)
}
func TestPureLoadBalancing(t *testing.T) {
tmpDir := t.TempDir()
store := NewConnectionStore(tmpDir)
if err := store.Load(); err != nil {
t.Fatal(err)
}
availableIPs := []string{"5.180.97.198", "5.180.97.199", "5.180.97.197"}
healthyIPs := map[string]bool{"5.180.97.198": true, "5.180.97.199": true, "5.180.97.197": true}
// 3 studios connect to 198 (overload it)
for i := 0; i < 3; i++ {
ip := "10.0.0." + string(rune('1'+i))
store.Connect(ip, "5.180.97.198", "nl-198", "windows", "")
}
// New studio should NOT get 198 (3 clients) — should get 199 or 197 (0 clients)
rec := store.GetRecommendation("99.99.99.99", availableIPs, healthyIPs)
if rec.RecommendedServerIP == "5.180.97.198" {
t.Fatalf("should not recommend overloaded server, got %s", rec.RecommendedServerIP)
}
t.Logf("new studio recommended: %s (reason: %s)", rec.RecommendedServerIP, rec.Reason)
// Even studio 1 (home=198) should get load-balanced recommendation
rec2 := store.GetRecommendation("10.0.0.1", availableIPs, healthyIPs)
t.Logf("studio 1 re-recommended: %s (reason: %s, isRebalance: %v)",
rec2.RecommendedServerIP, rec2.Reason, rec2.IsRebalance)
}
func TestHomeServerUnhealthy(t *testing.T) {
tmpDir := t.TempDir()
store := NewConnectionStore(tmpDir)
if err := store.Load(); err != nil {
t.Fatal(err)
}
availableIPs := []string{"5.180.97.198", "5.180.97.199"}
// 198 is NOT healthy
healthyIPs := map[string]bool{"5.180.97.199": true}
// Studio 1 has connected to 198
store.Connect("1.2.3.4", "5.180.97.198", "nl-198", "windows", "")
// But 198 is unhealthy — should recommend 199
rec := store.GetRecommendation("1.2.3.4", availableIPs, healthyIPs)
if rec.RecommendedServerIP == "5.180.97.198" {
t.Fatalf("should not recommend unhealthy server, got %s", rec.RecommendedServerIP)
}
if rec.RecommendedServerIP != "5.180.97.199" {
t.Fatalf("should recommend healthy server 199, got %s", rec.RecommendedServerIP)
}
}
func TestDisconnect(t *testing.T) {
tmpDir := t.TempDir()
store := NewConnectionStore(tmpDir)
if err := store.Load(); err != nil {
t.Fatal(err)
}
availableIPs := []string{"5.180.97.198"}
store.Connect("1.2.3.4", "5.180.97.198", "nl-198", "windows", "")
load := store.GetLoadInfo(availableIPs)
if len(load) == 0 || load[0].ActiveClients != 1 {
t.Fatalf("expected 1 active client, got %v", load)
}
store.Disconnect("1.2.3.4")
load = store.GetLoadInfo(availableIPs)
if len(load) == 0 || load[0].ActiveClients != 0 {
t.Fatalf("expected 0 active clients after disconnect, got %v", load)
}
}
func TestSessionExpiry(t *testing.T) {
tmpDir := t.TempDir()
store := NewConnectionStore(tmpDir)
store.staleAfter = 1 * time.Millisecond
if err := store.Load(); err != nil {
t.Fatal(err)
}
availableIPs := []string{"5.180.97.198"}
healthyIPs := map[string]bool{"5.180.97.198": true}
store.Connect("1.2.3.4", "5.180.97.198", "nl-198", "windows", "")
time.Sleep(10 * time.Millisecond)
rec := store.GetRecommendation("1.2.3.4", availableIPs, healthyIPs)
if rec.RecommendedServerIP != "5.180.97.198" {
t.Fatalf("expected recommendation to 198 after session expiry, got %s", rec.RecommendedServerIP)
}
load := store.GetLoadInfo(availableIPs)
if len(load) == 0 || load[0].ActiveClients != 0 {
t.Fatalf("expected 0 active clients after expiry, got %v", load)
}
}
func TestPersistence(t *testing.T) {
tmpDir := t.TempDir()
store1 := NewConnectionStore(tmpDir)
store1.Connect("1.2.3.4", "5.180.97.199", "nl-199", "windows", "")
store2 := NewConnectionStore(tmpDir)
if err := store2.Load(); err != nil {
t.Fatal(err)
}
availableIPs := []string{"5.180.97.199"}
healthyIPs := map[string]bool{"5.180.97.199": true}
rec := store2.GetRecommendation("1.2.3.4", availableIPs, healthyIPs)
if rec.RecommendedServerIP != "5.180.97.199" {
t.Fatalf("expected recommendation to 199, got %s", rec.RecommendedServerIP)
}
_, err := os.Stat(filepath.Join(tmpDir, "connections.json"))
if err != nil {
t.Fatal("expected connections.json to exist")
}
}
func TestLoadInfoFormat(t *testing.T) {
tmpDir := t.TempDir()
store := NewConnectionStore(tmpDir)
store.maxCap = 10
store.Connect("1.1.1.1", "5.180.97.198", "nl-198", "windows", "")
store.Connect("2.2.2.2", "5.180.97.198", "nl-198", "windows", "")
store.Connect("3.3.3.3", "5.180.97.199", "nl-199", "linux", "")
availableIPs := []string{"5.180.97.198", "5.180.97.199"}
load := store.GetLoadInfo(availableIPs)
if len(load) != 2 {
t.Fatalf("expected 2 server load entries, got %d", len(load))
}
for _, info := range load {
if info.ServerIP == "5.180.97.198" {
if info.ActiveClients != 2 {
t.Errorf("expected 2 clients on 198, got %d", info.ActiveClients)
}
if info.LoadPercent != 20 {
t.Errorf("expected 20%% load on 198, got %d", info.LoadPercent)
}
}
if info.ServerIP == "5.180.97.199" {
if info.ActiveClients != 1 {
t.Errorf("expected 1 client on 199, got %d", info.ActiveClients)
}
if info.LoadPercent != 10 {
t.Errorf("expected 10%% load on 199, got %d", info.LoadPercent)
}
}
}
}
|
