Define broker interface

heartbeat.go

@@ -9,14 +9,13 @@ import (
 	"time"
 
 	"github.com/hibiken/asynq/internal/base"
-	"github.com/hibiken/asynq/internal/rdb"
 )
 
 // heartbeater is responsible for writing process info to redis periodically to
 // indicate that the background worker process is up.
 type heartbeater struct {
 	logger Logger
-	rdb    *rdb.RDB
+	broker broker
 
 	ss *base.ServerState
 
@@ -27,10 +26,10 @@ type heartbeater struct {
 	interval time.Duration
 }
 
-func newHeartbeater(l Logger, rdb *rdb.RDB, ss *base.ServerState, interval time.Duration) *heartbeater {
+func newHeartbeater(l Logger, b broker, ss *base.ServerState, interval time.Duration) *heartbeater {
 	return &heartbeater{
 		logger:   l,
-		rdb:      rdb,
+		broker:   b,
 		ss:       ss,
 		done:     make(chan struct{}),
 		interval: interval,
@@ -53,7 +52,7 @@ func (h *heartbeater) start(wg *sync.WaitGroup) {
 		for {
 			select {
 			case <-h.done:
-				h.rdb.ClearServerState(h.ss)
+				h.broker.ClearServerState(h.ss)
 				h.logger.Info("Heartbeater done")
 				return
 			case <-time.After(h.interval):
@@ -66,7 +65,7 @@ func (h *heartbeater) start(wg *sync.WaitGroup) {
 func (h *heartbeater) beat() {
 	// Note: Set TTL to be long enough so that it won't expire before we write again
 	// and short enough to expire quickly once the process is shut down or killed.
-	err := h.rdb.WriteServerState(h.ss, h.interval*2)
+	err := h.broker.WriteServerState(h.ss, h.interval*2)
 	if err != nil {
 		h.logger.Error("could not write heartbeat data: %v", err)
 	}