Merge pull request #1 from muzzammilshahid/ping-manager

Implement ping manager

Author: muzzammilshahid

go.mod

@@ -1,3 +1,11 @@
 module github.com/xconnio/workerpool
 
 go 1.19
+
+require github.com/stretchr/testify v1.9.0
+
+require (
+	github.com/davecgh/go-spew v1.1.1 // indirect
+	github.com/pmezard/go-difflib v1.0.0 // indirect
+	gopkg.in/yaml.v3 v3.0.1 // indirect
+)

go.sum

@@ -0,0 +1,10 @@
+github.com/davecgh/go-spew v1.1.1 h1:vj9j/u1bqnvCEfJOwUhtlOARqs3+rkHYY13jYWTU97c=
+github.com/davecgh/go-spew v1.1.1/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38=
+github.com/pmezard/go-difflib v1.0.0 h1:4DBwDE0NGyQoBHbLQYPwSUPoCMWR5BEzIk/f1lZbAQM=
+github.com/pmezard/go-difflib v1.0.0/go.mod h1:iKH77koFhYxTK1pcRnkKkqfTogsbg7gZNVY4sRDYZ/4=
+github.com/stretchr/testify v1.9.0 h1:HtqpIVDClZ4nwg75+f6Lvsy/wHu+3BoSGCbBAcpTsTg=
+github.com/stretchr/testify v1.9.0/go.mod h1:r2ic/lqez/lEtzL7wO/rwa5dbSLXVDPFyf8C91i36aY=
+gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405 h1:yhCVgyC4o1eVCa2tZl7eS0r+SDo693bJlVdllGtEeKM=
+gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0=
+gopkg.in/yaml.v3 v3.0.1 h1:fxVm/GzAzEWqLHuvctI91KS9hhNmmWOoWu0XTYJS7CA=
+gopkg.in/yaml.v3 v3.0.1/go.mod h1:K4uyk7z7BCEPqu6E+C64Yfv1cQ7kz7rIZviUmN+EgEM=

workerpool.go

@@ -1 +1,116 @@
 package workerpool
+
+import (
+	"log"
+	"sync"
+	"time"
+)
+
+type wpConfig struct {
+	interval time.Duration
+	callback func() error
+}
+
+type WorkerPool struct {
+	workers            map[int64]map[int64]wpConfig
+	workersByIntervals map[int64]int64
+	mutex              sync.Mutex
+}
+
+// NewWorkerPool creates a new WorkerPool instance.
+func NewWorkerPool() *WorkerPool {
+	return &WorkerPool{
+		workers:            make(map[int64]map[int64]wpConfig),
+		workersByIntervals: make(map[int64]int64),
+	}
+}
+
+// Add registers a new worker.
+func (m *WorkerPool) Add(id int64, interval time.Duration, callback func() error) {
+	// Calculate the next worker time in seconds
+	timeToWork := time.Now().Add(interval).Unix()
+
+	m.mutex.Lock()
+	defer m.mutex.Unlock()
+	if m.workers[timeToWork] == nil {
+		m.workers[timeToWork] = make(map[int64]wpConfig)
+	}
+
+	m.workers[timeToWork][id] = wpConfig{
+		interval: interval,
+		callback: callback,
+	}
+	m.workersByIntervals[id] = timeToWork
+}
+
+// Remove unregisters a worker.
+func (m *WorkerPool) Remove(id int64) {
+	m.mutex.Lock()
+	defer m.mutex.Unlock()
+
+	nexTime, ok := m.workersByIntervals[id]
+	if ok {
+		delete(m.workers[nexTime], id)
+		if len(m.workers[nexTime]) == 0 {
+			delete(m.workers, nexTime)
+		}
+		delete(m.workersByIntervals, id)
+	}
+}
+
+// Reset reschedules the worker for a specific id.
+func (m *WorkerPool) Reset(id int64) {
+	m.mutex.Lock()
+	defer m.mutex.Unlock()
+
+	nextTime, ok := m.workersByIntervals[id]
+	if ok {
+		wpCfg := m.workers[nextTime][id]
+		delete(m.workers[nextTime], id)
+		if len(m.workers[nextTime]) == 0 {
+			delete(m.workers, nextTime)
+		}
+
+		newTime := time.Now().Add(wpCfg.interval).Unix()
+		if m.workers[newTime] == nil {
+			m.workers[newTime] = make(map[int64]wpConfig)
+		}
+		m.workers[newTime][id] = wpCfg
+		m.workersByIntervals[id] = newTime
+	}
+}
+
+// Start begins the worker process.
+func (m *WorkerPool) Start() {
+	go func() {
+		for {
+			nowSeconds := time.Now().Unix()
+
+			m.mutex.Lock()
+			workers, ok := m.workers[nowSeconds]
+			if ok {
+				for id, wpCfg := range workers {
+					err := wpCfg.callback()
+					if err != nil {
+						log.Printf("error pinging client: %v", err)
+						m.Remove(id)
+						continue
+					}
+
+					// Reschedule the next worker
+					timeToWork := time.Now().Add(wpCfg.interval).Unix()
+					if m.workers[timeToWork] == nil {
+						m.workers[timeToWork] = make(map[int64]wpConfig)
+					}
+					m.workers[timeToWork][id] = wpCfg
+					m.workersByIntervals[id] = timeToWork
+
+				}
+				delete(m.workers, nowSeconds)
+			}
+			m.mutex.Unlock()
+
+			time.Sleep(1 * time.Second)
+		}
+	}()
+}

workerpool_test.go

@@ -1,7 +1,69 @@
 package workerpool_test
 
-import "testing"
+import (
+	"testing"
+	"time"
 
-func TestHello(t *testing.T) {
+	"github.com/stretchr/testify/require"
 
+	"github.com/xconnio/workerpool"
+)
+
+func TestManager(t *testing.T) {
+	manager := workerpool.NewWorkerPool()
+
+	messages1 := 0
+	manager.Add(1, 2*time.Second, func() error {
+		messages1++
+		return nil
+	})
+
+	messages2 := 0
+	manager.Add(2, 1*time.Second, func() error {
+		messages2++
+		return nil
+	})
+
+	go manager.Start()
+
+	// Wait for worker
+	time.Sleep(3 * time.Second)
+
+	require.NotZero(t, messages1)
+	require.NotZero(t, messages2)
+
+	// Remove worker1 and check
+	manager.Remove(1)
+	time.Sleep(3 * time.Second)
+
+	require.Equal(t, messages1, 1)
+
+	require.Greater(t, messages2, 2)
+}
+
+func TestManagerReset(t *testing.T) {
+	manager := workerpool.NewWorkerPool()
+
+	messages := 0
+	manager.Add(1, 2*time.Second, func() error {
+		messages++
+		return nil
+	})
+
+	go manager.Start()
+
+	// Wait for first work
+	time.Sleep(3 * time.Second)
+	require.Equal(t, 1, messages)
+
+	time.Sleep(900 * time.Millisecond)
+	// Reset worker1 and check work
+	manager.Reset(1)
+
+	time.Sleep(1 * time.Second)
+	require.Equal(t, 1, messages)
+
+	// Wait for another work after reset
+	time.Sleep(2 * time.Second)
+	require.Equal(t, 2, messages)
 }