Add redis support for distributed caching (#83)

* Add redis support for distributed caching

* Move towards the self-provided Redis connection pool and RESP protocol implementation.
Official redis client library won't work with yaegi.

* fixup! Move towards the self-provided Redis connection pool and RESP protocol implementation. Official redis client library won't work with yaegi.

* fixup! fixup! Move towards the self-provided Redis connection pool and RESP protocol implementation. Official redis client library won't work with yaegi.

* fixup! fixup! fixup! Move towards the self-provided Redis connection pool and RESP protocol implementation. Official redis client library won't work with yaegi.

* fixup! fixup! fixup! fixup! Move towards the self-provided Redis connection pool and RESP protocol implementation. Official redis client library won't work with yaegi.

* fixup! fixup! fixup! fixup! fixup! Move towards the self-provided Redis connection pool and RESP protocol implementation. Official redis client library won't work with yaegi.

* ... and another all nighter.

* fixup! ... and another all nighter.

* fixup! fixup! ... and another all nighter.

* fixup! fixup! fixup! ... and another all nighter.

* Resolve issue #85 by adding ability to set custom claims in JWT tokens

* Remove redundant validation in auth middleware ( issue #89 )

* Add ability to set cookie prefix for session cookies ( #87 )

* fixup! Add ability to set cookie prefix for session cookies ( #87 )

* Add ability to set cookie max age - issue #91

* Potential fix for code scanning alert no. 10: Size computation for allocation may overflow

Co-authored-by: Copilot Autofix powered by AI <62310815+github-advanced-security[bot]@users.noreply.github.com>

* fixup! Merge main into 0.8.0-redis: resolve conflicts

---------

Co-authored-by: Copilot Autofix powered by AI <62310815+github-advanced-security[bot]@users.noreply.github.com>
This commit is contained in:
2025-11-30 02:18:46 +00:00
committed by GitHub
co-authored by Copilot Autofix powered by AI <62310815+github-advanced-security[bot]@users.noreply.github.com>
parent 5fcbd54955
commit e64fc7f730
318 changed files with 100989 additions and 948 deletions
+931
View File
@@ -0,0 +1,931 @@
//go:build !yaegi
package cleanup
import (
"sync"
"sync/atomic"
"testing"
"time"
)
// Mock logger for testing
type mockLogger struct {
mu sync.Mutex
logs []string
errLogs []string
debugLog []string
}
func (m *mockLogger) Logf(format string, args ...interface{}) {
m.mu.Lock()
defer m.mu.Unlock()
m.logs = append(m.logs, format)
}
func (m *mockLogger) ErrorLogf(format string, args ...interface{}) {
m.mu.Lock()
defer m.mu.Unlock()
m.errLogs = append(m.errLogs, format)
}
func (m *mockLogger) DebugLogf(format string, args ...interface{}) {
m.mu.Lock()
defer m.mu.Unlock()
m.debugLog = append(m.debugLog, format)
}
func (m *mockLogger) getLogCount() int {
m.mu.Lock()
defer m.mu.Unlock()
return len(m.logs)
}
// BackgroundTask tests
func TestNewBackgroundTask(t *testing.T) {
logger := &mockLogger{}
var wg sync.WaitGroup
runCount := 0
task := NewBackgroundTask("test-task", 100*time.Millisecond, func() {
runCount++
}, logger, &wg)
if task == nil {
t.Fatal("Expected NewBackgroundTask to return non-nil")
}
if task.name != "test-task" {
t.Errorf("Expected name 'test-task', got '%s'", task.name)
}
if task.interval != 100*time.Millisecond {
t.Errorf("Expected interval 100ms, got %v", task.interval)
}
if task.IsRunning() {
t.Error("Expected task to not be running initially")
}
}
func TestBackgroundTask_Start(t *testing.T) {
logger := &mockLogger{}
runCount := int32(0)
task := NewBackgroundTask("test-task", 50*time.Millisecond, func() {
atomic.AddInt32(&runCount, 1)
}, logger)
task.Start()
if !task.IsRunning() {
t.Error("Expected task to be running after Start()")
}
// Wait for at least 2 executions
time.Sleep(120 * time.Millisecond)
task.Stop()
count := atomic.LoadInt32(&runCount)
if count < 2 {
t.Errorf("Expected at least 2 executions, got %d", count)
}
}
func TestBackgroundTask_Stop(t *testing.T) {
logger := &mockLogger{}
task := NewBackgroundTask("test-task", 100*time.Millisecond, func() {}, logger)
task.Start()
time.Sleep(50 * time.Millisecond)
task.Stop()
if task.IsRunning() {
t.Error("Expected task to not be running after Stop()")
}
// Calling Stop again should not panic
task.Stop()
}
func TestBackgroundTask_DoubleStart(t *testing.T) {
logger := &mockLogger{}
task := NewBackgroundTask("test-task", 100*time.Millisecond, func() {}, logger)
task.Start()
logCountBefore := logger.getLogCount()
// Second start should be ignored
task.Start()
logCountAfter := logger.getLogCount()
if logCountAfter <= logCountBefore {
t.Error("Expected log message about task already running")
}
task.Stop()
}
func TestBackgroundTask_ExecuteWithPanic(t *testing.T) {
logger := &mockLogger{}
panicCount := int32(0)
task := NewBackgroundTask("panic-task", 50*time.Millisecond, func() {
count := atomic.AddInt32(&panicCount, 1)
if count == 1 {
panic("test panic")
}
}, logger)
task.Start()
time.Sleep(120 * time.Millisecond)
task.Stop()
// Task should recover from panic and continue
finalCount := atomic.LoadInt32(&panicCount)
if finalCount < 2 {
t.Errorf("Expected task to continue after panic, got %d executions", finalCount)
}
stats := task.GetStats()
if stats["errorCount"].(int64) < 1 {
t.Error("Expected error count to be at least 1")
}
}
func TestBackgroundTask_GetStats(t *testing.T) {
logger := &mockLogger{}
runCount := int32(0)
task := NewBackgroundTask("test-task", 50*time.Millisecond, func() {
atomic.AddInt32(&runCount, 1)
}, logger)
task.Start()
time.Sleep(120 * time.Millisecond)
task.Stop()
stats := task.GetStats()
if stats["name"] != "test-task" {
t.Errorf("Expected name 'test-task', got %v", stats["name"])
}
if !stats["isRunning"].(bool) == true {
// Task should be stopped
}
if stats["runCount"].(int64) < 2 {
t.Errorf("Expected runCount >= 2, got %v", stats["runCount"])
}
}
func TestBackgroundTask_WithWaitGroup(t *testing.T) {
logger := &mockLogger{}
var wg sync.WaitGroup
runCount := int32(0)
task := NewBackgroundTask("test-task", 50*time.Millisecond, func() {
atomic.AddInt32(&runCount, 1)
}, logger, &wg)
task.Start()
// Wait for task to start
time.Sleep(100 * time.Millisecond)
// Stop and wait
done := make(chan bool)
go func() {
task.Stop()
wg.Wait()
close(done)
}()
select {
case <-done:
// Success
case <-time.After(2 * time.Second):
t.Error("Timeout waiting for task to stop")
}
}
// TaskRegistry tests
func TestNewTaskRegistry(t *testing.T) {
logger := &mockLogger{}
registry := NewTaskRegistry(logger, 10)
if registry == nil {
t.Fatal("Expected NewTaskRegistry to return non-nil")
}
if registry.maxTasks != 10 {
t.Errorf("Expected maxTasks 10, got %d", registry.maxTasks)
}
if registry.GetTaskCount() != 0 {
t.Error("Expected initial task count to be 0")
}
}
func TestTaskRegistry_RegisterTask(t *testing.T) {
logger := &mockLogger{}
registry := NewTaskRegistry(logger, 10)
task := NewBackgroundTask("test-task", 100*time.Millisecond, func() {}, logger)
err := registry.RegisterTask("test-task", task)
if err != nil {
t.Errorf("Expected no error, got %v", err)
}
if registry.GetTaskCount() != 1 {
t.Error("Expected task count to be 1")
}
}
func TestTaskRegistry_RegisterTask_Duplicate(t *testing.T) {
logger := &mockLogger{}
registry := NewTaskRegistry(logger, 10)
task1 := NewBackgroundTask("test-task", 100*time.Millisecond, func() {}, logger)
task2 := NewBackgroundTask("test-task", 100*time.Millisecond, func() {}, logger)
err1 := registry.RegisterTask("test-task", task1)
if err1 != nil {
t.Errorf("Expected no error on first registration, got %v", err1)
}
err2 := registry.RegisterTask("test-task", task2)
if err2 == nil {
t.Error("Expected error when registering duplicate task")
}
}
func TestTaskRegistry_RegisterTask_Nil(t *testing.T) {
logger := &mockLogger{}
registry := NewTaskRegistry(logger, 10)
err := registry.RegisterTask("test-task", nil)
if err == nil {
t.Error("Expected error when registering nil task")
}
}
func TestTaskRegistry_RegisterTask_MaxLimit(t *testing.T) {
logger := &mockLogger{}
registry := NewTaskRegistry(logger, 2)
task1 := NewBackgroundTask("task1", 100*time.Millisecond, func() {}, logger)
task2 := NewBackgroundTask("task2", 100*time.Millisecond, func() {}, logger)
task3 := NewBackgroundTask("task3", 100*time.Millisecond, func() {}, logger)
registry.RegisterTask("task1", task1)
registry.RegisterTask("task2", task2)
err := registry.RegisterTask("task3", task3)
if err == nil {
t.Error("Expected error when exceeding max tasks")
}
}
func TestTaskRegistry_UnregisterTask(t *testing.T) {
logger := &mockLogger{}
registry := NewTaskRegistry(logger, 10)
task := NewBackgroundTask("test-task", 100*time.Millisecond, func() {}, logger)
registry.RegisterTask("test-task", task)
if registry.GetTaskCount() != 1 {
t.Error("Expected task count to be 1")
}
registry.UnregisterTask("test-task")
if registry.GetTaskCount() != 0 {
t.Error("Expected task count to be 0 after unregister")
}
}
func TestTaskRegistry_UnregisterTask_Running(t *testing.T) {
logger := &mockLogger{}
registry := NewTaskRegistry(logger, 10)
task := NewBackgroundTask("test-task", 100*time.Millisecond, func() {}, logger)
registry.RegisterTask("test-task", task)
task.Start()
time.Sleep(50 * time.Millisecond)
registry.UnregisterTask("test-task")
if task.IsRunning() {
t.Error("Expected task to be stopped after unregister")
}
}
func TestTaskRegistry_GetTask(t *testing.T) {
logger := &mockLogger{}
registry := NewTaskRegistry(logger, 10)
task := NewBackgroundTask("test-task", 100*time.Millisecond, func() {}, logger)
registry.RegisterTask("test-task", task)
retrieved, exists := registry.GetTask("test-task")
if !exists {
t.Error("Expected task to exist")
}
if retrieved != task {
t.Error("Expected to retrieve the same task")
}
_, exists = registry.GetTask("non-existent")
if exists {
t.Error("Expected non-existent task to not exist")
}
}
func TestTaskRegistry_StopAllTasks(t *testing.T) {
logger := &mockLogger{}
registry := NewTaskRegistry(logger, 10)
task1 := NewBackgroundTask("task1", 100*time.Millisecond, func() {}, logger)
task2 := NewBackgroundTask("task2", 100*time.Millisecond, func() {}, logger)
registry.RegisterTask("task1", task1)
registry.RegisterTask("task2", task2)
task1.Start()
task2.Start()
time.Sleep(50 * time.Millisecond)
registry.StopAllTasks()
if task1.IsRunning() || task2.IsRunning() {
t.Error("Expected all tasks to be stopped")
}
if registry.GetTaskCount() != 0 {
t.Error("Expected task count to be 0 after StopAllTasks")
}
}
func TestTaskRegistry_CreateSingletonTask(t *testing.T) {
logger := &mockLogger{}
registry := NewTaskRegistry(logger, 10)
runCount := int32(0)
task1, err1 := registry.CreateSingletonTask("singleton", 50*time.Millisecond, func() {
atomic.AddInt32(&runCount, 1)
}, logger)
if err1 != nil {
t.Errorf("Expected no error, got %v", err1)
}
if task1 == nil {
t.Fatal("Expected task to be created")
}
if !task1.IsRunning() {
t.Error("Expected task to be running")
}
// Try to create same task again
task2, err2 := registry.CreateSingletonTask("singleton", 50*time.Millisecond, func() {
atomic.AddInt32(&runCount, 1)
}, logger)
if err2 != nil {
t.Errorf("Expected no error on second call, got %v", err2)
}
if task2 != task1 {
t.Error("Expected to get the same task instance")
}
time.Sleep(120 * time.Millisecond)
task1.Stop()
if atomic.LoadInt32(&runCount) < 2 {
t.Error("Expected task to have run multiple times")
}
}
func TestTaskRegistry_GetAllTasks(t *testing.T) {
logger := &mockLogger{}
registry := NewTaskRegistry(logger, 10)
task1 := NewBackgroundTask("task1", 100*time.Millisecond, func() {}, logger)
task2 := NewBackgroundTask("task2", 100*time.Millisecond, func() {}, logger)
registry.RegisterTask("task1", task1)
registry.RegisterTask("task2", task2)
allTasks := registry.GetAllTasks()
if len(allTasks) != 2 {
t.Errorf("Expected 2 tasks, got %d", len(allTasks))
}
if _, ok := allTasks["task1"]; !ok {
t.Error("Expected task1 in results")
}
if _, ok := allTasks["task2"]; !ok {
t.Error("Expected task2 in results")
}
}
func TestTaskRegistry_GetStats(t *testing.T) {
logger := &mockLogger{}
registry := NewTaskRegistry(logger, 10)
task := NewBackgroundTask("test-task", 100*time.Millisecond, func() {}, logger)
registry.RegisterTask("test-task", task)
task.Start()
time.Sleep(50 * time.Millisecond)
stats := registry.GetStats()
if stats["totalTasks"].(int) != 1 {
t.Errorf("Expected totalTasks 1, got %v", stats["totalTasks"])
}
if stats["runningTasks"].(int) != 1 {
t.Errorf("Expected runningTasks 1, got %v", stats["runningTasks"])
}
if _, ok := stats["memory"]; !ok {
t.Error("Expected memory stats")
}
task.Stop()
}
func TestGlobalTaskRegistry(t *testing.T) {
// Reset before test
ResetGlobalTaskRegistry()
registry1 := GetGlobalTaskRegistry()
registry2 := GetGlobalTaskRegistry()
if registry1 != registry2 {
t.Error("Expected singleton to return same instance")
}
// Cleanup
ResetGlobalTaskRegistry()
}
func TestResetGlobalTaskRegistry(t *testing.T) {
ResetGlobalTaskRegistry()
registry := GetGlobalTaskRegistry()
logger := &mockLogger{}
task := NewBackgroundTask("test-task", 100*time.Millisecond, func() {}, logger)
registry.RegisterTask("test-task", task)
task.Start()
time.Sleep(50 * time.Millisecond)
ResetGlobalTaskRegistry()
// Should get a new instance
newRegistry := GetGlobalTaskRegistry()
if newRegistry.GetTaskCount() != 0 {
t.Error("Expected new registry to be empty")
}
}
// TaskCircuitBreaker tests
func TestNewTaskCircuitBreaker(t *testing.T) {
logger := &mockLogger{}
cb := NewTaskCircuitBreaker(5, 30*time.Second, logger)
if cb == nil {
t.Fatal("Expected NewTaskCircuitBreaker to return non-nil")
}
if cb.failureThreshold != 5 {
t.Errorf("Expected failureThreshold 5, got %d", cb.failureThreshold)
}
if cb.timeout != 30*time.Second {
t.Errorf("Expected timeout 30s, got %v", cb.timeout)
}
if cb.GetState() != CircuitBreakerClosed {
t.Error("Expected initial state to be closed")
}
}
func TestTaskCircuitBreaker_CanCreateTask(t *testing.T) {
logger := &mockLogger{}
cb := NewTaskCircuitBreaker(3, 100*time.Millisecond, logger)
err := cb.CanCreateTask("test-task")
if err != nil {
t.Errorf("Expected no error initially, got %v", err)
}
}
func TestTaskCircuitBreaker_OnTaskFailure(t *testing.T) {
logger := &mockLogger{}
cb := NewTaskCircuitBreaker(3, 100*time.Millisecond, logger)
// Record failures
for i := 0; i < 3; i++ {
cb.OnTaskFailure("test-task", nil)
}
// Circuit should be open
if cb.GetState() != CircuitBreakerOpen {
t.Error("Expected circuit breaker to be open after threshold failures")
}
// Should not be able to create task
err := cb.CanCreateTask("test-task")
if err == nil {
t.Error("Expected error when circuit breaker is open")
}
}
func TestTaskCircuitBreaker_OnTaskSuccess(t *testing.T) {
logger := &mockLogger{}
cb := NewTaskCircuitBreaker(5, 100*time.Millisecond, logger)
cb.OnTaskFailure("test-task", nil)
cb.OnTaskFailure("test-task", nil)
cb.OnTaskSuccess("test-task")
// Task-specific failures should be reset
err := cb.CanCreateTask("test-task")
if err != nil {
t.Errorf("Expected no error after success, got %v", err)
}
}
func TestTaskCircuitBreaker_Reset(t *testing.T) {
logger := &mockLogger{}
cb := NewTaskCircuitBreaker(2, 100*time.Millisecond, logger)
cb.OnTaskFailure("test-task", nil)
cb.OnTaskFailure("test-task", nil)
if cb.GetState() != CircuitBreakerOpen {
t.Error("Expected circuit breaker to be open")
}
cb.Reset()
if cb.GetState() != CircuitBreakerClosed {
t.Error("Expected circuit breaker to be closed after reset")
}
err := cb.CanCreateTask("test-task")
if err != nil {
t.Errorf("Expected no error after reset, got %v", err)
}
}
func TestTaskCircuitBreaker_TimeoutRecovery(t *testing.T) {
logger := &mockLogger{}
cb := NewTaskCircuitBreaker(2, 100*time.Millisecond, logger)
// Open circuit breaker
cb.OnTaskFailure("test-task", nil)
cb.OnTaskFailure("test-task", nil)
if cb.GetState() != CircuitBreakerOpen {
t.Error("Expected circuit breaker to be open")
}
// Wait for timeout
time.Sleep(150 * time.Millisecond)
// Circuit breaker should reset, but task-specific failures remain
// Need to check with a different task name
err := cb.CanCreateTask("different-task")
if err != nil {
t.Errorf("Expected no error for different task after timeout, got %v", err)
}
if cb.GetState() != CircuitBreakerClosed {
t.Error("Expected circuit breaker to be closed after timeout")
}
// Original task still has too many failures
err = cb.CanCreateTask("test-task")
if err == nil {
t.Error("Expected error for original task with too many failures")
}
}
// TaskMemoryMonitor tests
func TestNewTaskMemoryMonitor(t *testing.T) {
logger := &mockLogger{}
registry := NewTaskRegistry(logger, 10)
monitor := NewTaskMemoryMonitor(logger, registry)
if monitor == nil {
t.Fatal("Expected NewTaskMemoryMonitor to return non-nil")
}
if monitor.registry != registry {
t.Error("Expected registry to be set")
}
if monitor.memoryThreshold != 1024*1024*1024 {
t.Errorf("Expected default threshold 1GB, got %d", monitor.memoryThreshold)
}
}
func TestTaskMemoryMonitor_SetMemoryThreshold(t *testing.T) {
logger := &mockLogger{}
registry := NewTaskRegistry(logger, 10)
monitor := NewTaskMemoryMonitor(logger, registry)
monitor.SetMemoryThreshold(512 * 1024 * 1024)
stats := monitor.GetStats()
if stats["memoryThreshold"].(uint64) != 512*1024*1024 {
t.Error("Expected threshold to be updated")
}
}
func TestTaskMemoryMonitor_StartStop(t *testing.T) {
logger := &mockLogger{}
registry := NewTaskRegistry(logger, 10)
monitor := NewTaskMemoryMonitor(logger, registry)
monitor.StartMonitoring()
stats := monitor.GetStats()
if !stats["isMonitoring"].(bool) {
t.Error("Expected monitor to be running")
}
// Double start should be ignored
monitor.StartMonitoring()
monitor.StopMonitoring()
stats = monitor.GetStats()
if stats["isMonitoring"].(bool) {
t.Error("Expected monitor to be stopped")
}
// Double stop should be safe
monitor.StopMonitoring()
}
func TestTaskMemoryMonitor_GetStats(t *testing.T) {
logger := &mockLogger{}
registry := NewTaskRegistry(logger, 10)
monitor := NewTaskMemoryMonitor(logger, registry)
stats := monitor.GetStats()
if _, ok := stats["isMonitoring"]; !ok {
t.Error("Expected isMonitoring in stats")
}
if _, ok := stats["currentMemory"]; !ok {
t.Error("Expected currentMemory in stats")
}
if _, ok := stats["memoryThreshold"]; !ok {
t.Error("Expected memoryThreshold in stats")
}
}
// WorkerPool tests
func TestNewWorkerPool(t *testing.T) {
logger := &mockLogger{}
pool := NewWorkerPool(4, 10, logger)
if pool == nil {
t.Fatal("Expected NewWorkerPool to return non-nil")
}
if pool.workers != 4 {
t.Errorf("Expected 4 workers, got %d", pool.workers)
}
}
func TestWorkerPool_DefaultWorkers(t *testing.T) {
logger := &mockLogger{}
pool := NewWorkerPool(0, 0, logger)
// Should default to NumCPU
if pool.workers <= 0 {
t.Error("Expected positive number of workers")
}
}
func TestWorkerPool_StartStop(t *testing.T) {
logger := &mockLogger{}
pool := NewWorkerPool(2, 5, logger)
pool.Start()
metrics := pool.GetMetrics()
if !metrics["isRunning"].(bool) {
t.Error("Expected worker pool to be running")
}
// Double start should be ignored
pool.Start()
pool.Stop()
metrics = pool.GetMetrics()
if metrics["isRunning"].(bool) {
t.Error("Expected worker pool to be stopped")
}
// Double stop should be safe
pool.Stop()
}
func TestWorkerPool_Submit(t *testing.T) {
logger := &mockLogger{}
pool := NewWorkerPool(2, 5, logger)
pool.Start()
defer pool.Stop()
executed := int32(0)
var wg sync.WaitGroup
for i := 0; i < 3; i++ {
wg.Add(1)
err := pool.Submit(func() {
defer wg.Done()
atomic.AddInt32(&executed, 1)
})
if err != nil {
t.Errorf("Expected no error submitting task, got %v", err)
}
}
// Wait for tasks to complete
done := make(chan bool)
go func() {
wg.Wait()
close(done)
}()
select {
case <-done:
// Success
case <-time.After(2 * time.Second):
t.Error("Timeout waiting for tasks to complete")
}
if atomic.LoadInt32(&executed) != 3 {
t.Errorf("Expected 3 tasks executed, got %d", atomic.LoadInt32(&executed))
}
}
func TestWorkerPool_SubmitWhenStopped(t *testing.T) {
logger := &mockLogger{}
pool := NewWorkerPool(2, 5, logger)
err := pool.Submit(func() {})
if err == nil {
t.Error("Expected error when submitting to stopped pool")
}
}
func TestWorkerPool_TaskPanic(t *testing.T) {
logger := &mockLogger{}
pool := NewWorkerPool(2, 5, logger)
pool.Start()
defer pool.Stop()
executed := int32(0)
var wg sync.WaitGroup
wg.Add(2)
// Submit task that panics
pool.Submit(func() {
defer wg.Done()
panic("test panic")
})
// Submit normal task
pool.Submit(func() {
defer wg.Done()
atomic.AddInt32(&executed, 1)
})
// Wait for tasks
done := make(chan bool)
go func() {
wg.Wait()
close(done)
}()
select {
case <-done:
// Success
case <-time.After(2 * time.Second):
t.Error("Timeout waiting for tasks")
}
// Pool should still be functional
metrics := pool.GetMetrics()
if metrics["tasksFailed"].(int64) < 1 {
t.Error("Expected at least one failed task")
}
}
func TestWorkerPool_GetMetrics(t *testing.T) {
logger := &mockLogger{}
pool := NewWorkerPool(2, 5, logger)
pool.Start()
defer pool.Stop()
var wg sync.WaitGroup
wg.Add(2)
pool.Submit(func() {
defer wg.Done()
time.Sleep(10 * time.Millisecond)
})
pool.Submit(func() {
defer wg.Done()
time.Sleep(10 * time.Millisecond)
})
wg.Wait()
metrics := pool.GetMetrics()
if metrics["workers"].(int) != 2 {
t.Errorf("Expected 2 workers, got %v", metrics["workers"])
}
if metrics["tasksProcessed"].(int64) != 2 {
t.Errorf("Expected 2 processed tasks, got %v", metrics["tasksProcessed"])
}
if metrics["tasksQueued"].(int64) != 2 {
t.Errorf("Expected 2 queued tasks, got %v", metrics["tasksQueued"])
}
}
func TestWorkerPool_Concurrent(t *testing.T) {
logger := &mockLogger{}
pool := NewWorkerPool(4, 20, logger)
pool.Start()
defer pool.Stop()
executed := int32(0)
var wg sync.WaitGroup
taskCount := 10
for i := 0; i < taskCount; i++ {
wg.Add(1)
err := pool.Submit(func() {
defer wg.Done()
atomic.AddInt32(&executed, 1)
time.Sleep(10 * time.Millisecond)
})
if err != nil {
wg.Done()
t.Errorf("Failed to submit task: %v", err)
}
}
// Wait for all tasks
done := make(chan bool)
go func() {
wg.Wait()
close(done)
}()
select {
case <-done:
// Success
case <-time.After(5 * time.Second):
t.Error("Timeout waiting for concurrent tasks")
}
if atomic.LoadInt32(&executed) != int32(taskCount) {
t.Errorf("Expected %d tasks executed, got %d", taskCount, atomic.LoadInt32(&executed))
}
}
+407
View File
@@ -0,0 +1,407 @@
// Package cleanup provides background task management and cleanup functionality.
package cleanup
import (
"context"
"fmt"
"runtime"
"sync"
"sync/atomic"
"time"
)
// Logger defines the logging interface
type Logger interface {
Logf(format string, args ...interface{})
ErrorLogf(format string, args ...interface{})
DebugLogf(format string, args ...interface{})
}
// BackgroundTask represents a recurring background task
type BackgroundTask struct {
name string
interval time.Duration
taskFunc func()
ticker *time.Ticker
stopChan chan bool
isRunning int32
logger Logger
waitGroup *sync.WaitGroup
lastRun time.Time
runCount int64
errorCount int64
mu sync.RWMutex
ctx context.Context
cancelFunc context.CancelFunc
}
// NewBackgroundTask creates a new background task
func NewBackgroundTask(name string, interval time.Duration, taskFunc func(), logger Logger, wg ...*sync.WaitGroup) *BackgroundTask {
var waitGroup *sync.WaitGroup
if len(wg) > 0 && wg[0] != nil {
waitGroup = wg[0]
}
ctx, cancel := context.WithCancel(context.Background())
return &BackgroundTask{
name: name,
interval: interval,
taskFunc: taskFunc,
stopChan: make(chan bool, 1),
isRunning: 0,
logger: logger,
waitGroup: waitGroup,
ctx: ctx,
cancelFunc: cancel,
}
}
// Start begins executing the background task
func (bt *BackgroundTask) Start() {
if !atomic.CompareAndSwapInt32(&bt.isRunning, 0, 1) {
if bt.logger != nil {
bt.logger.Logf("Background task %s is already running", bt.name)
}
return
}
bt.ticker = time.NewTicker(bt.interval)
if bt.waitGroup != nil {
bt.waitGroup.Add(1)
}
go bt.run()
if bt.logger != nil {
bt.logger.Logf("Started background task: %s (interval: %v)", bt.name, bt.interval)
}
}
// Stop stops the background task
func (bt *BackgroundTask) Stop() {
if !atomic.CompareAndSwapInt32(&bt.isRunning, 1, 0) {
if bt.logger != nil {
bt.logger.Logf("Background task %s is not running", bt.name)
}
return
}
// Cancel context
if bt.cancelFunc != nil {
bt.cancelFunc()
}
// Stop ticker
if bt.ticker != nil {
bt.ticker.Stop()
}
// Send stop signal
select {
case bt.stopChan <- true:
case <-time.After(5 * time.Second):
if bt.logger != nil {
bt.logger.ErrorLogf("Timeout stopping background task: %s", bt.name)
}
}
if bt.logger != nil {
bt.logger.Logf("Stopped background task: %s", bt.name)
}
}
// run is the main loop for the background task
func (bt *BackgroundTask) run() {
defer func() {
if bt.waitGroup != nil {
bt.waitGroup.Done()
}
if r := recover(); r != nil {
atomic.AddInt64(&bt.errorCount, 1)
if bt.logger != nil {
bt.logger.ErrorLogf("Background task %s panicked: %v", bt.name, r)
}
}
}()
// Run task immediately on start
bt.executeTask()
for {
select {
case <-bt.ticker.C:
bt.executeTask()
case <-bt.stopChan:
return
case <-bt.ctx.Done():
return
}
}
}
// executeTask runs the task function with error handling
func (bt *BackgroundTask) executeTask() {
defer func() {
if r := recover(); r != nil {
atomic.AddInt64(&bt.errorCount, 1)
if bt.logger != nil {
bt.logger.ErrorLogf("Task %s panicked: %v", bt.name, r)
}
}
}()
bt.mu.Lock()
bt.lastRun = time.Now()
bt.mu.Unlock()
atomic.AddInt64(&bt.runCount, 1)
bt.taskFunc()
}
// GetStats returns statistics about the task
func (bt *BackgroundTask) GetStats() map[string]interface{} {
bt.mu.RLock()
lastRun := bt.lastRun
bt.mu.RUnlock()
return map[string]interface{}{
"name": bt.name,
"interval": bt.interval.String(),
"isRunning": atomic.LoadInt32(&bt.isRunning) == 1,
"lastRun": lastRun.Format(time.RFC3339),
"runCount": atomic.LoadInt64(&bt.runCount),
"errorCount": atomic.LoadInt64(&bt.errorCount),
}
}
// IsRunning returns whether the task is currently running
func (bt *BackgroundTask) IsRunning() bool {
return atomic.LoadInt32(&bt.isRunning) == 1
}
// TaskRegistry manages all background tasks
type TaskRegistry struct {
tasks map[string]*BackgroundTask
mu sync.RWMutex
logger Logger
maxTasks int
circuitBreaker *TaskCircuitBreaker
}
// globalTaskRegistry is the singleton task registry
var (
globalTaskRegistry *TaskRegistry
registryOnce sync.Once
registryMutex sync.Mutex
)
// GetGlobalTaskRegistry returns the global task registry singleton
func GetGlobalTaskRegistry() *TaskRegistry {
registryOnce.Do(func() {
globalTaskRegistry = &TaskRegistry{
tasks: make(map[string]*BackgroundTask),
maxTasks: 100, // Default maximum tasks
}
})
return globalTaskRegistry
}
// ResetGlobalTaskRegistry resets the global task registry (mainly for testing)
func ResetGlobalTaskRegistry() {
registryMutex.Lock()
defer registryMutex.Unlock()
if globalTaskRegistry != nil {
globalTaskRegistry.StopAllTasks()
globalTaskRegistry = nil
}
registryOnce = sync.Once{}
}
// NewTaskRegistry creates a new task registry
func NewTaskRegistry(logger Logger, maxTasks int) *TaskRegistry {
return &TaskRegistry{
tasks: make(map[string]*BackgroundTask),
logger: logger,
maxTasks: maxTasks,
circuitBreaker: NewTaskCircuitBreaker(5, 30*time.Second, logger),
}
}
// RegisterTask registers a new background task
func (tr *TaskRegistry) RegisterTask(name string, task *BackgroundTask) error {
if task == nil {
return fmt.Errorf("task cannot be nil")
}
tr.mu.Lock()
defer tr.mu.Unlock()
// Check if task already exists
if _, exists := tr.tasks[name]; exists {
return fmt.Errorf("task with name %s already exists", name)
}
// Check task limit
if len(tr.tasks) >= tr.maxTasks {
return fmt.Errorf("maximum number of tasks (%d) reached", tr.maxTasks)
}
// Check circuit breaker
if tr.circuitBreaker != nil {
if err := tr.circuitBreaker.CanCreateTask(name); err != nil {
return err
}
}
tr.tasks[name] = task
if tr.logger != nil {
tr.logger.Logf("Registered task: %s", name)
}
return nil
}
// UnregisterTask removes a task from the registry
func (tr *TaskRegistry) UnregisterTask(name string) {
tr.mu.Lock()
defer tr.mu.Unlock()
if task, exists := tr.tasks[name]; exists {
if task.IsRunning() {
task.Stop()
}
delete(tr.tasks, name)
if tr.logger != nil {
tr.logger.Logf("Unregistered task: %s", name)
}
}
}
// GetTask returns a task by name
func (tr *TaskRegistry) GetTask(name string) (*BackgroundTask, bool) {
tr.mu.RLock()
defer tr.mu.RUnlock()
task, exists := tr.tasks[name]
return task, exists
}
// StopAllTasks stops all registered tasks
func (tr *TaskRegistry) StopAllTasks() {
tr.mu.RLock()
tasks := make([]*BackgroundTask, 0, len(tr.tasks))
for _, task := range tr.tasks {
tasks = append(tasks, task)
}
tr.mu.RUnlock()
var wg sync.WaitGroup
for _, task := range tasks {
if task.IsRunning() {
wg.Add(1)
go func(t *BackgroundTask) {
defer wg.Done()
t.Stop()
}(task)
}
}
wg.Wait()
// Clear all tasks from the registry after stopping them
tr.mu.Lock()
tr.tasks = make(map[string]*BackgroundTask)
tr.mu.Unlock()
if tr.logger != nil {
tr.logger.Logf("Stopped all tasks")
}
}
// GetTaskCount returns the number of registered tasks
func (tr *TaskRegistry) GetTaskCount() int {
tr.mu.RLock()
defer tr.mu.RUnlock()
return len(tr.tasks)
}
// CreateSingletonTask creates or retrieves an existing task
func (tr *TaskRegistry) CreateSingletonTask(name string, interval time.Duration,
taskFunc func(), logger Logger, wg ...*sync.WaitGroup) (*BackgroundTask, error) {
// Check if task already exists
if existingTask, exists := tr.GetTask(name); exists {
if existingTask.IsRunning() {
if logger != nil {
logger.Logf("Task %s already exists and is running", name)
}
return existingTask, nil
}
// Task exists but not running, start it
existingTask.Start()
return existingTask, nil
}
// Create new task
task := NewBackgroundTask(name, interval, taskFunc, logger, wg...)
// Register task
if err := tr.RegisterTask(name, task); err != nil {
return nil, err
}
// Start task
task.Start()
return task, nil
}
// GetAllTasks returns all registered tasks
func (tr *TaskRegistry) GetAllTasks() map[string]*BackgroundTask {
tr.mu.RLock()
defer tr.mu.RUnlock()
tasks := make(map[string]*BackgroundTask)
for name, task := range tr.tasks {
tasks[name] = task
}
return tasks
}
// GetStats returns statistics for all tasks
func (tr *TaskRegistry) GetStats() map[string]interface{} {
tr.mu.RLock()
defer tr.mu.RUnlock()
stats := make(map[string]interface{})
stats["totalTasks"] = len(tr.tasks)
runningCount := 0
taskStats := make(map[string]interface{})
for name, task := range tr.tasks {
if task.IsRunning() {
runningCount++
}
taskStats[name] = task.GetStats()
}
stats["runningTasks"] = runningCount
stats["tasks"] = taskStats
// Add memory stats
var m runtime.MemStats
runtime.ReadMemStats(&m)
stats["memory"] = map[string]interface{}{
"alloc": m.Alloc,
"totalAlloc": m.TotalAlloc,
"sys": m.Sys,
"numGC": m.NumGC,
"goroutines": runtime.NumGoroutine(),
}
return stats
}
+449
View File
@@ -0,0 +1,449 @@
// Package cleanup provides background task management and cleanup functionality.
package cleanup
import (
"fmt"
"runtime"
"sync"
"sync/atomic"
"time"
)
// TaskCircuitBreaker prevents task creation failures from cascading
type TaskCircuitBreaker struct {
failureThreshold int32
failureCount int32
lastFailureTime time.Time
timeout time.Duration
state int32 // 0: closed, 1: open
logger Logger
mu sync.RWMutex
taskFailures map[string]int32
}
// CircuitBreakerState represents the state of the circuit breaker
type CircuitBreakerState int32
const (
CircuitBreakerClosed CircuitBreakerState = iota
CircuitBreakerOpen
)
// NewTaskCircuitBreaker creates a new circuit breaker for task management
func NewTaskCircuitBreaker(failureThreshold int32, timeout time.Duration, logger Logger) *TaskCircuitBreaker {
return &TaskCircuitBreaker{
failureThreshold: failureThreshold,
timeout: timeout,
logger: logger,
taskFailures: make(map[string]int32),
}
}
// CanCreateTask checks if a new task can be created
func (cb *TaskCircuitBreaker) CanCreateTask(taskName string) error {
cb.mu.RLock()
defer cb.mu.RUnlock()
// Check circuit breaker state
if atomic.LoadInt32(&cb.state) == int32(CircuitBreakerOpen) {
// Check if timeout has elapsed
if time.Since(cb.lastFailureTime) < cb.timeout {
return fmt.Errorf("circuit breaker open: too many task failures")
}
// Reset circuit breaker
atomic.StoreInt32(&cb.state, int32(CircuitBreakerClosed))
atomic.StoreInt32(&cb.failureCount, 0)
if cb.logger != nil {
cb.logger.Logf("Circuit breaker reset after timeout")
}
}
// Check task-specific failures
if failures, exists := cb.taskFailures[taskName]; exists {
if failures >= cb.failureThreshold {
return fmt.Errorf("task %s has too many failures (%d)", taskName, failures)
}
}
return nil
}
// OnTaskStart records that a task has started
func (cb *TaskCircuitBreaker) OnTaskStart(taskName string) {
// Currently just for tracking, could add rate limiting here
if cb.logger != nil {
cb.logger.DebugLogf("Task %s started", taskName)
}
}
// OnTaskComplete records that a task completed (success or failure)
func (cb *TaskCircuitBreaker) OnTaskComplete(taskName string) {
// Currently just for tracking
if cb.logger != nil {
cb.logger.DebugLogf("Task %s completed", taskName)
}
}
// OnTaskSuccess records a successful task execution
func (cb *TaskCircuitBreaker) OnTaskSuccess(taskName string) {
cb.mu.Lock()
defer cb.mu.Unlock()
// Reset task-specific failure count on success
delete(cb.taskFailures, taskName)
}
// OnTaskFailure records a task failure
func (cb *TaskCircuitBreaker) OnTaskFailure(taskName string, err error) {
cb.mu.Lock()
defer cb.mu.Unlock()
// Increment task-specific failure count
cb.taskFailures[taskName]++
// Increment overall failure count
failures := atomic.AddInt32(&cb.failureCount, 1)
cb.lastFailureTime = time.Now()
if cb.logger != nil {
cb.logger.ErrorLogf("Task %s failed: %v (failure count: %d)", taskName, err, cb.taskFailures[taskName])
}
// Open circuit breaker if threshold reached
if failures >= cb.failureThreshold {
atomic.StoreInt32(&cb.state, int32(CircuitBreakerOpen))
if cb.logger != nil {
cb.logger.ErrorLogf("Circuit breaker opened due to %d failures", failures)
}
}
}
// Reset resets the circuit breaker
func (cb *TaskCircuitBreaker) Reset() {
cb.mu.Lock()
defer cb.mu.Unlock()
atomic.StoreInt32(&cb.state, int32(CircuitBreakerClosed))
atomic.StoreInt32(&cb.failureCount, 0)
cb.taskFailures = make(map[string]int32)
cb.lastFailureTime = time.Time{}
if cb.logger != nil {
cb.logger.Logf("Circuit breaker reset")
}
}
// GetState returns the current state of the circuit breaker
func (cb *TaskCircuitBreaker) GetState() CircuitBreakerState {
return CircuitBreakerState(atomic.LoadInt32(&cb.state))
}
// TaskMemoryMonitor monitors memory usage and can trigger cleanup
type TaskMemoryMonitor struct {
logger Logger
registry *TaskRegistry
memoryThreshold uint64
checkInterval time.Duration
isMonitoring int32
stopChan chan bool
lastCheck time.Time
mu sync.RWMutex
}
var (
globalMemoryMonitor *TaskMemoryMonitor
monitorOnce sync.Once
)
// GetGlobalTaskMemoryMonitor returns the global memory monitor singleton
func GetGlobalTaskMemoryMonitor(logger Logger) *TaskMemoryMonitor {
monitorOnce.Do(func() {
globalMemoryMonitor = NewTaskMemoryMonitor(logger, GetGlobalTaskRegistry())
})
return globalMemoryMonitor
}
// NewTaskMemoryMonitor creates a new memory monitor
func NewTaskMemoryMonitor(logger Logger, registry *TaskRegistry) *TaskMemoryMonitor {
return &TaskMemoryMonitor{
logger: logger,
registry: registry,
memoryThreshold: 1024 * 1024 * 1024, // 1GB default
checkInterval: 1 * time.Minute,
stopChan: make(chan bool, 1),
}
}
// SetMemoryThreshold sets the memory threshold for triggering cleanup
func (tmm *TaskMemoryMonitor) SetMemoryThreshold(bytes uint64) {
tmm.mu.Lock()
defer tmm.mu.Unlock()
tmm.memoryThreshold = bytes
}
// StartMonitoring starts the memory monitoring routine
func (tmm *TaskMemoryMonitor) StartMonitoring() {
if !atomic.CompareAndSwapInt32(&tmm.isMonitoring, 0, 1) {
if tmm.logger != nil {
tmm.logger.Logf("Memory monitor is already running")
}
return
}
go tmm.monitorLoop()
if tmm.logger != nil {
tmm.logger.Logf("Started memory monitoring (threshold: %d bytes, interval: %v)",
tmm.memoryThreshold, tmm.checkInterval)
}
}
// StopMonitoring stops the memory monitoring routine
func (tmm *TaskMemoryMonitor) StopMonitoring() {
if !atomic.CompareAndSwapInt32(&tmm.isMonitoring, 1, 0) {
if tmm.logger != nil {
tmm.logger.Logf("Memory monitor is not running")
}
return
}
select {
case tmm.stopChan <- true:
case <-time.After(5 * time.Second):
if tmm.logger != nil {
tmm.logger.ErrorLogf("Timeout stopping memory monitor")
}
}
if tmm.logger != nil {
tmm.logger.Logf("Stopped memory monitoring")
}
}
// monitorLoop is the main monitoring loop
func (tmm *TaskMemoryMonitor) monitorLoop() {
ticker := time.NewTicker(tmm.checkInterval)
defer ticker.Stop()
for {
select {
case <-ticker.C:
tmm.checkMemory()
case <-tmm.stopChan:
return
}
}
}
// checkMemory checks current memory usage and triggers cleanup if needed
func (tmm *TaskMemoryMonitor) checkMemory() {
tmm.mu.Lock()
tmm.lastCheck = time.Now()
tmm.mu.Unlock()
var m runtime.MemStats
runtime.ReadMemStats(&m)
if tmm.logger != nil {
tmm.logger.DebugLogf("Memory check - Alloc: %d MB, Sys: %d MB, NumGC: %d",
m.Alloc/1024/1024, m.Sys/1024/1024, m.NumGC)
}
// Check if memory usage exceeds threshold
if m.Alloc > tmm.memoryThreshold {
if tmm.logger != nil {
tmm.logger.Logf("Memory usage (%d MB) exceeds threshold (%d MB), triggering cleanup",
m.Alloc/1024/1024, tmm.memoryThreshold/1024/1024)
}
// Trigger garbage collection
runtime.GC()
// Could also trigger task-specific cleanup here
tmm.triggerTaskCleanup()
}
}
// triggerTaskCleanup triggers cleanup operations on tasks
func (tmm *TaskMemoryMonitor) triggerTaskCleanup() {
if tmm.registry == nil {
return
}
// Get all tasks and potentially pause non-critical ones
tasks := tmm.registry.GetAllTasks()
for name, task := range tasks {
// Could implement task priority here
if tmm.logger != nil {
tmm.logger.DebugLogf("Checking task %s for cleanup opportunities", name)
}
// Tasks could implement a Cleanup() method
_ = task // Placeholder for future cleanup logic
}
}
// GetStats returns memory monitor statistics
func (tmm *TaskMemoryMonitor) GetStats() map[string]interface{} {
tmm.mu.RLock()
lastCheck := tmm.lastCheck
tmm.mu.RUnlock()
var m runtime.MemStats
runtime.ReadMemStats(&m)
return map[string]interface{}{
"isMonitoring": atomic.LoadInt32(&tmm.isMonitoring) == 1,
"lastCheck": lastCheck.Format(time.RFC3339),
"checkInterval": tmm.checkInterval.String(),
"memoryThreshold": tmm.memoryThreshold,
"currentMemory": map[string]interface{}{
"alloc": m.Alloc,
"totalAlloc": m.TotalAlloc,
"sys": m.Sys,
"mallocs": m.Mallocs,
"frees": m.Frees,
"numGC": m.NumGC,
"goroutines": runtime.NumGoroutine(),
},
}
}
// WorkerPool manages a pool of worker goroutines for task execution
type WorkerPool struct {
workers int
taskQueue chan func()
workerWg sync.WaitGroup
isRunning int32
logger Logger
stopChan chan bool
metrics WorkerPoolMetrics
}
// WorkerPoolMetrics tracks worker pool performance
type WorkerPoolMetrics struct {
tasksProcessed int64
tasksQueued int64
tasksFailed int64
avgProcessTime int64 // nanoseconds
}
// NewWorkerPool creates a new worker pool
func NewWorkerPool(workers int, queueSize int, logger Logger) *WorkerPool {
if workers <= 0 {
workers = runtime.NumCPU()
}
if queueSize <= 0 {
queueSize = workers * 10
}
return &WorkerPool{
workers: workers,
taskQueue: make(chan func(), queueSize),
stopChan: make(chan bool),
logger: logger,
}
}
// Start starts the worker pool
func (wp *WorkerPool) Start() {
if !atomic.CompareAndSwapInt32(&wp.isRunning, 0, 1) {
if wp.logger != nil {
wp.logger.Logf("Worker pool is already running")
}
return
}
for i := 0; i < wp.workers; i++ {
wp.workerWg.Add(1)
go wp.worker(i)
}
if wp.logger != nil {
wp.logger.Logf("Started worker pool with %d workers", wp.workers)
}
}
// Stop stops the worker pool
func (wp *WorkerPool) Stop() {
if !atomic.CompareAndSwapInt32(&wp.isRunning, 1, 0) {
if wp.logger != nil {
wp.logger.Logf("Worker pool is not running")
}
return
}
close(wp.stopChan)
close(wp.taskQueue)
wp.workerWg.Wait()
if wp.logger != nil {
wp.logger.Logf("Stopped worker pool")
}
}
// Submit submits a task to the worker pool
func (wp *WorkerPool) Submit(task func()) error {
if atomic.LoadInt32(&wp.isRunning) != 1 {
return fmt.Errorf("worker pool is not running")
}
select {
case wp.taskQueue <- task:
atomic.AddInt64(&wp.metrics.tasksQueued, 1)
return nil
default:
return fmt.Errorf("worker pool queue is full")
}
}
// worker is the main worker routine
func (wp *WorkerPool) worker(id int) {
defer wp.workerWg.Done()
for {
select {
case task, ok := <-wp.taskQueue:
if !ok {
return // Channel closed
}
wp.executeTask(task)
case <-wp.stopChan:
return
}
}
}
// executeTask executes a task with error handling
func (wp *WorkerPool) executeTask(task func()) {
startTime := time.Now()
defer func() {
if r := recover(); r != nil {
atomic.AddInt64(&wp.metrics.tasksFailed, 1)
if wp.logger != nil {
wp.logger.ErrorLogf("Worker pool task panicked: %v", r)
}
}
// Update average process time
duration := time.Since(startTime).Nanoseconds()
processed := atomic.AddInt64(&wp.metrics.tasksProcessed, 1)
currentAvg := atomic.LoadInt64(&wp.metrics.avgProcessTime)
newAvg := (currentAvg*(processed-1) + duration) / processed
atomic.StoreInt64(&wp.metrics.avgProcessTime, newAvg)
}()
task()
}
// GetMetrics returns worker pool metrics
func (wp *WorkerPool) GetMetrics() map[string]interface{} {
return map[string]interface{}{
"workers": wp.workers,
"isRunning": atomic.LoadInt32(&wp.isRunning) == 1,
"queueSize": len(wp.taskQueue),
"queueCapacity": cap(wp.taskQueue),
"tasksProcessed": atomic.LoadInt64(&wp.metrics.tasksProcessed),
"tasksQueued": atomic.LoadInt64(&wp.metrics.tasksQueued),
"tasksFailed": atomic.LoadInt64(&wp.metrics.tasksFailed),
"avgProcessTime": time.Duration(atomic.LoadInt64(&wp.metrics.avgProcessTime)),
}
}