diff --git a/config_example.yaml b/config_example.yaml index 664f2e6..5a3e064 100644 --- a/config_example.yaml +++ b/config_example.yaml @@ -25,6 +25,7 @@ chihaya: write_timeout: 10s client_store: memory ip_store: memory + string_store: memory peer_store: memory peer_store_config: gcAfter: 30m diff --git a/server/store/memory/string_store.go b/server/store/memory/string_store.go new file mode 100644 index 0000000..e1bb0ce --- /dev/null +++ b/server/store/memory/string_store.go @@ -0,0 +1,57 @@ +// 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/server/store" +) + +func init() { + store.RegisterStringStoreDriver("memory", &stringStoreDriver{}) +} + +type stringStoreDriver struct{} + +func (d *stringStoreDriver) New(cfg *store.Config) (store.StringStore, error) { + return &stringStore{ + strings: make(map[string]struct{}), + }, nil +} + +type stringStore struct { + strings map[string]struct{} + sync.RWMutex +} + +var _ store.StringStore = &stringStore{} + +func (ss *stringStore) PutString(s string) error { + ss.Lock() + defer ss.Unlock() + + ss.strings[s] = struct{}{} + + return nil +} + +func (ss *stringStore) HasString(s string) (bool, error) { + ss.RLock() + defer ss.RUnlock() + + _, ok := ss.strings[s] + + return ok, nil +} + +func (ss *stringStore) RemoveString(s string) error { + ss.Lock() + defer ss.Unlock() + + delete(ss.strings, s) + + return nil +} diff --git a/server/store/memory/string_store_test.go b/server/store/memory/string_store_test.go new file mode 100644 index 0000000..49e56c8 --- /dev/null +++ b/server/store/memory/string_store_test.go @@ -0,0 +1,73 @@ +// 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 ( + "github.com/chihaya/chihaya/server/store" + "github.com/stretchr/testify/assert" + "testing" +) + +var ( + driver = &stringStoreDriver{} + s1 = "abc" + s2 = "def" +) + +func TestStringStore(t *testing.T) { + ss, err := driver.New(&store.Config{}) + assert.Nil(t, err) + assert.NotNil(t, ss) + + has, err := ss.HasString(s1) + assert.Nil(t, err) + assert.False(t, has) + + has, err = ss.HasString(s2) + assert.Nil(t, err) + assert.False(t, has) + + err = ss.RemoveString(s1) + assert.Nil(t, err) + + err = ss.PutString(s1) + assert.Nil(t, err) + + has, err = ss.HasString(s1) + assert.Nil(t, err) + assert.True(t, has) + + has, err = ss.HasString(s2) + assert.Nil(t, err) + assert.False(t, has) + + err = ss.PutString(s1) + assert.Nil(t, err) + + err = ss.PutString(s2) + assert.Nil(t, err) + + has, err = ss.HasString(s1) + assert.Nil(t, err) + assert.True(t, has) + + has, err = ss.HasString(s2) + assert.Nil(t, err) + assert.True(t, has) + + err = ss.RemoveString(s1) + assert.Nil(t, err) + + err = ss.RemoveString(s2) + assert.Nil(t, err) + + has, err = ss.HasString(s1) + assert.Nil(t, err) + assert.False(t, has) + + has, err = ss.HasString(s2) + assert.Nil(t, err) + assert.False(t, has) +} diff --git a/server/store/store.go b/server/store/store.go index 801cd56..df5dfe3 100644 --- a/server/store/store.go +++ b/server/store/store.go @@ -45,6 +45,11 @@ func constructor(srvcfg *chihaya.ServerConfig, tkr *tracker.Tracker) (server.Ser return nil, err } + ss, err := OpenStringStore(cfg) + if err != nil { + return nil, err + } + theStore = &Store{ cfg: cfg, tkr: tkr, @@ -52,6 +57,7 @@ func constructor(srvcfg *chihaya.ServerConfig, tkr *tracker.Tracker) (server.Ser ClientStore: cs, PeerStore: ps, IPStore: ips, + StringStore: ss, } } return theStore, nil @@ -69,6 +75,8 @@ type Config struct { PeerStoreConfig interface{} `yaml:"peer_store_config"` IPStore string `yaml:"ip_store"` IPStoreConfig interface{} `yaml:"ip_store_config"` + StringStore string `yaml:"string_store"` + StringStoreConfig interface{} `yaml:"string_store_config"` } func newConfig(srvcfg *chihaya.ServerConfig) (*Config, error) { @@ -106,6 +114,7 @@ type Store struct { PeerStore ClientStore IPStore + StringStore } func (s *Store) Start() { diff --git a/server/store/string_store.go b/server/store/string_store.go new file mode 100644 index 0000000..1f530e9 --- /dev/null +++ b/server/store/string_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" + +// PrefixInfohash is the prefix to be used for infohashes. +const PrefixInfohash = "ih-" + +var stringStoreDrivers = make(map[string]StringStoreDriver) + +// StringStore represents an interface for manipulating strings. +type StringStore interface { + PutString(s string) error + HasString(s string) (bool, error) + RemoveString(s string) error +} + +// StringStoreDriver represents an interface for creating a handle to the +// storage of swarms. +type StringStoreDriver interface { + New(*Config) (StringStore, error) +} + +// RegisterStringStoreDriver 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 RegisterStringStoreDriver(name string, driver StringStoreDriver) { + if driver == nil { + panic("store: could not register nil StringStoreDriver") + } + if _, dup := stringStoreDrivers[name]; dup { + panic("store: could not register duplicate StringStoreDriver: " + name) + } + stringStoreDrivers[name] = driver +} + +// OpenStringStore returns a StringStore specified by a configuration. +func OpenStringStore(cfg *Config) (StringStore, error) { + driver, ok := stringStoreDrivers[cfg.StringStore] + if !ok { + return nil, fmt.Errorf("store: unknown driver %q (forgotten import?)", cfg.StringStore) + } + + return driver.New(cfg) +}