Handle waiting for a redis instance to be available

Author: Edward Muller

Procfile

@@ -1,2 +1 @@
-web: go-websocket-chat-demo
-release: bin/addon-wait
+web: go-websocket-chat-demo

app.json

@@ -10,12 +10,7 @@
   "mount_dir": "src/github.com/heroku-examples/go-websocket-chat-demo",
   "website": "http://github.com/heroku-examples/go-websocket-chat-demo",
   "repository": "http://github.com/heroku-examples/go-websocket-chat-demol",
-  "scripts": {},
-  "buildpacks":[
-   { "url" : "https://github.com/heroku/heroku-buildpack-addon-wait.git" },
-   { "url" : "heroku/go" }
-  ],
   "addons": [
     "heroku-redis"
   ]
-}
+}

main.go

@@ -3,33 +3,62 @@ package main
 import (
 	"net/http"
 	"os"
+	"time"
 
 	"github.com/Sirupsen/logrus"
 	"github.com/heroku/gobits/redis"
 )
 
 var (
-	log = logrus.WithField("cmd", "go-websocket-chat-demo")
-	rr  redisReceiver
-	rw  redisWriter
+	waitTimeout = time.Minute * 10
+	log         = logrus.WithField("cmd", "go-websocket-chat-demo")
+	rr          redisReceiver
+	rw          redisWriter
 )
 
 func main() {
+	port := os.Getenv("PORT")
+	if port == "" {
+		log.WithField("PORT", port).Fatal("$PORT must be set")
+	}
+
 	redisURL := os.Getenv("REDIS_URL")
 	redisPool, err := redis.NewRedisPoolFromURL(redisURL)
 	if err != nil {
 		log.WithField("url", redisURL).Fatal("Unable to create Redis pool")
 	}
 
 	rr = newRedisReceiver(redisPool)
-	go rr.run()
 	rw = newRedisWriter(redisPool)
-	go rw.run()
 
-	port := os.Getenv("PORT")
-	if port == "" {
-		log.WithField("PORT", port).Fatal("$PORT must be set")
-	}
+	go func() {
+		for {
+			waited, err := redis.WaitForAvailability(redisURL, waitTimeout, rr.wait)
+			if !waited || err != nil {
+				log.WithFields(logrus.Fields{"waitTimeout": waitTimeout, "err": err}).Fatal("Redis not available by timeout!")
+			}
+			rr.broadcast(availableMessage)
+			err = rr.run()
+			if err == nil {
+				break
+			}
+			log.Error(err)
+		}
+	}()
+
+	go func() {
+		for {
+			waited, err := redis.WaitForAvailability(redisURL, waitTimeout, nil)
+			if !waited || err != nil {
+				log.WithFields(logrus.Fields{"waitTimeout": waitTimeout, "err": err}).Fatal("Redis not available by timeout!")
+			}
+			err = rw.run()
+			if err == nil {
+				break
+			}
+			log.Error(err)
+		}
+	}()
 
 	http.Handle("/", http.FileServer(http.Dir("./public")))
 	http.HandleFunc("/ws", handleWebsocket)

redis.go

@@ -1,11 +1,14 @@
 package main
 
 import (
+	"encoding/json"
 	"sync"
+	"time"
 
 	"github.com/Sirupsen/logrus"
 	"github.com/garyburd/redigo/redis"
 	"github.com/gorilla/websocket"
+	"github.com/pkg/errors"
 	"github.com/satori/go.uuid"
 )
 
@@ -14,12 +17,36 @@ const (
 	Channel = "chat"
 )
 
+var (
+	waitingMessage, availableMessage []byte
+	waitSleep                        = time.Second * 10
+)
+
+func init() {
+	var err error
+	waitingMessage, err = json.Marshal(message{
+		Handle: "system",
+		Text:   "Waiting for redis to be available. Messaging won't work until redis is available",
+	})
+	if err != nil {
+		panic(err)
+	}
+	availableMessage, err = json.Marshal(message{
+		Handle: "system",
+		Text:   "Redis is now available & messaging is now possible",
+	})
+	if err != nil {
+		panic(err)
+	}
+}
+
 // redisReceiver receives messages from Redis and broadcasts them to all
 // registered websocket connections that are Registered.
 type redisReceiver struct {
-	pool       *redis.Pool
-	sync.Mutex // Protects the conns map
-	conns      map[string]*websocket.Conn
+	pool *redis.Pool
+
+	mu    sync.Mutex
+	conns map[string]*websocket.Conn
 }
 
 // newRedisReceiver creates a redisReceiver that will use the provided
@@ -31,9 +58,15 @@ func newRedisReceiver(pool *redis.Pool) redisReceiver {
 	}
 }
 
+func (rr *redisReceiver) wait(_ time.Time) error {
+	rr.broadcast(waitingMessage)
+	time.Sleep(waitSleep)
+	return nil
+}
+
 // run receives pubsub messages from Redis after establishing a connection.
 // When a valid message is received it is broadcast to all connected websockets
-func (rr *redisReceiver) run() {
+func (rr *redisReceiver) run() error {
 	l := log.WithField("channel", Channel)
 	conn := rr.pool.Get()
 	defer conn.Close()
@@ -54,7 +87,7 @@ func (rr *redisReceiver) run() {
 				"count": v.Count,
 			}).Println("Redis Subscription Received")
 		case error:
-			l.WithField("err", v).Error("Error while subscribed to Redis channel")
+			return errors.Wrap(v, "Error while subscribed to Redis channel")
 		default:
 			l.WithField("v", v).Info("Unknown Redis receive during subscription")
 		}
@@ -65,8 +98,8 @@ func (rr *redisReceiver) run() {
 // If an error occurs while writting a message to a websocket connection it is
 // closed and deregistered.
 func (rr *redisReceiver) broadcast(data []byte) {
-	rr.Mutex.Lock()
-	defer rr.Mutex.Unlock()
+	rr.mu.Lock()
+	defer rr.mu.Unlock()
 	for id, conn := range rr.conns {
 		if err := conn.WriteMessage(websocket.TextMessage, data); err != nil {
 			log.WithFields(logrus.Fields{
@@ -84,17 +117,17 @@ func (rr *redisReceiver) broadcast(data []byte) {
 // identifier for the connection. This identifier can be used to deregister the
 // connection later
 func (rr *redisReceiver) register(conn *websocket.Conn) string {
-	rr.Mutex.Lock()
-	defer rr.Mutex.Unlock()
+	rr.mu.Lock()
+	defer rr.mu.Unlock()
 	id := uuid.NewV4().String()
 	rr.conns[id] = conn
 	return id
 }
 
 // deRegister the connection by closing it and removing it from our list.
 func (rr *redisReceiver) deRegister(id string) {
-	rr.Mutex.Lock()
-	defer rr.Mutex.Unlock()
+	rr.mu.Lock()
+	defer rr.mu.Unlock()
 	conn, ok := rr.conns[id]
 	if ok {
 		conn.Close()
@@ -111,24 +144,32 @@ type redisWriter struct {
 func newRedisWriter(pool *redis.Pool) redisWriter {
 	return redisWriter{
 		pool:     pool,
-		messages: make(chan []byte),
+		messages: make(chan []byte, 10000),
 	}
 }
 
 // run the main redisWriter loop that publishes incoming messages to Redis.
-func (rw *redisWriter) run() {
+func (rw *redisWriter) run() error {
 	conn := rw.pool.Get()
 	defer conn.Close()
 
 	for data := range rw.messages {
-		l := log.WithField("data", data)
-		if err := conn.Send("PUBLISH", Channel, data); err != nil {
-			l.WithField("err", err).Fatalf("Unable to publish message to Redis")
-		}
-		if err := conn.Flush(); err != nil {
-			l.WithField("err", err).Fatalf("Unable to flush published message to Redis")
+		if err := writeToRedis(conn, data); err != nil {
+			rw.publish(data) // attempt to redeliver later
+			return err
 		}
 	}
+	return nil
+}
+
+func writeToRedis(conn redis.Conn, data []byte) error {
+	if err := conn.Send("PUBLISH", Channel, data); err != nil {
+		return errors.Wrap(err, "Unable to publish message to Redis")
+	}
+	if err := conn.Flush(); err != nil {
+		return errors.Wrap(err, "Unable to flush published message to Redis")
+	}
+	return nil
 }
 
 // publish to Redis via channel.

vendor/github.com/heroku/gobits/redis/redis.go

@@ -32,10 +32,10 @@
 			"revisionTime": "2016-08-17T22:48:30Z"
 		},
 		{
-			"checksumSHA1": "KwTwEV5aWQbYFnbdsPWjV++9Y5c=",
+			"checksumSHA1": "G2M6++gheNWcKWclENJPEF7mtsI=",
 			"path": "github.com/heroku/gobits/redis",
-			"revision": "c55d8aab15a05920fbc6b9ef474122e2c0acee90",
-			"revisionTime": "2016-08-17T22:48:30Z"
+			"revision": "7ec1b023d6ef07ea214fa234ca9952ae01b9135a",
+			"revisionTime": "2016-08-18T00:10:40Z"
 		},
 		{
 			"checksumSHA1": "QoVjlQFru1ixgV8vh63T4/JAtLI=",