mirror of
https://github.com/sot-tech/mochi.git
synced 2026-07-24 00:38:15 -07:00
initial middleware refactor
This commit is contained in:
Jimmy Zelinskie
@@ -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
|
||||
}
|
||||
@@ -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
|
||||
}
|
||||
Reference in New Issue
Block a user