diff --git a/redis.go b/redis.go index d81fdef..e366a42 100644 --- a/redis.go +++ b/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 { diff --git a/shield.go b/shield.go index 6fcc5aa..397fdf8 100644 --- a/shield.go +++ b/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() +} diff --git a/types.go b/types.go index 9f337f8..6de1a6d 100644 --- a/types.go +++ b/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() }