Fixing RedisRTC implementation and test

go/redisutil/rtc.go

 package redisutil
 
 import (
         "fmt"
-        "time"
+        "sync"
 
         "github.com/garyburd/redigo/redis"
         "github.com/skia-dev/glog"
+
         "go.skia.org/infra/go/rtcache"
         "go.skia.org/infra/go/util"
 )
 
 const (
         // QUEUE_PREFIX is the prefix for the task queue which is a sorted set in Redis.
         QUEUE_PREFIX = "rc:"
 )
 
 // RedisRTC implements a rtcache.ReadThroughCache with a Redis backend.
 type RedisRTC struct {
         // redisPool is the connection to redis this read-through-cache uses.
         redisPool *RedisPool
 
         // queueKey is the key of the work queue (sorted set) in Redis.
         queueKey string
 
         // inProgressKey is the key of the set of current tasks in progress.
         inProgressKey string
 
         // keyPrefix is the prefix of the results that are cached in redis.
         keyPrefix string
 
         // errKeyPrefix is the prefix of the error message if the task failed to
         // produce the desired resul.
         errKeyPrefix string
 
-        // queueSubChannel is the channel where keyspace notifications are generated
-        // when a new item is added to the task queue.
-        queueSubChannel string
+        // workReadyKey is the name of the queue that indicates that there is work
+        // available.
+        workReadyKey string
 
         // finishedChannel is the key of the PubSub channel where finished tasked are
         // announced.
         finishedChannel string
 
         // codec is used to serialize and desirialize the cached items.
         codec util.LRUCodec
 
         // worker is the function that is called to produce an item when it is not
         // cached.
         worker rtcache.ReadThroughFunc
 
         // doneCh is a channel to subscribe to be notified when a task is done.
         doneCh chan<- *doneSub
+
+        // shutdownCh is used to signal by closing it that all go-routines need to shut down.
+        shutdownCh chan bool
+
+        // shutdownWg is the WaitGroup used to synchronize the shutdown.
+        shutdownWg sync.WaitGroup
 }
 
 // workerTask is an auxiliary struct that contains the ID of a task
 // and it's priority.
 type workerTask struct {
         id       string
         priority int64
 }
 
 // doneSub is an auxiliary struct that is used by tasks to subscribe to be
@@ skipping 18 matching lines @@
         }
 
         queueKey := QUEUE_PREFIX + queueName
         keyPrefix := queueKey + ":k:"
         errKeyPrefix := queueKey + ":er:"
         finishedChannel := queueKey + ":ch"
 
         ret := &RedisRTC{
                 redisPool:       redisPool,
                 queueKey:        queueKey,
+                workReadyKey:    queueKey + ":wr",
                 inProgressKey:   queueKey + ":inp",
                 finishedChannel: finishedChannel,
                 keyPrefix:       keyPrefix,
                 errKeyPrefix:    errKeyPrefix,
-                queueSubChannel: fmt.Sprintf("__keyspace@%d__:%s", redisPool.db, queueKey),
                 codec:           codec,
                 worker:          worker,
+                shutdownCh:      make(chan bool),
         }
 
         // Start the feeder process.
         var err error
         if ret.doneCh, err = ret.startWorkScheduler(); err != nil {
                 return nil, err
         }
 
         // Start the workers if we have a worker specified.
         if worker != nil {
@@ skipping 18 matching lines @@
         return r.waitFor(id, priority, returnBytes)
 }
 
 func (r *RedisRTC) startWorkers(nWorkers int) error {
         workCh, err := r.getWorkChannel()
         if err != nil {
                 return err
         }
 
         for i := 0; i < nWorkers; i++ {
+                r.shutdownWg.Add(1)
                 go func() {
+                        var task *workerTask
+                WorkerLoop:
                         for {
-                                task := <-workCh
+                                // Wait for a task or the signal to shut down.
+                                select {
+                                case task = <-workCh:
+                                case <-r.shutdownCh:
+                                        break WorkerLoop
+                                }
+
                                 r.writeWorkerResult(task.priority, task.id)
                         }
+                        r.shutdownWg.Done()
                 }()
         }
 
         return nil
 }
 
+// shutdown terminates all go-routines started by this instance.
+func (r *RedisRTC) shutdown() {
+        close(r.shutdownCh)
+        r.shutdownWg.Wait()
+}
+
 // getWorkChannel returns a channel that sends tasks to be processed.
 func (r *RedisRTC) getWorkChannel() (<-chan *workerTask, error) {
         ret := make(chan *workerTask)
 
-        subCh, err := r.getKeyEventChannel(r.queueSubChannel, false)
-        if err != nil {
-                return nil, err
-        }
-
+        r.shutdownWg.Add(1)
         go func() {
-                hasMore := true
+                hasMore := false
+        WorkLoop:
                 for {
                         if !hasMore {
-                                <-subCh
+                                hasMore = r.checkForWork()
+                        }
+
+                        // Check if we need to shutdown.
+                        select {
+                        case <-r.shutdownCh:
+                                break WorkLoop
+                        default:
+                        }
+
+                        if !hasMore {
+                                continue
                         }
 
                         // Get the next task.
                         workTask, itemsLeft, err := r.dequeue()
                         if err != nil {
                                 glog.Errorf("Error moving work ids: %s", err)
                                 continue
                         }
 
                         hasMore = itemsLeft > 0
                         if workTask != nil {
                                 ret <- workTask
                         }
                 }
+                r.shutdownWg.Done()
         }()
 
         return ret, nil
 }
 
 // dequeue returns a task to be performed by the worker.
 // It returns a triple: workerTask, itemsLeft_in_the_task_queue, error.
 func (r *RedisRTC) dequeue() (*workerTask, int, error) {
         c := r.redisPool.Get()
         defer util.Close(c)
 
-        // TODO: this needs to be a transaction.
+        // Remove the first element in the transaction.
         util.LogErr(c.Send("MULTI"))
         util.LogErr(c.Send("ZRANGE", r.queueKey, 0, 0, "WITHSCORES"))
         util.LogErr(c.Send("ZREMRANGEBYRANK", r.queueKey, 0, 0))
         util.LogErr(c.Send("ZCARD", r.queueKey))
         combinedVals, err := redis.Values(c.Do("EXEC"))
 
         if err != nil {
                 return nil, 0, err
         }
 
@@ skipping 50 matching lines @@
         // whether we need to retry.
         if err != nil {
                 util.LogErr(c.Send("EXPIRE", writeKey, 10))
         }
         util.LogErr(c.Send("PUBLISH", r.finishedChannel, id))
         if _, err = c.Do("EXEC"); err != nil {
                 glog.Errorf("Error writing result to redis: %s", err)
         }
 }
 
-// getKeyEventChannel returns a channel that sends events for key changes in
-// Redis.
-func (r *RedisRTC) getKeyEventChannel(channelOrPattern string, isPattern bool) (<-chan string, error) {
-        // Listen for changes on the queue continously.
-        psc := redis.PubSubConn{Conn: r.redisPool.Get()}
+// waitForWork waits until there is either a work indicator in the queue
+// or 1 second has passed. It will only return false if there was an error.
+func (r *RedisRTC) checkForWork() bool {
+        c := r.redisPool.Get()
+        defer util.Close(c)
+        _, err := redis.Values(c.Do("BLPOP", r.workReadyKey, 1))
 
-        subscribe := func() error {
-                if isPattern {
-                        return psc.PSubscribe(channelOrPattern)
-                } else {
-                        return psc.Subscribe(channelOrPattern)
-                }
+        // If we timed out return true because there could be work and it will
+        // cause the caller to poll the queue.

[dogben, 2016/08/09 13:50:49]

Why checkForWork() at all? Why not just poll the queue continuously?

+        if err == redis.ErrNil {
+                return true
         }
 
-        // Subscribe to the key events
-        if err := subscribe(); err != nil {
-                return nil, err
+        // If there was an error.
+        if err != nil {
+                glog.Errorf("Error retrieving list: %s", err)
+                return false
         }
 
-        readyCh := make(chan bool)
-        ret := make(chan string)
-        go func() {
-                for {
-                Loop:
-                        for {
-                                switch v := psc.Receive().(type) {
-                                case redis.PMessage:
-                                        ret <- string(v.Data)
-                                case redis.Message:
-                                        ret <- string(v.Data)
-                                case redis.Subscription:
-                                        if readyCh != nil {
-                                                readyCh <- true
-                                                close(readyCh)
-                                        }
-                                case error:
-                                        glog.Errorf("Error waiting for key events: %s", v)
-                                        glog.Infof("Reconnecting.")
-                                        util.Close(psc)
-                                        break Loop
-                                }
-                        }
-
-                        readyCh = nil
-                        psc = redis.PubSubConn{Conn: r.redisPool.Get()}
-                        if err := subscribe(); err != nil {
-                                glog.Errorf("Error re-connecting: %s", err)
-                                time.Sleep(time.Second)
-                        }
-                }
-        }()
-        <-readyCh
-
-        return ret, nil
+        // Now we got something from the queue and should check for work.
+        return true
 }
 
 // enqueue adds the given task and priority to the task queue. Updating the
 // priority if necessary.
 func (r *RedisRTC) enqueue(id string, priority int64) (bool, error) {
         c := r.redisPool.Get()
         defer util.Close(c)
 
         util.LogErr(c.Send("MULTI"))
         util.LogErr(c.Send("ZSCORE", r.queueKey, id))
@@ skipping 14 matching lines @@
         if !isInProgress && !found {
                 saveId := true
 
                 // Only update the queue if this has a lower score.
                 if inQueueScore != nil {
                         oldPriority, _ := redis.Int64(inQueueScore, nil)
                         saveId = priority < oldPriority
                 }
 
                 if saveId {
-                        if _, err = c.Do("ZADD", r.queueKey, priority, id); err != nil {
+                        util.LogErr(c.Send("MULTI"))
+                        util.LogErr(c.Send("ZADD", r.queueKey, priority, id))
+                        util.LogErr(c.Send("RPUSH", r.workReadyKey, []byte("W")))
+                        if _, err = c.Do("EXEC"); err != nil {
                                 return false, err
                         }
                 }
         }
 
         return found, nil
 }
 
 // inQueue returns up to 'maxElements' ids that are currently in the work
 // queue. This is primarily for testing.
@@ skipping 32 matching lines @@
 
 // startWorkScheduler starts the background progress that takes tasks
 // from the doneCh and adds them to the Redis priority queue.
 func (r *RedisRTC) startWorkScheduler() (chan<- *doneSub, error) {
         finishedCh, err := r.redisPool.subscribeToChannel(r.finishedChannel)
         if err != nil {
                 return nil, err
         }
 
         doneCh := make(chan *doneSub)
+        r.shutdownWg.Add(1)
         go func() {
                 watchIds := map[string][]chan bool{}
                 notifyChannels := func(id string) {
                         for _, ch := range watchIds[id] {
                                 ch <- true
                                 close(ch)
                         }
                         delete(watchIds, id)
                 }
 
                 notifyAll := func() {
                         for id := range watchIds {
                                 if r.isFinished(id) {
                                         notifyChannels(id)
                                 }
                         }
                 }
 
+        WorkLoop:
                 for {
                         select {
+                        case <-r.shutdownCh:
+                                break WorkLoop
                         case subscription := <-doneCh:
                                 watchIds[subscription.id] = append(watchIds[subscription.id], subscription.notifyCh)
                                 if r.isFinished(subscription.id) {
                                         notifyChannels(subscription.id)
                                 }
                         case finishedId := <-finishedCh:
                                 // An emtpy string indicates that we (re)connected.
                                 if string(finishedId) == "" {
                                         notifyAll()
                                 } else {
                                         notifyChannels(string(finishedId))
                                 }
                         }
                 }
+                r.shutdownWg.Done()
         }()
 
         return doneCh, nil
 }
 
 // waitFor blocks until the key identified by id is available in Redis.
 func (r *RedisRTC) waitFor(id string, priority int64, returnBytes bool) (interface{}, error) {
         var found bool
         var err error
         if found, err = r.enqueue(id, priority); err != nil {
@@ skipping 50 matching lines @@
 // key returns the Redis key for the given id.
 func (r *RedisRTC) key(id string) string {
         return r.keyPrefix + id
 }
 
 // errorKey returns the key of the error message for the given ID if the
 // the worker call failed.
 func (r *RedisRTC) errorKey(id string) string {
         return r.errKeyPrefix + id
 }