Add Ability to Test and Reconnect to Data Store

redis.go

@@ -1,7 +1,7 @@
 package shield
 
 import (
-	"github.com/garyburd/redigo/redis"
+	"github.com/gomodule/redigo/redis"
 	"log"
 	"strconv"
 )
@@ -50,6 +50,19 @@ func (rs *RedisStore) conn() (conn redis.Conn, err error) {
 	return rs.redis, nil
 }
 
+func (rs *RedisStore) TestConnection() {
+	if c, err := rs.conn(); err == nil {
+		if _, err := c.Do("PING"); err == nil {
+			return
+		}
+	}
+
+	log.Println("Resetting redis connection...")
+
+	rs.redis = nil
+	rs.conn()
+}
+
 func (rs *RedisStore) Classes() (a []string, err error) {
 	c, err := rs.conn()
 	if err != nil {

shield.go

@@ -158,3 +158,7 @@ func (s *shield) Classify(text string) (class string, err error) {
 func (sh *shield) Reset() error {
 	return sh.store.Reset()
 }
+
+func (sh *shield) TestConnection() {
+	sh.store.TestConnection()
+}

types.go

@@ -27,6 +27,9 @@ type Shield interface {
 
 	// Reset clears the storage
 	Reset() error
+
+	// Checks and reestablishes connection to the data store
+	TestConnection()
 }
 
 type Store interface {
@@ -36,4 +39,5 @@ type Store interface {
 	IncrementClassWordCounts(m map[string]map[string]int64) error
 	TotalClassWordCounts() (map[string]int64, error)
 	Reset() error
+	TestConnection()
 }