diff --git a/db.go b/db.go index 897504f..232cfc6 100644 --- a/db.go +++ b/db.go @@ -78,8 +78,8 @@ func (db *DB) NewDHT(dbs ...int) { db.dht = shardid.NewDHT(dbs...) } -// AddDHT add new databases into current DHT -func (db *DB) AddDHT(dbs ...int) { +// DHTAdd add new databases into current DHT +func (db *DB) DHTAdd(dbs ...int) { db.Lock() defer db.Unlock() if db.dht == nil { @@ -88,14 +88,14 @@ func (db *DB) AddDHT(dbs ...int) { db.dht.Add(dbs...) } -// EndDHT end current DHT scaling, and reload. -func (db *DB) EndDHT() { +// DHTAdded get databases added on current DHT, and reload it. +func (db *DB) DHTAdded() { db.Lock() defer db.Unlock() if db.dht == nil { return } - db.dht.End() + db.dht.Done() } // OnDHT select database from DHT diff --git a/db_test.go b/db_test.go index 7a393e1..b049b6a 100644 --- a/db_test.go +++ b/db_test.go @@ -83,6 +83,25 @@ func TestOn(t *testing.T) { } +func TestDHT(t *testing.T) { + db := Open(createSQLite3()) + + // always work if only single server + ctx, err := db.OnDHT("") + require.Equal(t, 0, ctx.index) + require.Nil(t, err) + + // MUST NOT panic even DHT is missing + db.DHTAdd(1) + db.DHTAdded() + + db.Add(createSQLite3()) + + ctx, err = db.OnDHT("") + require.ErrorIs(t, err, ErrMissingDHT) + require.Nil(t, ctx) +} + func TestOnDHT(t *testing.T) { dbs := make([]*sql.DB, 0, 10) @@ -240,7 +259,7 @@ func TestDHTScaling(t *testing.T) { require.Equal(t, it.current, ctx.index) } - db.AddDHT(2) + db.DHTAdd(2) for v, it := range values { ctx, err := db.OnDHT(v) @@ -253,7 +272,7 @@ func TestDHTScaling(t *testing.T) { } - db.EndDHT() + db.DHTAdded() for v, it := range values { ctx, err := db.OnDHT(v) require.NoError(t, err) diff --git a/shardid/dht.go b/shardid/dht.go index 188ed6b..5584502 100644 --- a/shardid/dht.go +++ b/shardid/dht.go @@ -33,7 +33,7 @@ func NewDHT(dbs ...int) *DHT { m.dbs[i] = db } - m.current = NewHR(m.dbsCount) + m.current = NewHR(m.dbsCount, WithReplicas(defaultReplicas...)) return m } @@ -60,8 +60,8 @@ func (m *DHT) On(v string) (int, int, error) { return current, current, nil } -// End end added, and reset current and next HashRings -func (m *DHT) End() { +// Done dbs are added, then reset current/next HashRing +func (m *DHT) Done() { m.Lock() defer m.Unlock() @@ -81,7 +81,7 @@ func (m *DHT) Add(dbs ...int) []int { } m.dbsCount += len(dbs) - m.next = NewHR(m.dbsCount) + m.next = NewHR(m.dbsCount, WithReplicas(defaultReplicas...)) var ( db1 int db2 int diff --git a/shardid/dht_test.go b/shardid/dht_test.go index aa8918f..2a25472 100644 --- a/shardid/dht_test.go +++ b/shardid/dht_test.go @@ -113,7 +113,7 @@ func TestDHT(t *testing.T) { require.Equal(t, 1, next) require.Nil(t, err) // > Q1 last node => E0! - m.End() + m.Done() cur, next, err = m.On("E1") require.Equal(t, 3, cur)