Merge pull request #24 from chihaya/mock-tracker-storage

Mock Tracker Storage
This commit is contained in:
Jimmy Zelinskie 2013-11-30 19:33:57 -08:00
commit af6dbcc3cd
4 changed files with 299 additions and 1 deletions

View file

@ -0,0 +1,238 @@
// Copyright 2013 The Chihaya Authors. All rights reserved.
// Use of this source code is governed by the BSD 2-Clause license,
// which can be found in the LICENSE file.
package mock
import (
"github.com/pushrax/chihaya/storage"
"github.com/pushrax/chihaya/storage/tracker"
)
type Conn struct {
*Pool
}
func (c *Conn) FindUser(passkey string) (*storage.User, bool, error) {
c.usersM.RLock()
defer c.usersM.RUnlock()
user, ok := c.users[passkey]
if !ok {
return nil, false, nil
}
u := *user
return &u, true, nil
}
func (c *Conn) FindTorrent(infohash string) (*storage.Torrent, bool, error) {
c.torrentsM.RLock()
defer c.torrentsM.RUnlock()
torrent, ok := c.torrents[infohash]
if !ok {
return nil, false, nil
}
t := *torrent
return &t, true, nil
}
func (c *Conn) ClientWhitelisted(peerID string) (bool, error) {
c.whitelistM.RLock()
defer c.whitelistM.RUnlock()
_, ok := c.whitelist[peerID]
if !ok {
return false, nil
}
return true, nil
}
func (c *Conn) RecordSnatch(u *storage.User, t *storage.Torrent) error {
c.torrentsM.Lock()
defer c.torrentsM.Unlock()
torrent, ok := c.torrents[t.Infohash]
if !ok {
return tracker.ErrMissingResource
}
torrent.Snatches++
t.Snatches++
return nil
}
func (c *Conn) MarkActive(t *storage.Torrent) error {
c.torrentsM.Lock()
defer c.torrentsM.Unlock()
torrent, ok := c.torrents[t.Infohash]
if !ok {
return tracker.ErrMissingResource
}
torrent.Active = true
t.Active = true
return nil
}
func (c *Conn) MarkInactive(t *storage.Torrent) error {
c.torrentsM.Lock()
defer c.torrentsM.Unlock()
torrent, ok := c.torrents[t.Infohash]
if !ok {
return tracker.ErrMissingResource
}
torrent.Active = false
t.Active = false
return nil
}
func (c *Conn) AddLeecher(t *storage.Torrent, p *storage.Peer) error {
c.torrentsM.Lock()
defer c.torrentsM.Unlock()
torrent, ok := c.torrents[t.Infohash]
if !ok {
return tracker.ErrMissingResource
}
torrent.Leechers[storage.PeerMapKey(p)] = *p
t.Leechers[storage.PeerMapKey(p)] = *p
return nil
}
func (c *Conn) AddSeeder(t *storage.Torrent, p *storage.Peer) error {
c.torrentsM.Lock()
defer c.torrentsM.Unlock()
torrent, ok := c.torrents[t.Infohash]
if !ok {
return tracker.ErrMissingResource
}
torrent.Leechers[storage.PeerMapKey(p)] = *p
t.Leechers[storage.PeerMapKey(p)] = *p
return nil
}
func (c *Conn) RemoveLeecher(t *storage.Torrent, p *storage.Peer) error {
c.torrentsM.Lock()
defer c.torrentsM.Unlock()
torrent, ok := c.torrents[t.Infohash]
if !ok {
return tracker.ErrMissingResource
}
delete(torrent.Leechers, storage.PeerMapKey(p))
delete(t.Leechers, storage.PeerMapKey(p))
return nil
}
func (c *Conn) RemoveSeeder(t *storage.Torrent, p *storage.Peer) error {
c.torrentsM.Lock()
defer c.torrentsM.Unlock()
torrent, ok := c.torrents[t.Infohash]
if !ok {
return tracker.ErrMissingResource
}
delete(torrent.Seeders, storage.PeerMapKey(p))
delete(t.Seeders, storage.PeerMapKey(p))
return nil
}
func (c *Conn) SetLeecher(t *storage.Torrent, p *storage.Peer) error {
c.torrentsM.Lock()
defer c.torrentsM.Unlock()
torrent, ok := c.torrents[t.Infohash]
if !ok {
return tracker.ErrMissingResource
}
torrent.Leechers[storage.PeerMapKey(p)] = *p
t.Leechers[storage.PeerMapKey(p)] = *p
return nil
}
func (c *Conn) SetSeeder(t *storage.Torrent, p *storage.Peer) error {
c.torrentsM.Lock()
defer c.torrentsM.Unlock()
torrent, ok := c.torrents[t.Infohash]
if !ok {
return tracker.ErrMissingResource
}
torrent.Seeders[storage.PeerMapKey(p)] = *p
t.Seeders[storage.PeerMapKey(p)] = *p
return nil
}
func (c *Conn) IncrementSlots(u *storage.User) error {
c.usersM.Lock()
defer c.usersM.Unlock()
user, ok := c.users[u.Passkey]
if !ok {
return tracker.ErrMissingResource
}
user.Slots++
u.Slots++
return nil
}
func (c *Conn) DecrementSlots(u *storage.User) error {
c.usersM.Lock()
defer c.usersM.Unlock()
user, ok := c.users[u.Passkey]
if !ok {
return tracker.ErrMissingResource
}
user.Slots--
u.Slots--
return nil
}
func (c *Conn) LeecherFinished(t *storage.Torrent, p *storage.Peer) error {
c.torrentsM.Lock()
defer c.torrentsM.Unlock()
torrent, ok := c.torrents[t.Infohash]
if !ok {
return tracker.ErrMissingResource
}
torrent.Seeders[storage.PeerMapKey(p)] = *p
delete(torrent.Leechers, storage.PeerMapKey(p))
t.Seeders[storage.PeerMapKey(p)] = *p
delete(t.Leechers, storage.PeerMapKey(p))
return nil
}
func (c *Conn) AddTorrent(t *storage.Torrent) error {
c.torrentsM.Lock()
defer c.torrentsM.Unlock()
torrent := *t
c.torrents[t.Infohash] = &torrent
return nil
}
func (c *Conn) RemoveTorrent(t *storage.Torrent) error {
c.torrentsM.Lock()
defer c.torrentsM.Unlock()
delete(c.torrents, t.Infohash)
return nil
}
func (c *Conn) AddUser(u *storage.User) error {
c.usersM.Lock()
defer c.usersM.Unlock()
user := *u
c.users[u.Passkey] = &user
return nil
}
func (c *Conn) RemoveUser(u *storage.User) error {
c.usersM.Lock()
defer c.usersM.Unlock()
delete(c.users, u.Passkey)
return nil
}
func (c *Conn) WhitelistClient(peerID string) error {
c.whitelistM.Lock()
defer c.whitelistM.Unlock()
c.whitelist[peerID] = true
return nil
}
func (c *Conn) UnWhitelistClient(peerID string) error {
c.whitelistM.Lock()
defer c.whitelistM.Unlock()
delete(c.whitelist, peerID)
return nil
}

View file

@ -0,0 +1,24 @@
// Copyright 2013 The Chihaya Authors. All rights reserved.
// Use of this source code is governed by the BSD 2-Clause license,
// which can be found in the LICENSE file.
// Package mock implements the storage interface for a BitTorrent tracker
// within memory. It can be used in production, but isn't recommended.
// Stored values will not persist if the tracker is restarted.
package mock
import (
"github.com/pushrax/chihaya/config"
"github.com/pushrax/chihaya/storage"
"github.com/pushrax/chihaya/storage/tracker"
)
type driver struct{}
func (d *driver) New(conf *config.DataStore) tracker.Pool {
return &Pool{
users: make(map[string]*storage.User),
torrents: make(map[string]*storage.Torrent),
whitelist: make(map[string]bool),
}
}

View file

@ -0,0 +1,33 @@
// Copyright 2013 The Chihaya Authors. All rights reserved.
// Use of this source code is governed by the BSD 2-Clause license,
// which can be found in the LICENSE file.
package mock
import (
"sync"
"github.com/pushrax/chihaya/storage"
"github.com/pushrax/chihaya/storage/tracker"
)
type Pool struct {
users map[string]*storage.User
usersM sync.RWMutex
torrents map[string]*storage.Torrent
torrentsM sync.RWMutex
whitelist map[string]bool
whitelistM sync.RWMutex
}
func (p *Pool) Get() (tracker.Conn, error) {
return &Conn{
Pool: p,
}, nil
}
func (p *Pool) Close() error {
return nil
}

View file

@ -7,6 +7,7 @@
package tracker package tracker
import ( import (
"errors"
"fmt" "fmt"
"github.com/pushrax/chihaya/config" "github.com/pushrax/chihaya/config"
@ -14,7 +15,8 @@ import (
) )
var ( var (
drivers = make(map[string]Driver) ErrMissingResource = errors.New("tracker: resource missing")
drivers = make(map[string]Driver)
) )
type Driver interface { type Driver interface {
@ -65,6 +67,7 @@ type Conn interface {
// Writes // Writes
RecordSnatch(u *storage.User, t *storage.Torrent) error RecordSnatch(u *storage.User, t *storage.Torrent) error
MarkActive(t *storage.Torrent) error MarkActive(t *storage.Torrent) error
MarkInactive(t *storage.Torrent) error
AddLeecher(t *storage.Torrent, p *storage.Peer) error AddLeecher(t *storage.Torrent, p *storage.Peer) error
AddSeeder(t *storage.Torrent, p *storage.Peer) error AddSeeder(t *storage.Torrent, p *storage.Peer) error
RemoveLeecher(t *storage.Torrent, p *storage.Peer) error RemoveLeecher(t *storage.Torrent, p *storage.Peer) error