Storage API matured, docs, tests & more

storage/redis/redis.go

@@ -2,6 +2,7 @@
 // Use of this source code is governed by the BSD 2-Clause license,
 // which can be found in the LICENSE file.
 
+// Package redis implements the storage interface for a BitTorrent tracker.
 package redis
 
 import (
@@ -15,12 +16,12 @@ import (
 
 type driver struct{}
 
-func (d *driver) New(conf *config.Storage) storage.Pool {
-	return &Pool{
+func (d *driver) New(conf *config.Storage) storage.DS {
+	return &DS{
 		conf: conf,
-		pool: redis.Pool{
-			MaxIdle:      3,
-			IdleTimeout:  240 * time.Second,
+		Pool: redis.Pool{
+			MaxIdle:      conf.MaxIdleConn,
+			IdleTimeout:  conf.IdleTimeout.Duration,
 			Dial:         makeDialFunc(conf),
 			TestOnBorrow: testOnBorrow,
 		},
@@ -34,16 +35,13 @@ func makeDialFunc(conf *config.Storage) func() (redis.Conn, error) {
 			err  error
 		)
 
-		if conf.ConnectTimeout != nil &&
-			conf.ReadTimeout != nil &&
-			conf.WriteTimeout != nil {
-
+		if conf.ConnTimeout != nil {
 			conn, err = redis.DialTimeout(
 				conf.Network,
 				conf.Addr,
-				conf.ConnectTimeout.Duration,
-				conf.ReadTimeout.Duration,
-				conf.WriteTimeout.Duration,
+				conf.ConnTimeout.Duration, // Connect Timeout
+				conf.ConnTimeout.Duration, // Read Timeout
+				conf.ConnTimeout.Duration, // Write Timeout
 			)
 		} else {
 			conn, err = redis.Dial(conf.Network, conf.Addr)
@@ -60,42 +58,26 @@ func testOnBorrow(c redis.Conn, t time.Time) error {
 	return err
 }
 
-type Pool struct {
+type DS struct {
 	conf *config.Storage
-	pool redis.Pool
+	redis.Pool
 }
 
-func (p *Pool) Get() storage.Conn {
-	return &Conn{
-		conf: p.conf,
-		Conn: p.pool.Get(),
-	}
-}
+func (ds *DS) FindUser(passkey string) (*storage.User, bool, error) {
+	conn := ds.Get()
+	defer conn.Close()
 
-func (p *Pool) Close() error {
-	return p.pool.Close()
-}
-
-type Conn struct {
-	conf *config.Storage
-	redis.Conn
-}
-
-func (c *Conn) FindUser(passkey string) (*storage.User, bool, error) {
-	key := c.conf.Prefix + "User:" + passkey
-
-	exists, err := redis.Bool(c.Do("EXISTS", key))
-	if err != nil {
-		return nil, false, err
-	}
-	if !exists {
-		return nil, false, nil
-	}
-
-	reply, err := redis.Values(c.Do("HGETALL", key))
+	key := ds.conf.Prefix + "user:" + passkey
+	reply, err := redis.Values(conn.Do("HGETALL", key))
 	if err != nil {
 		return nil, true, err
 	}
+
+	// If we get nothing back, the user isn't found.
+	if len(reply) == 0 {
+		return nil, false, nil
+	}
+
 	user := &storage.User{}
 	err = redis.ScanStruct(reply, user)
 	if err != nil {
@@ -104,21 +86,21 @@ func (c *Conn) FindUser(passkey string) (*storage.User, bool, error) {
 	return user, true, nil
 }
 
-func (c *Conn) FindTorrent(infohash string) (*storage.Torrent, bool, error) {
-	key := c.conf.Prefix + "Torrent:" + infohash
+func (ds *DS) FindTorrent(infohash string) (*storage.Torrent, bool, error) {
+	conn := ds.Get()
+	defer conn.Close()
 
-	exists, err := redis.Bool(c.Do("EXISTS", key))
+	key := ds.conf.Prefix + "torrent:" + infohash
+	reply, err := redis.Values(conn.Do("HGETALL", key))
 	if err != nil {
 		return nil, false, err
 	}
-	if !exists {
+
+	// If we get nothing back, the torrent isn't found.
+	if len(reply) == 0 {
 		return nil, false, nil
 	}
 
-	reply, err := redis.Values(c.Do("HGETALL", key))
-	if err != nil {
-		return nil, true, err
-	}
 	torrent := &storage.Torrent{}
 	err = redis.ScanStruct(reply, torrent)
 	if err != nil {
@@ -127,21 +109,50 @@ func (c *Conn) FindTorrent(infohash string) (*storage.Torrent, bool, error) {
 	return torrent, true, nil
 }
 
-type Tx struct {
-	conn *Conn
+func (ds *DS) ClientWhitelisted(peerID string) (bool, error) {
+	conn := ds.Get()
+	defer conn.Close()
+
+	key := ds.conf.Prefix + "whitelist:" + peerID
+	exists, err := redis.Bool(conn.Do("EXISTS", key))
+	if err != nil {
+		return false, err
+	}
+	return exists, nil
 }
 
-func (c *Conn) NewTx() (storage.Tx, error) {
-	err := c.Send("MULTI")
+type Tx struct {
+	conf *config.Storage
+	done bool
+	redis.Conn
+}
+
+func (ds *DS) Begin() (storage.Tx, error) {
+	conn := ds.Get()
+	err := conn.Send("MULTI")
 	if err != nil {
 		return nil, err
 	}
-	return &Tx{c}, nil
+	return &Tx{
+		conf: ds.conf,
+		Conn: conn,
+	}, nil
+}
+
+func (t *Tx) Close() {
+	if t.done {
+		panic("redis: transaction closed twice")
+	}
+	t.done = true
+	t.Conn.Close()
 }
 
 func (t *Tx) UnpruneTorrent(torrent *storage.Torrent) error {
-	key := t.conn.conf.Prefix + "Torrent:" + torrent.Infohash
-	err := t.conn.Send("HSET " + key + " Status 0")
+	if t.done {
+		return storage.ErrTxDone
+	}
+	key := t.conf.Prefix + "torrent:" + torrent.Infohash
+	err := t.Send("HSET " + key + " Status 0")
 	if err != nil {
 		return err
 	}
@@ -149,10 +160,24 @@ func (t *Tx) UnpruneTorrent(torrent *storage.Torrent) error {
 }
 
 func (t *Tx) Commit() error {
-	_, err := t.conn.Do("EXEC")
+	if t.done {
+		return storage.ErrTxDone
+	}
+	_, err := t.Do("EXEC")
 	if err != nil {
 		return err
 	}
+
+	t.Close()
+	return nil
+}
+
+// Redis doesn't need to rollback. Exec is atomic.
+func (t *Tx) Rollback() error {
+	if t.done {
+		return storage.ErrTxDone
+	}
+	t.Close()
 	return nil
 }
 

storage/storage.go

@@ -7,17 +7,24 @@
 package storage
 
 import (
+	"errors"
 	"fmt"
 
 	"github.com/pushrax/chihaya/config"
 )
 
-var drivers = make(map[string]Driver)
+var (
+	drivers   = make(map[string]Driver)
+	ErrTxDone = errors.New("storage: Transaction has already been committed or rolled back")
+)
 
 type Driver interface {
-	New(*config.Storage) Pool
+	New(*config.Storage) DS
 }
 
+// Register makes a database driver available by the provided name.
+// If Register is called twice with the same name or if driver is nil,
+// it panics.
 func Register(name string, driver Driver) {
 	if driver == nil {
 		panic("storage: Register driver is nil")
@@ -28,7 +35,8 @@ func Register(name string, driver Driver) {
 	drivers[name] = driver
 }
 
-func Open(conf *config.Storage) (Pool, error) {
+// Open opens a data store specified by a storage configuration.
+func Open(conf *config.Storage) (DS, error) {
 	driver, ok := drivers[conf.Driver]
 	if !ok {
 		return nil, fmt.Errorf(
@@ -40,25 +48,29 @@ func Open(conf *config.Storage) (Pool, error) {
 	return pool, nil
 }
 
-// ConnPool represents a pool of connections to the data store.
-type Pool interface {
-	Close() error
-	Get() Conn
-}
-
-// Conn represents a single connection to the data store.
-type Conn interface {
+// DS represents a data store handle. It's expected to be safe for concurrent
+// use by multiple goroutines.
+//
+// A pool of connections or a database/sql.DB is a great concrete type to
+// implement the DS interface.
+type DS interface {
 	Close() error
 
-	NewTx() (Tx, error)
+	Begin() (Tx, error)
 
 	FindUser(passkey string) (*User, bool, error)
 	FindTorrent(infohash string) (*Torrent, bool, error)
+	ClientWhitelisted(peerID string) (bool, error)
 }
 
-// Tx represents a data store transaction.
+// Tx represents an in-progress data store transaction.
+// A transaction must end with a call to Commit or Rollback.
+//
+// After a call to Commit or Rollback, all operations on the
+// transaction must fail with ErrTxDone.
 type Tx interface {
 	Commit() error
+	Rollback() error
 
 	UnpruneTorrent(torrent *Torrent) error
 }