initial middleware refactor

server/store/client_store.go

@@ -0,0 +1,49 @@
+// Copyright 2016 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 store
+
+import "fmt"
+
+var clientStoreDrivers = make(map[string]ClientStoreDriver)
+
+// ClientStore represents an interface for manipulating clientIDs.
+type ClientStore interface {
+	CreateClient(clientID string) error
+	FindClient(peerID string) (bool, error)
+	DeleteClient(clientID string) error
+}
+
+// ClientStoreDriver represents an interface for creating a handle to the
+// storage of swarms.
+type ClientStoreDriver interface {
+	New(*Config) (ClientStore, error)
+}
+
+// RegisterClientStoreDriver makes a driver available by the provided name.
+//
+// If this function is called twice with the same name or if the driver is nil,
+// it panics.
+func RegisterClientStoreDriver(name string, driver ClientStoreDriver) {
+	if driver == nil {
+		panic("store: could not register nil ClientStoreDriver")
+	}
+	if _, dup := clientStoreDrivers[name]; dup {
+		panic("store: could not register duplicate ClientStoreDriver: " + name)
+	}
+	clientStoreDrivers[name] = driver
+}
+
+// OpenClientStore returns a ClientStore specified by a configuration.
+func OpenClientStore(name string, cfg *Config) (ClientStore, error) {
+	driver, ok := clientStoreDrivers[name]
+	if !ok {
+		return nil, fmt.Errorf(
+			"store: unknown driver %q (forgotten import?)",
+			name,
+		)
+	}
+
+	return driver.New(cfg)
+}

server/store/memory/client_store.go

@@ -0,0 +1,59 @@
+// Copyright 2016 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 memory
+
+import (
+	"sync"
+
+	"github.com/chihaya/chihaya/pkg/clientid"
+	"github.com/chihaya/chihaya/server/store"
+)
+
+func init() {
+	store.RegisterClientStoreDriver("memory", &clientStoreDriver{})
+}
+
+type clientStoreDriver struct{}
+
+func (d *clientStoreDriver) New(cfg *store.Config) (store.ClientStore, error) {
+	return &clientStore{
+		clientIDs: make(map[string]struct{}),
+	}, nil
+}
+
+type clientStore struct {
+	clientIDs map[string]struct{}
+	sync.RWMutex
+}
+
+var _ store.ClientStore = &clientStore{}
+
+func (s *clientStore) CreateClient(clientID string) error {
+	s.Lock()
+	defer s.Unlock()
+
+	s.clientIDs[clientID] = struct{}{}
+
+	return nil
+}
+
+func (s *clientStore) FindClient(peerID string) (bool, error) {
+	clientID := clientid.New(peerID)
+	s.Lock()
+	defer s.Unlock()
+
+	_, ok := s.clientIDs[clientID]
+
+	return ok, nil
+}
+
+func (s *clientStore) DeleteClient(clientID string) error {
+	s.Lock()
+	defer s.Unlock()
+
+	delete(s.clientIDs, clientID)
+
+	return nil
+}

server/store/memory/peer_store.go

@@ -0,0 +1,272 @@
+// Copyright 2016 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 memory
+
+import (
+	"hash/fnv"
+	"runtime"
+	"sync"
+	"time"
+
+	"gopkg.in/yaml.v2"
+
+	"github.com/chihaya/chihaya"
+	"github.com/chihaya/chihaya/server/store"
+)
+
+func init() {
+	store.RegisterPeerStoreDriver("memory", &peerStoreDriver{})
+}
+
+type peerStoreDriver struct{}
+
+func (d *peerStoreDriver) New(storecfg *store.Config) (store.PeerStore, error) {
+	cfg, err := newPeerStoreConfig(storecfg)
+	if err != nil {
+		return nil, err
+	}
+
+	return &peerStore{
+		shards: make([]*peerShard, cfg.Shards),
+	}, nil
+}
+
+type peerStoreConfig struct {
+	Shards int `yaml:"shards"`
+}
+
+func newPeerStoreConfig(storecfg *store.Config) (*peerStoreConfig, error) {
+	bytes, err := yaml.Marshal(storecfg.PeerStoreConfig)
+	if err != nil {
+		return nil, err
+	}
+
+	var cfg peerStoreConfig
+	err = yaml.Unmarshal(bytes, &cfg)
+	if err != nil {
+		return nil, err
+	}
+
+	return &cfg, nil
+}
+
+const seedersSuffix = "-seeders"
+const leechersSuffix = "-leechers"
+
+type peer struct {
+	chihaya.Peer
+	LastAction time.Time
+}
+
+type peerShard struct {
+	peers map[string]map[string]peer
+	sync.RWMutex
+}
+
+type peerStore struct {
+	shards []*peerShard
+}
+
+var _ store.PeerStore = &peerStore{}
+
+func (s *peerStore) shardIndex(infohash string) uint32 {
+	idx := fnv.New32()
+	idx.Write([]byte(infohash))
+	return idx.Sum32() % uint32(len(s.shards))
+}
+
+func (s *peerStore) PutSeeder(infohash string, p chihaya.Peer) error {
+	key := infohash + seedersSuffix
+
+	shard := s.shards[s.shardIndex(infohash)]
+	shard.Lock()
+	defer shard.Unlock()
+
+	if shard.peers[key] == nil {
+		shard.peers[key] = make(map[string]peer)
+	}
+
+	shard.peers[key][p.ID] = peer{
+		Peer:       p,
+		LastAction: time.Now(),
+	}
+
+	return nil
+}
+
+func (s *peerStore) DeleteSeeder(infohash, peerID string) error {
+	key := infohash + seedersSuffix
+
+	shard := s.shards[s.shardIndex(infohash)]
+	shard.Lock()
+	defer shard.Unlock()
+
+	if shard.peers[key] == nil {
+		return nil
+	}
+
+	delete(shard.peers[key], peerID)
+
+	if len(shard.peers[key]) == 0 {
+		shard.peers[key] = nil
+	}
+
+	return nil
+}
+
+func (s *peerStore) PutLeecher(infohash string, p chihaya.Peer) error {
+	key := infohash + leechersSuffix
+
+	shard := s.shards[s.shardIndex(infohash)]
+	shard.Lock()
+	defer shard.Unlock()
+
+	if shard.peers[key] == nil {
+		shard.peers[key] = make(map[string]peer)
+	}
+
+	shard.peers[key][p.ID] = peer{
+		Peer:       p,
+		LastAction: time.Now(),
+	}
+
+	return nil
+}
+
+func (s *peerStore) DeleteLeecher(infohash, peerID string) error {
+	key := infohash + leechersSuffix
+
+	shard := s.shards[s.shardIndex(infohash)]
+	shard.Lock()
+	defer shard.Unlock()
+
+	if shard.peers[key] == nil {
+		return nil
+	}
+
+	delete(shard.peers[key], peerID)
+
+	if len(shard.peers[key]) == 0 {
+		shard.peers[key] = nil
+	}
+
+	return nil
+}
+
+func (s *peerStore) GraduateLeecher(infohash string, p chihaya.Peer) error {
+	leecherKey := infohash + leechersSuffix
+	seederKey := infohash + seedersSuffix
+
+	shard := s.shards[s.shardIndex(infohash)]
+	shard.Lock()
+	defer shard.Unlock()
+
+	if shard.peers[leecherKey] != nil {
+		delete(shard.peers[leecherKey], p.ID)
+	}
+
+	if shard.peers[seederKey] == nil {
+		shard.peers[seederKey] = make(map[string]peer)
+	}
+
+	shard.peers[seederKey][p.ID] = peer{
+		Peer:       p,
+		LastAction: time.Now(),
+	}
+
+	return nil
+}
+
+func (s *peerStore) CollectGarbage(cutoff time.Time) error {
+	for _, shard := range s.shards {
+		shard.RLock()
+		var keys []string
+		for key := range shard.peers {
+			keys = append(keys, key)
+		}
+		shard.RUnlock()
+		runtime.Gosched()
+
+		for _, key := range keys {
+			shard.Lock()
+			var peersToDelete []string
+			for peerID, p := range shard.peers[key] {
+				if p.LastAction.Before(cutoff) {
+					peersToDelete = append(peersToDelete, peerID)
+				}
+			}
+
+			for _, peerID := range peersToDelete {
+				delete(shard.peers[key], peerID)
+			}
+			shard.Unlock()
+			runtime.Gosched()
+		}
+
+		runtime.Gosched()
+	}
+
+	return nil
+}
+
+func (s *peerStore) AnnouncePeers(infohash string, seeder bool, numWant int) (peers, peers6 []chihaya.Peer, err error) {
+	leecherKey := infohash + leechersSuffix
+	seederKey := infohash + seedersSuffix
+
+	shard := s.shards[s.shardIndex(infohash)]
+	shard.RLock()
+	defer shard.RUnlock()
+
+	if seeder {
+		// Append leechers as possible.
+		leechers := shard.peers[leecherKey]
+		for _, p := range leechers {
+			if numWant == 0 {
+				break
+			}
+
+			if p.IP.To4() == nil {
+				peers6 = append(peers6, p.Peer)
+			} else {
+				peers = append(peers, p.Peer)
+			}
+			numWant--
+		}
+	} else {
+		// Append as many seeders as possible.
+		seeders := shard.peers[seederKey]
+		for _, p := range seeders {
+			if numWant == 0 {
+				break
+			}
+
+			if p.IP.To4() == nil {
+				peers6 = append(peers6, p.Peer)
+			} else {
+				peers = append(peers, p.Peer)
+			}
+			numWant--
+		}
+
+		// Append leechers until we reach numWant.
+		leechers := shard.peers[leecherKey]
+		if numWant > 0 {
+			for _, p := range leechers {
+				if numWant == 0 {
+					break
+				}
+
+				if p.IP.To4() == nil {
+					peers6 = append(peers6, p.Peer)
+				} else {
+					peers = append(peers, p.Peer)
+				}
+				numWant--
+			}
+		}
+	}
+
+	return
+}

server/store/peer_store.go

@@ -0,0 +1,62 @@
+// Copyright 2016 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 store
+
+import (
+	"fmt"
+	"time"
+
+	"github.com/chihaya/chihaya"
+)
+
+var peerStoreDrivers = make(map[string]PeerStoreDriver)
+
+// PeerStore represents an interface for manipulating peers.
+type PeerStore interface {
+	PutSeeder(infohash string, p chihaya.Peer) error
+	DeleteSeeder(infohash, peerID string) error
+
+	PutLeecher(infohash string, p chihaya.Peer) error
+	DeleteLeecher(infohash, peerID string) error
+
+	GraduateLeecher(infohash string, p chihaya.Peer) error
+	AnnouncePeers(infohash string, seeder bool, numWant int) (peers, peers6 []chihaya.Peer, err error)
+	CollectGarbage(cutoff time.Time) error
+}
+
+// PeerStoreDriver represents an interface for creating a handle to the storage
+// of peers.
+type PeerStoreDriver interface {
+	New(*Config) (PeerStore, error)
+}
+
+// RegisterPeerStoreDriver makes a driver available by the provided name.
+//
+// If this function is called twice with the same name or if the driver is nil,
+// it panics.
+func RegisterPeerStoreDriver(name string, driver PeerStoreDriver) {
+	if driver == nil {
+		panic("storage: could not register nil PeerStoreDriver")
+	}
+
+	if _, dup := peerStoreDrivers[name]; dup {
+		panic("storage: could not register duplicate PeerStoreDriver: " + name)
+	}
+
+	peerStoreDrivers[name] = driver
+}
+
+// OpenPeerStore returns a PeerStore specified by a configuration.
+func OpenPeerStore(name string, cfg *Config) (PeerStore, error) {
+	driver, ok := peerStoreDrivers[name]
+	if !ok {
+		return nil, fmt.Errorf(
+			"storage: unknown driver %q (forgotten import?)",
+			name,
+		)
+	}
+
+	return driver.New(cfg)
+}

server/store/store.go

@@ -0,0 +1,95 @@
+// Copyright 2016 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 store
+
+import (
+	"errors"
+	"log"
+	"sync"
+	"time"
+
+	"gopkg.in/yaml.v2"
+
+	"github.com/chihaya/chihaya/config"
+	"github.com/chihaya/chihaya/server"
+	"github.com/chihaya/chihaya/tracker"
+)
+
+var theStore *Store
+
+func init() {
+	server.Register("store", constructor)
+}
+
+func constructor(srvcfg *config.ServerConfig, tkr *tracker.Tracker) (server.Server, error) {
+	if theStore == nil {
+		cfg, err := newConfig(srvcfg)
+		if err != nil {
+			return nil, errors.New("store: invalid store config: " + err.Error())
+		}
+
+		theStore = &Store{
+			cfg: cfg,
+			tkr: tkr,
+		}
+	}
+	return theStore, nil
+}
+
+type Config struct {
+	Addr              string        `yaml:"addr"`
+	RequestTimeout    time.Duration `yaml:"requestTimeout"`
+	ReadTimeout       time.Duration `yaml:"readTimeout"`
+	WriteTimeout      time.Duration `yaml:"writeTimeout"`
+	GCAfter           time.Duration `yaml:"gcAfter"`
+	ClientStore       string        `yaml:"clientStore"`
+	ClientStoreConfig interface{}   `yaml:"clienStoreConfig"`
+	PeerStore         string        `yaml:"peerStore"`
+	PeerStoreConfig   interface{}   `yaml:"peerStoreConfig"`
+}
+
+func newConfig(srvcfg interface{}) (*Config, error) {
+	bytes, err := yaml.Marshal(srvcfg)
+	if err != nil {
+		return nil, err
+	}
+
+	var cfg Config
+	err = yaml.Unmarshal(bytes, &cfg)
+	if err != nil {
+		return nil, err
+	}
+
+	return &cfg, nil
+}
+
+// MustGetStore is used by middleware to access the store.
+//
+// This function calls log.Fatal if a server hasn't been already created by
+// the server package.
+func MustGetStore() *Store {
+	if theStore == nil {
+		log.Fatal("store middleware used without store server")
+	}
+	return theStore
+}
+
+type Store struct {
+	cfg      *Config
+	tkr      *tracker.Tracker
+	shutdown chan struct{}
+	wg       sync.WaitGroup
+
+	PeerStore
+	ClientStore
+}
+
+func (s *Store) Start() {
+}
+
+func (s *Store) Stop() {
+	close(s.shutdown)
+	s.wg.Wait()
+}