mirror of
https://github.com/lukaszraczylo/traefikoidc.git
synced 2026-10-14 04:07:46 +00:00
... and another all nighter.
This commit is contained in:
@@ -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