mirror of
https://github.com/lukaszraczylo/traefikoidc.git
synced 2026-10-14 04:07:46 +00:00
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:
co-authored by
Copilot Autofix powered by AI <62310815+github-advanced-security[bot]@users.noreply.github.com>
parent
5fcbd54955
commit
e64fc7f730
@@ -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))
|
||||
}
|
||||
}
|
||||
@@ -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
|
||||
}
|
||||
@@ -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)),
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user