Skip to content

Commit 03d79af

Browse files
hhftechnologyhhftechnology
committed
update
1 parent f4f1929 commit 03d79af

File tree

2 files changed

+112
-63
lines changed

2 files changed

+112
-63
lines changed

‎queue-manager.go

Lines changed: 98 additions & 60 deletions
Original file line numberDiff line numberDiff line change
@@ -1,3 +1,7 @@
1+
// Package queuemanager implements a Traefik middleware plugin that manages
2+
// access to services by implementing a queue when capacity is reached.
3+
// It functions similar to a virtual waiting room, allowing a controlled number
4+
// of users to access the service while placing others in a structured queue.
15
package queuemanager
26

37
import (
@@ -134,47 +138,64 @@ func (qm *QueueManager) ServeHTTP(rw http.ResponseWriter, req *http.Request) {
134138
log.Printf("[Queue Manager] Active sessions: %d, Max entries: %d", len(qm.activeSessionIDs), qm.config.MaxEntries)
135139
}
136140

137-
// Check if client is already in active sessions
138-
if qm.activeSessionIDs[clientID] {
139-
// Update last seen timestamp
140-
if session, found := qm.cache.Get(clientID); found {
141-
sessionData, ok := session.(Session)
142-
if !ok {
143-
if qm.config.Debug {
144-
log.Printf("[Queue Manager] Error: Failed to convert session to Session type")
145-
}
146-
} else {
147-
sessionData.LastSeen = time.Now()
148-
qm.cache.Set(clientID, sessionData, cache.DefaultExpiration)
149-
}
150-
}
151-
152-
// Allow access
141+
// Check if the client can proceed directly
142+
if qm.canClientProceed(clientID) {
153143
qm.next.ServeHTTP(rw, req)
154144
return
155145
}
156146

157-
// If there's room for more active sessions, add this client
147+
// Handle client that needs to wait
148+
position := qm.placeClientInQueue(clientID)
149+
150+
// Serve queue page
151+
qm.serveQueuePage(rw, position)
152+
}
153+
154+
// canClientProceed checks if a client can proceed without waiting.
155+
func (qm *QueueManager) canClientProceed(clientID string) bool {
156+
// If client is in active sessions, update timestamp and proceed
157+
if qm.activeSessionIDs[clientID] {
158+
qm.updateClientTimestamp(clientID)
159+
return true
160+
}
161+
162+
// If there's room for more active sessions, add client and proceed
158163
if len(qm.activeSessionIDs) < qm.config.MaxEntries {
159-
// Create new session
160164
session := Session{
161165
ID: clientID,
162166
CreatedAt: time.Now(),
163167
LastSeen: time.Now(),
164168
Position: 0,
165169
}
166170

167-
// Add to active sessions
168171
qm.activeSessionIDs[clientID] = true
169172
qm.cache.Set(clientID, session, cache.DefaultExpiration)
173+
return true
174+
}
175+
176+
// No capacity available
177+
return false
178+
}
179+
180+
// updateClientTimestamp updates the last seen timestamp for a client.
181+
func (qm *QueueManager) updateClientTimestamp(clientID string) {
182+
if session, found := qm.cache.Get(clientID); found {
183+
sessionData, ok := session.(Session)
184+
if !ok {
185+
if qm.config.Debug {
186+
log.Printf("[Queue Manager] Error: Failed to convert session to Session type")
187+
}
188+
return
189+
}
170190

171-
// Allow access
172-
qm.next.ServeHTTP(rw, req)
173-
return
191+
sessionData.LastSeen = time.Now()
192+
qm.cache.Set(clientID, sessionData, cache.DefaultExpiration)
174193
}
194+
}
175195

176-
// At this point, we know the client needs to wait
177-
// Let's check if they're already in the queue
196+
// placeClientInQueue places a client in the waiting queue and returns their position.
197+
func (qm *QueueManager) placeClientInQueue(clientID string) int {
198+
// Check if they're already in the queue
178199
position := -1
179200
for i, session := range qm.queue {
180201
if session.ID == clientID {
@@ -197,20 +218,9 @@ func (qm *QueueManager) ServeHTTP(rw http.ResponseWriter, req *http.Request) {
197218
}
198219

199220
// Update last seen timestamp
200-
if session, found := qm.cache.Get(clientID); found {
201-
sessionData, ok := session.(Session)
202-
if !ok {
203-
if qm.config.Debug {
204-
log.Printf("[Queue Manager] Error: Failed to convert session to Session type")
205-
}
206-
} else {
207-
sessionData.LastSeen = time.Now()
208-
qm.cache.Set(clientID, sessionData, cache.DefaultExpiration)
209-
}
210-
}
211-
212-
// Serve queue page
213-
qm.serveQueuePage(rw, position)
221+
qm.updateClientTimestamp(clientID)
222+
223+
return position
214224
}
215225

216226
// getClientID generates or retrieves a unique client identifier.
@@ -310,6 +320,20 @@ func getClientIP(req *http.Request) string {
310320

311321
// serveQueuePage serves the queue page HTML.
312322
func (qm *QueueManager) serveQueuePage(rw http.ResponseWriter, position int) {
323+
// Prepare template data
324+
data := qm.prepareQueuePageData(position)
325+
326+
// Try to use the template file
327+
if qm.serveCustomTemplate(rw, data) {
328+
return
329+
}
330+
331+
// Fall back to the default template
332+
qm.serveFallbackTemplate(rw, data)
333+
}
334+
335+
// prepareQueuePageData creates the data structure for the queue page template.
336+
func (qm *QueueManager) prepareQueuePageData(position int) QueuePageData {
313337
// Calculate estimated wait time (rough estimate: 30 seconds per position)
314338
estimatedWaitTime := int(math.Ceil(float64(position) * 0.5)) // in minutes
315339

@@ -319,40 +343,54 @@ func (qm *QueueManager) serveQueuePage(rw http.ResponseWriter, position int) {
319343
progressPercentage = int(100 - (float64(position) / float64(len(qm.queue)) * 100))
320344
}
321345

322-
// Prepare template data
323-
data := QueuePageData{
346+
return QueuePageData{
324347
Position: position + 1, // 1-based position for users
325348
QueueSize: len(qm.queue),
326349
EstimatedWaitTime: estimatedWaitTime,
327350
RefreshInterval: qm.config.RefreshInterval,
328351
ProgressPercentage: progressPercentage,
329352
Message: "Please wait while we process your request.",
330353
}
354+
}
355+
356+
// serveCustomTemplate attempts to serve the custom queue page template.
357+
// Returns true if successful, false if the template could not be served.
358+
func (qm *QueueManager) serveCustomTemplate(rw http.ResponseWriter, data QueuePageData) bool {
359+
if !fileExists(qm.config.QueuePageFile) {
360+
return false
361+
}
331362

332-
// Try to use the template file
333-
if fileExists(qm.config.QueuePageFile) {
334-
content, err := os.ReadFile(qm.config.QueuePageFile)
335-
if err == nil {
336-
queueTemplate, parseErr := template.New("QueuePage").Delims("[[", "]]").Parse(string(content))
337-
if parseErr == nil {
338-
// Set content type
339-
rw.Header().Set("Content-Type", qm.config.HTTPContentType)
340-
rw.WriteHeader(qm.config.HTTPResponseCode)
341-
342-
// Execute template
343-
if execErr := queueTemplate.Execute(rw, data); execErr != nil && qm.config.Debug {
344-
log.Printf("[Queue Manager] Error executing template: %v", execErr)
345-
}
346-
return
347-
} else if qm.config.Debug {
348-
log.Printf("[Queue Manager] Error parsing template: %v", parseErr)
349-
}
350-
} else if qm.config.Debug {
363+
content, err := os.ReadFile(qm.config.QueuePageFile)
364+
if err != nil {
365+
if qm.config.Debug {
351366
log.Printf("[Queue Manager] Error reading template file: %v", err)
352367
}
368+
return false
369+
}
370+
371+
queueTemplate, parseErr := template.New("QueuePage").Delims("[[", "]]").Parse(string(content))
372+
if parseErr != nil {
373+
if qm.config.Debug {
374+
log.Printf("[Queue Manager] Error parsing template: %v", parseErr)
375+
}
376+
return false
377+
}
378+
379+
// Set content type
380+
rw.Header().Set("Content-Type", qm.config.HTTPContentType)
381+
rw.WriteHeader(qm.config.HTTPResponseCode)
382+
383+
// Execute template
384+
if execErr := queueTemplate.Execute(rw, data); execErr != nil && qm.config.Debug {
385+
log.Printf("[Queue Manager] Error executing template: %v", execErr)
386+
return false
353387
}
354388

355-
// Fall back to a simple built-in template if there's an issue with the file
389+
return true
390+
}
391+
392+
// serveFallbackTemplate serves the built-in default queue page template.
393+
func (qm *QueueManager) serveFallbackTemplate(rw http.ResponseWriter, data QueuePageData) {
356394
fallbackTemplate := `<!DOCTYPE html>
357395
<html>
358396
<head>

‎queue-manager_test.go

Lines changed: 14 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -1,9 +1,11 @@
1+
// Package queuemanager provides a Traefik middleware to control traffic flow
2+
// by implementing a queue system similar to a virtual waiting room.
13
package queuemanager
24

35
import (
46
"net/http"
5-
"testing"
67
"os"
8+
"testing"
79
)
810

911
func TestGenerateUniqueID(t *testing.T) {
@@ -130,10 +132,19 @@ func TestGetClientIP(t *testing.T) {
130132
func TestFileExists(t *testing.T) {
131133
// Create a temporary file
132134
tempFile := "temp_test_file.txt"
133-
if _, err := os.Create(tempFile); err != nil {
135+
f, err := os.Create(tempFile)
136+
if err != nil {
134137
t.Fatalf("Failed to create temporary file: %v", err)
135138
}
136-
defer os.Remove(tempFile)
139+
f.Close()
140+
141+
// Clean up after the test
142+
defer func() {
143+
err := os.Remove(tempFile)
144+
if err != nil {
145+
t.Logf("Warning: Failed to remove temporary file: %v", err)
146+
}
147+
}()
137148

138149
// Test that the file exists
139150
if !fileExists(tempFile) {

0 commit comments

Comments
 (0)