pkg/storage/redis: do not use timeout to signal redis conn to stop (#1155)

Instead, we run the loop in goroutine, and when context was done,
closing the underlying connection of PubSubConn, so the Receive will
return.

Fixes #1154

pkg/storage/redis/redis.go

@@ -8,7 +8,6 @@ import (
 	"strconv"
 	"time"
 
-	"github.com/cenkalti/backoff/v4"
 	"github.com/golang/protobuf/proto"
 	"github.com/golang/protobuf/ptypes"
 	"github.com/gomodule/redigo/redis"
@@ -212,68 +211,48 @@ func (db *DB) ClearDeleted(_ context.Context, cutoff time.Time) {
 	}
 }
 
-// doNotify receives event from redis and signal the channel that something happened.
+// doNotifyLoop receives event from redis and send signal to the channel.
-func doNotify(ctx context.Context, psc *redis.PubSubConn, ch chan struct{}) error {
+func (db *DB) doNotifyLoop(ctx context.Context, ch chan struct{}, psc *redis.PubSubConn) {
-	switch v := psc.ReceiveWithTimeout(time.Second).(type) {
-	case redis.Message:
-		log.Debug().Str("action", string(v.Data)).Msg("got redis message")
-		if string(v.Data) != watchAction {
-			return nil
-		}
-		select {
-		case <-ctx.Done():
-			log.Warn().Err(ctx.Err()).Msg("unable to notify channel")
-			return ctx.Err()
-		case ch <- struct{}{}:
-		}
-	case error:
-		log.Debug().Err(v).Msg("redis subscribe error")
-		return v
-	}
-	return nil
-}
-
-// doNotifyLoop tries to run doNotify forever.
-//
-// Because redis.PubSubConn does not support context, so it will block until it receives event, we can not use
-// context to signal it stops. We mitigate this case by using PubSubConn.ReceiveWithTimeout. In case of timeout
-// occurred, we return a nil error, so the caller of doNotifyLoop will re-create new connection to start new loop.
-func (db *DB) doNotifyLoop(ctx context.Context, ch chan struct{}, psc *redis.PubSubConn, eb *backoff.ExponentialBackOff) error {
 	for {
-		err, ok := doNotify(ctx, psc, ch).(net.Error)
+		switch v := psc.Receive().(type) {
-		if !ok && err != nil {
+		case redis.Message:
-			log.Error().Err(ctx.Err()).Msg("failed to notify channel")
+			log.Debug().Str("action", string(v.Data)).Msg("got redis message")
-			return err
+			if string(v.Data) != watchAction {
-		}
+				continue
-		if ok && err.Timeout() {
+			}
 			select {
 			case <-ctx.Done():
-				return ctx.Err()
+				log.Warn().Err(ctx.Err()).Msg("context done, stop receive from redis channel")
-			case <-time.After(eb.NextBackOff()):
+				return
+			case ch <- struct{}{}:
+			}
+		case error:
+			log.Error().Err(v).Msg("failed to receive from redis channel")
+			if _, ok := v.(net.Error); ok {
+				return
 			}
-			return nil
 		}
 	}
 }
 
-// watchLoop runs the doNotifyLoop forever.
+// watch runs the doNotifyLoop. It returns when ctx was done or doNotifyLoop exits.
-//
+func (db *DB) watch(ctx context.Context, ch chan struct{}) {
-// If doNotifyLoop returns a nil error, watchLoop re-create the PubSubConn and start new iteration.
+	psConn := db.pool.Get()
-func (db *DB) watchLoop(ctx context.Context, ch chan struct{}) {
+	defer psConn.Close()
-	var psConn redis.Conn
+	psc := redis.PubSubConn{Conn: psConn}
-	eb := backoff.NewExponentialBackOff()
+	if err := psc.PSubscribe("__keyspace*__:" + db.versionSet); err != nil {
-	for {
+		log.Error().Err(err).Msg("failed to subscribe to version set channel")
-		psConn = db.pool.Get()
+		return
-		psc := redis.PubSubConn{Conn: psConn}
+	}
-		if err := psc.PSubscribe("__keyspace*__:" + db.versionSet); err != nil {
+
-			log.Error().Err(err).Msg("failed to subscribe to version set channel")
+	done := make(chan struct{})
-			psConn.Close()
+	go func() {
-			return
+		defer close(done)
-		}
+		db.doNotifyLoop(ctx, ch, &psc)
-		if err := db.doNotifyLoop(ctx, ch, &psc, eb); err != nil {
+	}()
-			psConn.Close()
+	select {
-			return
+	case <-ctx.Done():
-		}
+	case <-done:
 	}
 }
 
@@ -294,7 +273,7 @@ func (db *DB) Watch(ctx context.Context) chan struct{} {
 			return
 		}
 		c.Close()
-		db.watchLoop(ctx, ch)
+		db.watch(ctx, ch)
 	}()
 
 	return ch