Skip to content

Commit 560fac5

Browse files
Sayan-Sayan-
authored
Simpler stz manager (#8)
## Overview We noticed issues where implementation would have < 0 pending. It seems like the naive assumptions that events would be deduped and delivered exactly once weren't really reliable. Instead swap to a model where we only disable/enable around edges rather than every event - Use global state from the session manager rather than tracking internally - Centralize logic in a simple helper. Call the helper right at start up to be safe - Check errors and only update internal state on success ## Testing Ran locally + against a few unikernels in staging. I would connect and disconnect a bunch of tabs to simulate flaky connections (this was a pretty consistent repro of the original issue) Example logs from staging: ``` [neko] 11:19PM INF plugin enabled module=scaletozero [neko] 11:19PM INF no operation needed; skipping toggle currently_connected_sessions=0 currently_disabled=false module=scaletozero previously_disabled=false [neko] 11:19PM INF plugin started module=plugins plugin=scaletozero [neko] 11:29PM INF disabling scale-to-zero connected_sessions=1 module=scaletozero [neko] 11:29PM INF no operation needed; skipping toggle currently_connected_sessions=2 currently_disabled=true module=scaletozero previously_disabled=true [neko] 11:29PM INF no operation needed; skipping toggle currently_connected_sessions=3 currently_disabled=true module=scaletozero previously_disabled=true [neko] 11:29PM INF no operation needed; skipping toggle currently_connected_sessions=4 currently_disabled=true module=scaletozero previously_disabled=true [neko] 11:29PM INF no operation needed; skipping toggle currently_connected_sessions=5 currently_disabled=true module=scaletozero previously_disabled=true [neko] 11:29PM INF no operation needed; skipping toggle currently_connected_sessions=6 currently_disabled=true module=scaletozero previously_disabled=true [neko] 11:29PM INF no operation needed; skipping toggle currently_connected_sessions=7 currently_disabled=true module=scaletozero previously_disabled=true [neko] 11:29PM INF no operation needed; skipping toggle currently_connected_sessions=8 currently_disabled=true module=scaletozero previously_disabled=true [neko] 11:29PM INF no operation needed; skipping toggle currently_connected_sessions=9 currently_disabled=true module=scaletozero previously_disabled=true [neko] 11:29PM INF no operation needed; skipping toggle currently_connected_sessions=10 currently_disabled=true module=scaletozero previously_disabled=true [neko] 11:29PM INF no operation needed; skipping toggle currently_connected_sessions=11 currently_disabled=true module=scaletozero previously_disabled=true [neko] 11:29PM INF no operation needed; skipping toggle currently_connected_sessions=10 currently_disabled=true module=scaletozero previously_disabled=true [neko] 11:29PM INF no operation needed; skipping toggle currently_connected_sessions=9 currently_disabled=true module=scaletozero previously_disabled=true [neko] 11:29PM INF no operation needed; skipping toggle currently_connected_sessions=8 currently_disabled=true module=scaletozero previously_disabled=true [neko] 11:29PM INF no operation needed; skipping toggle currently_connected_sessions=7 currently_disabled=true module=scaletozero previously_disabled=true [neko] 11:29PM INF no operation needed; skipping toggle currently_connected_sessions=6 currently_disabled=true module=scaletozero previously_disabled=true [neko] 11:29PM INF no operation needed; skipping toggle currently_connected_sessions=5 currently_disabled=true module=scaletozero previously_disabled=true [neko] 11:29PM INF no operation needed; skipping toggle currently_connected_sessions=5 currently_disabled=true module=scaletozero previously_disabled=true [neko] 11:29PM INF no operation needed; skipping toggle currently_connected_sessions=4 currently_disabled=true module=scaletozero previously_disabled=true [neko] 11:29PM INF no operation needed; skipping toggle currently_connected_sessions=3 currently_disabled=true module=scaletozero previously_disabled=true [neko] 11:29PM INF no operation needed; skipping toggle currently_connected_sessions=2 currently_disabled=true module=scaletozero previously_disabled=true [neko] 11:29PM INF no operation needed; skipping toggle currently_connected_sessions=1 currently_disabled=true module=scaletozero previously_disabled=true [neko] 11:29PM INF enabling scale-to-zero module=scaletozero [neko] 11:29PM INF disabling scale-to-zero connected_sessions=1 module=scaletozero [neko] 11:30PM INF no operation needed; skipping toggle currently_connected_sessions=2 currently_disabled=true module=scaletozero previously_disabled=true [neko] 11:30PM INF no operation needed; skipping toggle currently_connected_sessions=3 currently_disabled=true module=scaletozero previously_disabled=true [neko] 11:30PM INF no operation needed; skipping toggle currently_connected_sessions=4 currently_disabled=true module=scaletozero previously_disabled=true [neko] 11:30PM INF no operation needed; skipping toggle currently_connected_sessions=5 currently_disabled=true module=scaletozero previously_disabled=true [neko] 11:30PM INF no operation needed; skipping toggle currently_connected_sessions=6 currently_disabled=true module=scaletozero previously_disabled=true [neko] 11:30PM INF no operation needed; skipping toggle currently_connected_sessions=7 currently_disabled=true module=scaletozero previously_disabled=true [neko] 11:30PM INF no operation needed; skipping toggle currently_connected_sessions=8 currently_disabled=true module=scaletozero previously_disabled=true [neko] 11:30PM INF no operation needed; skipping toggle currently_connected_sessions=9 currently_disabled=true module=scaletozero previously_disabled=true [neko] 11:30PM INF no operation needed; skipping toggle currently_connected_sessions=10 currently_disabled=true module=scaletozero previously_disabled=true [neko] 11:30PM INF no operation needed; skipping toggle currently_connected_sessions=11 currently_disabled=true module=scaletozero previously_disabled=true [neko] 11:30PM INF no operation needed; skipping toggle currently_connected_sessions=10 currently_disabled=true module=scaletozero previously_disabled=true [neko] 11:30PM INF no operation needed; skipping toggle currently_connected_sessions=9 currently_disabled=true module=scaletozero previously_disabled=true [neko] 11:30PM INF no operation needed; skipping toggle currently_connected_sessions=8 currently_disabled=true module=scaletozero previously_disabled=true [neko] 11:30PM INF no operation needed; skipping toggle currently_connected_sessions=7 currently_disabled=true module=scaletozero previously_disabled=true [neko] 11:30PM INF no operation needed; skipping toggle currently_connected_sessions=6 currently_disabled=true module=scaletozero previously_disabled=true [neko] 11:30PM INF no operation needed; skipping toggle currently_connected_sessions=5 currently_disabled=true module=scaletozero previously_disabled=true [neko] 11:30PM INF no operation needed; skipping toggle currently_connected_sessions=5 currently_disabled=true module=scaletozero previously_disabled=true [neko] 11:30PM INF no operation needed; skipping toggle currently_connected_sessions=4 currently_disabled=true module=scaletozero previously_disabled=true [neko] 11:30PM INF no operation needed; skipping toggle currently_connected_sessions=3 currently_disabled=true module=scaletozero previously_disabled=true [neko] 11:30PM INF no operation needed; skipping toggle currently_connected_sessions=2 currently_disabled=true module=scaletozero previously_disabled=true [neko] 11:30PM INF no operation needed; skipping toggle currently_connected_sessions=1 currently_disabled=true module=scaletozero previously_disabled=true [neko] 11:30PM INF enabling scale-to-zero module=scaletozero [neko] 11:30PM INF disabling scale-to-zero connected_sessions=1 module=scaletozero [neko] 11:30PM INF no operation needed; skipping toggle currently_connected_sessions=2 currently_disabled=true module=scaletozero previously_disabled=true [neko] 11:30PM INF no operation needed; skipping toggle currently_connected_sessions=3 currently_disabled=true module=scaletozero previously_disabled=true [neko] 11:30PM INF no operation needed; skipping toggle currently_connected_sessions=4 currently_disabled=true module=scaletozero previously_disabled=true [neko] 11:30PM INF no operation needed; skipping toggle currently_connected_sessions=5 currently_disabled=true module=scaletozero previously_disabled=true [neko] 11:30PM INF no operation needed; skipping toggle currently_connected_sessions=6 currently_disabled=true module=scaletozero previously_disabled=true [neko] 11:30PM INF no operation needed; skipping toggle currently_connected_sessions=7 currently_disabled=true module=scaletozero previously_disabled=true [neko] 11:30PM INF no operation needed; skipping toggle currently_connected_sessions=8 currently_disabled=true module=scaletozero previously_disabled=true [neko] 11:30PM INF no operation needed; skipping toggle currently_connected_sessions=9 currently_disabled=true module=scaletozero previously_disabled=true [neko] 11:30PM INF no operation needed; skipping toggle currently_connected_sessions=10 currently_disabled=true module=scaletozero previously_disabled=true [neko] 11:30PM INF no operation needed; skipping toggle currently_connected_sessions=11 currently_disabled=true module=scaletozero previously_disabled=true [neko] 11:30PM INF no operation needed; skipping toggle currently_connected_sessions=10 currently_disabled=true module=scaletozero previously_disabled=true [neko] 11:30PM INF no operation needed; skipping toggle currently_connected_sessions=9 currently_disabled=true module=scaletozero previously_disabled=true [neko] 11:30PM INF no operation needed; skipping toggle currently_connected_sessions=8 currently_disabled=true module=scaletozero previously_disabled=true [neko] 11:30PM INF no operation needed; skipping toggle currently_connected_sessions=7 currently_disabled=true module=scaletozero previously_disabled=true [neko] 11:30PM INF no operation needed; skipping toggle currently_connected_sessions=6 currently_disabled=true module=scaletozero previously_disabled=true [neko] 11:30PM INF no operation needed; skipping toggle currently_connected_sessions=5 currently_disabled=true module=scaletozero previously_disabled=true [neko] 11:30PM INF no operation needed; skipping toggle currently_connected_sessions=4 currently_disabled=true module=scaletozero previously_disabled=true [neko] 11:30PM INF no operation needed; skipping toggle currently_connected_sessions=3 currently_disabled=true module=scaletozero previously_disabled=true [neko] 11:30PM INF no operation needed; skipping toggle currently_connected_sessions=2 currently_disabled=true module=scaletozero previously_disabled=true [neko] 11:30PM INF no operation needed; skipping toggle currently_connected_sessions=1 currently_disabled=true module=scaletozero previously_disabled=true [neko] 11:30PM INF enabling scale-to-zero module=scaletozero [neko] 11:30PM INF disabling scale-to-zero connected_sessions=1 module=scaletozero [neko] 11:30PM INF no operation needed; skipping toggle currently_connected_sessions=2 currently_disabled=true module=scaletozero previously_disabled=true [neko] 11:30PM INF no operation needed; skipping toggle currently_connected_sessions=3 currently_disabled=true module=scaletozero previously_disabled=true [neko] 11:30PM INF no operation needed; skipping toggle currently_connected_sessions=4 currently_disabled=true module=scaletozero previously_disabled=true [neko] 11:30PM INF no operation needed; skipping toggle currently_connected_sessions=5 currently_disabled=true module=scaletozero previously_disabled=true [neko] 11:30PM INF no operation needed; skipping toggle currently_connected_sessions=6 currently_disabled=true module=scaletozero previously_disabled=true [neko] 11:30PM INF no operation needed; skipping toggle currently_connected_sessions=7 currently_disabled=true module=scaletozero previously_disabled=true [neko] 11:30PM INF no operation needed; skipping toggle currently_connected_sessions=8 currently_disabled=true module=scaletozero previously_disabled=true [neko] 11:30PM INF no operation needed; skipping toggle currently_connected_sessions=9 currently_disabled=true module=scaletozero previously_disabled=true [neko] 11:30PM INF no operation needed; skipping toggle currently_connected_sessions=10 currently_disabled=true module=scaletozero previously_disabled=true [neko] 11:30PM INF no operation needed; skipping toggle currently_connected_sessions=11 currently_disabled=true module=scaletozero previously_disabled=true [neko] 11:30PM INF no operation needed; skipping toggle currently_connected_sessions=10 currently_disabled=true module=scaletozero previously_disabled=true [neko] 11:30PM INF no operation needed; skipping toggle currently_connected_sessions=9 currently_disabled=true module=scaletozero previously_disabled=true [neko] 11:30PM INF no operation needed; skipping toggle currently_connected_sessions=8 currently_disabled=true module=scaletozero previously_disabled=true [neko] 11:30PM INF no operation needed; skipping toggle currently_connected_sessions=7 currently_disabled=true module=scaletozero previously_disabled=true [neko] 11:30PM INF no operation needed; skipping toggle currently_connected_sessions=6 currently_disabled=true module=scaletozero previously_disabled=true [neko] 11:30PM INF no operation needed; skipping toggle currently_connected_sessions=5 currently_disabled=true module=scaletozero previously_disabled=true [neko] 11:30PM INF no operation needed; skipping toggle currently_connected_sessions=4 currently_disabled=true module=scaletozero previously_disabled=true [neko] 11:30PM INF no operation needed; skipping toggle currently_connected_sessions=3 currently_disabled=true module=scaletozero previously_disabled=true [neko] 11:30PM INF no operation needed; skipping toggle currently_connected_sessions=2 currently_disabled=true module=scaletozero previously_disabled=true [neko] 11:30PM INF no operation needed; skipping toggle currently_connected_sessions=1 currently_disabled=true module=scaletozero previously_disabled=true [neko] 11:30PM INF enabling scale-to-zero module=scaletozero [neko] 11:30PM INF disabling scale-to-zero connected_sessions=1 module=scaletozero [neko] 11:30PM INF no operation needed; skipping toggle currently_connected_sessions=2 currently_disabled=true module=scaletozero previously_disabled=true [neko] 11:30PM INF no operation needed; skipping toggle currently_connected_sessions=3 currently_disabled=true module=scaletozero previously_disabled=true [neko] 11:30PM INF no operation needed; skipping toggle currently_connected_sessions=4 currently_disabled=true module=scaletozero previously_disabled=true [neko] 11:30PM INF no operation needed; skipping toggle currently_connected_sessions=5 currently_disabled=true module=scaletozero previously_disabled=true [neko] 11:30PM INF no operation needed; skipping toggle currently_connected_sessions=6 currently_disabled=true module=scaletozero previously_disabled=true [neko] 11:30PM INF no operation needed; skipping toggle currently_connected_sessions=7 currently_disabled=true module=scaletozero previously_disabled=true [neko] 11:30PM INF no operation needed; skipping toggle currently_connected_sessions=8 currently_disabled=true module=scaletozero previously_disabled=true [neko] 11:30PM INF no operation needed; skipping toggle currently_connected_sessions=9 currently_disabled=true module=scaletozero previously_disabled=true [neko] 11:30PM INF no operation needed; skipping toggle currently_connected_sessions=10 currently_disabled=true module=scaletozero previously_disabled=true [neko] 11:30PM INF no operation needed; skipping toggle currently_connected_sessions=11 currently_disabled=true module=scaletozero previously_disabled=true [neko] 11:31PM INF no operation needed; skipping toggle currently_connected_sessions=10 currently_disabled=true module=scaletozero previously_disabled=true [neko] 11:31PM INF no operation needed; skipping toggle currently_connected_sessions=9 currently_disabled=true module=scaletozero previously_disabled=true [neko] 11:31PM INF no operation needed; skipping toggle currently_connected_sessions=8 currently_disabled=true module=scaletozero previously_disabled=true [neko] 11:31PM INF no operation needed; skipping toggle currently_connected_sessions=7 currently_disabled=true module=scaletozero previously_disabled=true [neko] 11:31PM INF no operation needed; skipping toggle currently_connected_sessions=6 currently_disabled=true module=scaletozero previously_disabled=true [neko] 11:31PM INF no operation needed; skipping toggle currently_connected_sessions=5 currently_disabled=true module=scaletozero previously_disabled=true [neko] 11:31PM INF no operation needed; skipping toggle currently_connected_sessions=4 currently_disabled=true module=scaletozero previously_disabled=true [neko] 11:31PM INF no operation needed; skipping toggle currently_connected_sessions=3 currently_disabled=true module=scaletozero previously_disabled=true [neko] 11:31PM INF no operation needed; skipping toggle currently_connected_sessions=2 currently_disabled=true module=scaletozero previously_disabled=true [neko] 11:31PM INF no operation needed; skipping toggle currently_connected_sessions=1 currently_disabled=true module=scaletozero previously_disabled=true [neko] 11:31PM INF enabling scale-to-zero module=scaletozero [neko] 11:32PM INF disabling scale-to-zero connected_sessions=2 module=scaletozero [neko] 11:32PM INF no operation needed; skipping toggle currently_connected_sessions=2 currently_disabled=true module=scaletozero previously_disabled=true [neko] 11:32PM INF no operation needed; skipping toggle currently_connected_sessions=3 currently_disabled=true module=scaletozero previously_disabled=true [neko] 11:32PM INF no operation needed; skipping toggle currently_connected_sessions=4 currently_disabled=true module=scaletozero previously_disabled=true [neko] 11:32PM INF no operation needed; skipping toggle currently_connected_sessions=5 currently_disabled=true module=scaletozero previously_disabled=true [neko] 11:32PM INF no operation needed; skipping toggle currently_connected_sessions=6 currently_disabled=true module=scaletozero previously_disabled=true [neko] 11:32PM INF no operation needed; skipping toggle currently_connected_sessions=7 currently_disabled=true module=scaletozero previously_disabled=true [neko] 11:32PM INF no operation needed; skipping toggle currently_connected_sessions=8 currently_disabled=true module=scaletozero previously_disabled=true [neko] 11:32PM INF no operation needed; skipping toggle currently_connected_sessions=9 currently_disabled=true module=scaletozero previously_disabled=true [neko] 11:32PM INF no operation needed; skipping toggle currently_connected_sessions=10 currently_disabled=true module=scaletozero previously_disabled=true [neko] 11:32PM INF no operation needed; skipping toggle currently_connected_sessions=11 currently_disabled=true module=scaletozero previously_disabled=true [neko] 11:33PM INF no operation needed; skipping toggle currently_connected_sessions=10 currently_disabled=true module=scaletozero previously_disabled=true [neko] 11:33PM INF no operation needed; skipping toggle currently_connected_sessions=9 currently_disabled=true module=scaletozero previously_disabled=true [neko] 11:33PM INF no operation needed; skipping toggle currently_connected_sessions=8 currently_disabled=true module=scaletozero previously_disabled=true [neko] 11:33PM INF no operation needed; skipping toggle currently_connected_sessions=8 currently_disabled=true module=scaletozero previously_disabled=true [neko] 11:33PM INF no operation needed; skipping toggle currently_connected_sessions=7 currently_disabled=true module=scaletozero previously_disabled=true [neko] 11:33PM INF no operation needed; skipping toggle currently_connected_sessions=6 currently_disabled=true module=scaletozero previously_disabled=true [neko] 11:33PM INF no operation needed; skipping toggle currently_connected_sessions=5 currently_disabled=true module=scaletozero previously_disabled=true [neko] 11:33PM INF no operation needed; skipping toggle currently_connected_sessions=4 currently_disabled=true module=scaletozero previously_disabled=true [neko] 11:33PM INF no operation needed; skipping toggle currently_connected_sessions=3 currently_disabled=true module=scaletozero previously_disabled=true [neko] 11:33PM INF no operation needed; skipping toggle currently_connected_sessions=2 currently_disabled=true module=scaletozero previously_disabled=true [neko] 11:33PM INF no operation needed; skipping toggle currently_connected_sessions=2 currently_disabled=true module=scaletozero previously_disabled=true [neko] 11:33PM INF no operation needed; skipping toggle currently_connected_sessions=1 currently_disabled=true module=scaletozero previously_disabled=true [neko] 11:33PM INF enabling scale-to-zero module=scaletozero [neko] 11:33PM INF no operation needed; skipping toggle currently_connected_sessions=0 currently_disabled=false module=scaletozero previously_disabled=false ```
1 parent 46e2fc7 commit 560fac5

File tree

3 files changed

+246
-29
lines changed

3 files changed

+246
-29
lines changed

‎.github/workflows/server_test.yml‎

Lines changed: 56 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -42,3 +42,59 @@ jobs:
4242
with:
4343
context: ./server
4444
platforms: linux/arm64
45+
46+
server-tests:
47+
name: Server tests
48+
runs-on: ubuntu-latest
49+
steps:
50+
- uses: actions/checkout@v4
51+
with:
52+
fetch-depth: 0
53+
54+
- name: Set up Go
55+
uses: actions/setup-go@v5
56+
with:
57+
go-version-file: "server/go.mod"
58+
cache: true
59+
60+
- name: Cache apt packages
61+
uses: awalsh128/cache-apt-pkgs-action@v1
62+
with:
63+
packages: >
64+
libgstreamer1.0-dev
65+
libgstreamer-plugins-base1.0-dev
66+
libgtk-3-dev
67+
libx11-dev
68+
libxrandr-dev
69+
libxtst-dev
70+
libxfixes-dev
71+
libxcvt-dev
72+
pkg-config
73+
version: ${{ runner.os }}-ubuntu-24.04
74+
75+
- name: Install system dependencies
76+
run: |
77+
sudo apt-get update
78+
sudo apt-get install -y \
79+
libgstreamer1.0-dev \
80+
libgstreamer-plugins-base1.0-dev \
81+
libgtk-3-dev \
82+
libx11-dev \
83+
libxrandr-dev \
84+
libxtst-dev \
85+
libxfixes-dev \
86+
libxcvt-dev \
87+
pkg-config
88+
89+
- name: Cache Go modules
90+
uses: actions/cache@v4
91+
with:
92+
path: |
93+
~/.cache/go-build
94+
~/go/pkg/mod
95+
key: ${{ runner.os }}-go-${{ hashFiles('**/go.sum') }}
96+
restore-keys: ${{ runner.os }}-go-
97+
98+
- name: Run tests
99+
working-directory: server
100+
run: go test ./... -v

‎server/internal/plugins/scaletozero/manager.go‎

Lines changed: 57 additions & 29 deletions
Original file line numberDiff line numberDiff line change
@@ -25,43 +25,30 @@ func NewManager(
2525
}
2626

2727
type Manager struct {
28-
logger zerolog.Logger
29-
config *Config
30-
sessions types.SessionManager
31-
ctrl scaletozero.Controller
32-
mu sync.Mutex
33-
shutdown bool
34-
pending int
28+
logger zerolog.Logger
29+
config *Config
30+
sessions types.SessionManager
31+
ctrl scaletozero.Controller
32+
mu sync.Mutex
33+
shutdown bool
34+
disabledScaleToZero bool
3535
}
3636

3737
func (m *Manager) Start() error {
3838
if !m.config.Enabled {
3939
return nil
4040
}
41-
m.logger.Info().Msg("scale-to-zero plugin enabled")
41+
m.logger.Info().Msg("plugin enabled")
4242

43-
m.sessions.OnConnected(func(session types.Session) {
44-
m.mu.Lock()
45-
defer m.mu.Unlock()
46-
if m.shutdown {
47-
return
48-
}
43+
// compute initial state and toggle if needed
44+
m.manage()
4945

50-
m.pending++
51-
m.logger.Info().Msgf("connection started, disabling scale-to-zero (pending: %d)", m.pending)
52-
m.ctrl.Disable(context.Background())
46+
m.sessions.OnConnected(func(session types.Session) {
47+
m.manage()
5348
})
5449

5550
m.sessions.OnDisconnected(func(session types.Session) {
56-
m.mu.Lock()
57-
defer m.mu.Unlock()
58-
if m.shutdown {
59-
return
60-
}
61-
62-
m.pending--
63-
m.logger.Info().Msgf("connection started, disabling scale-to-zero (pending: %d)", m.pending)
64-
m.ctrl.Enable(context.Background())
51+
m.manage()
6552
})
6653

6754
return nil
@@ -72,10 +59,51 @@ func (m *Manager) Shutdown() error {
7259
defer m.mu.Unlock()
7360
m.shutdown = true
7461

75-
m.logger.Info().Msgf("shutdown started, re-enabling scale-to-zero (pending: %d)", m.pending)
76-
for i := 0; i < m.pending; i++ {
77-
m.ctrl.Enable(context.Background())
62+
if m.disabledScaleToZero {
63+
return m.ctrl.Enable(context.Background())
7864
}
7965

8066
return nil
8167
}
68+
69+
func (m *Manager) manage() {
70+
m.mu.Lock()
71+
defer m.mu.Unlock()
72+
73+
if m.shutdown {
74+
return
75+
}
76+
77+
connectedSessions := 0
78+
for _, s := range m.sessions.List() {
79+
if s.State().IsConnected {
80+
connectedSessions++
81+
}
82+
}
83+
hasConnectedSessions := connectedSessions > 0
84+
85+
if hasConnectedSessions == m.disabledScaleToZero {
86+
m.logger.Info().Bool("previously_disabled", m.disabledScaleToZero).
87+
Bool("currently_disabled", hasConnectedSessions).
88+
Int("currently_connected_sessions", connectedSessions).
89+
Msg("no operation needed; skipping toggle")
90+
return
91+
}
92+
93+
// toggle if needed but only update internal state if successful
94+
if hasConnectedSessions {
95+
m.logger.Info().Int("connected_sessions", connectedSessions).Msg("disabling scale-to-zero")
96+
if err := m.ctrl.Disable(context.Background()); err != nil {
97+
m.logger.Error().Err(err).Msg("failed to disable scale-to-zero")
98+
return
99+
}
100+
} else {
101+
m.logger.Info().Msg("enabling scale-to-zero")
102+
if err := m.ctrl.Enable(context.Background()); err != nil {
103+
m.logger.Error().Err(err).Msg("failed to enable scale-to-zero")
104+
return
105+
}
106+
}
107+
108+
m.disabledScaleToZero = hasConnectedSessions
109+
}

‎server/internal/plugins/scaletozero/manager_test.go‎

Lines changed: 133 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,133 @@
1+
package scaletozero
2+
3+
import (
4+
"context"
5+
"testing"
6+
"time"
7+
8+
"github.com/stretchr/testify/require"
9+
10+
"github.com/m1k1o/neko/server/internal/config"
11+
intsession "github.com/m1k1o/neko/server/internal/session"
12+
"github.com/m1k1o/neko/server/pkg/types"
13+
)
14+
15+
func TestSingleSessionConnectDisconnectReconnect(t *testing.T) {
16+
sm := newSessionManager(t)
17+
m, fc := newPluginWithFakeCtrl(sm)
18+
19+
require.NoError(t, m.Start())
20+
require.Equal(t, 0, fc.disableCalls)
21+
require.Equal(t, 0, fc.enableCalls)
22+
23+
s, p := connect(t, sm, "1")
24+
require.Equal(t, 1, fc.disableCalls)
25+
require.Equal(t, 0, fc.enableCalls)
26+
27+
s.DisconnectWebSocketPeer(p, true)
28+
require.Equal(t, 1, fc.disableCalls)
29+
require.Equal(t, 0, fc.enableCalls)
30+
31+
// wait for an arbitrary fraction of the delay duration to mimic client behavior
32+
start := time.Now()
33+
time.Sleep(intsession.WS_DELAYED_DURATION / 10)
34+
35+
// safeguard to prevent flake
36+
if time.Since(start) >= intsession.WS_DELAYED_DURATION {
37+
return
38+
}
39+
40+
// connect and ensure no subsequent stz calls
41+
_, p2 := connect(t, sm, "1")
42+
require.Equal(t, 1, fc.disableCalls)
43+
require.Equal(t, 0, fc.enableCalls)
44+
_ = p2
45+
}
46+
47+
func TestMultipleSessionsConnectDisconnect(t *testing.T) {
48+
sm := newSessionManager(t)
49+
m, fc := newPluginWithFakeCtrl(sm)
50+
require.NoError(t, m.Start())
51+
52+
s1, p1 := connect(t, sm, "1")
53+
require.Equal(t, 1, fc.disableCalls)
54+
require.Equal(t, 0, fc.enableCalls)
55+
_, p2 := connect(t, sm, "2")
56+
require.Equal(t, 1, fc.disableCalls)
57+
require.Equal(t, 0, fc.enableCalls)
58+
59+
// enable only after both sessions are disconnected
60+
s1.DisconnectWebSocketPeer(p1, false)
61+
require.Equal(t, 1, fc.disableCalls)
62+
require.Equal(t, 0, fc.enableCalls)
63+
s2, ok := sm.Get("2")
64+
require.True(t, ok)
65+
s2.DisconnectWebSocketPeer(p2, false)
66+
require.Equal(t, 1, fc.disableCalls)
67+
require.Equal(t, 1, fc.enableCalls)
68+
}
69+
70+
func TestSingleSessionReplacementDoesNotDoubleDisable(t *testing.T) {
71+
sm := newSessionManager(t)
72+
m, fc := newPluginWithFakeCtrl(sm)
73+
require.NoError(t, m.Start())
74+
75+
s, _ := connect(t, sm, "1")
76+
require.Equal(t, 1, fc.disableCalls)
77+
require.Equal(t, 0, fc.enableCalls)
78+
79+
// replacement: connect again while connected; should not trigger another disable
80+
p2 := &mockWebsocketPeer{}
81+
s.ConnectWebSocketPeer(p2)
82+
require.Equal(t, 1, fc.disableCalls)
83+
require.Equal(t, 0, fc.enableCalls)
84+
}
85+
86+
type mockScaleToZeroer struct {
87+
disableCalls int
88+
enableCalls int
89+
disableErr error
90+
enableErr error
91+
}
92+
93+
func (f *mockScaleToZeroer) Disable(ctx context.Context) error {
94+
f.disableCalls++
95+
return f.disableErr
96+
}
97+
98+
func (f *mockScaleToZeroer) Enable(ctx context.Context) error {
99+
f.enableCalls++
100+
return f.enableErr
101+
}
102+
103+
type mockWebsocketPeer struct{}
104+
105+
func (mockWebsocketPeer) Send(event string, payload any) {}
106+
func (mockWebsocketPeer) Ping() error { return nil }
107+
func (mockWebsocketPeer) Destroy(reason string) {}
108+
109+
func newSessionManager(t *testing.T) *intsession.SessionManagerCtx {
110+
t.Helper()
111+
return intsession.New(&config.Session{})
112+
}
113+
114+
func newPluginWithFakeCtrl(sm types.SessionManager) (*Manager, *mockScaleToZeroer) {
115+
fc := &mockScaleToZeroer{}
116+
m := NewManager(sm, &Config{Enabled: true})
117+
m.ctrl = fc
118+
return m, fc
119+
}
120+
121+
func connect(t *testing.T, sm types.SessionManager, id string) (types.Session, types.WebSocketPeer) {
122+
t.Helper()
123+
s, ok := sm.Get(id)
124+
if !ok {
125+
var err error
126+
s, _, err = sm.Create(id, types.MemberProfile{CanLogin: true, CanConnect: true, CanWatch: true})
127+
require.NoError(t, err)
128+
}
129+
p := &mockWebsocketPeer{}
130+
s.ConnectWebSocketPeer(p)
131+
132+
return s, p
133+
}

0 commit comments

Comments
 (0)