165 lines
3.5 KiB
Go
165 lines
3.5 KiB
Go
/*
|
|
* Copyright 2019 Dgraph Labs, Inc. and Contributors
|
|
*
|
|
* Licensed under the Apache License, Version 2.0 (the "License");
|
|
* you may not use this file except in compliance with the License.
|
|
* You may obtain a copy of the License at
|
|
*
|
|
* http://www.apache.org/licenses/LICENSE-2.0
|
|
*
|
|
* Unless required by applicable law or agreed to in writing, software
|
|
* distributed under the License is distributed on an "AS IS" BASIS,
|
|
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
|
* See the License for the specific language governing permissions and
|
|
* limitations under the License.
|
|
*/
|
|
|
|
package badger
|
|
|
|
import (
|
|
"sync"
|
|
|
|
"github.com/dgraph-io/badger/pb"
|
|
"github.com/dgraph-io/badger/trie"
|
|
"github.com/dgraph-io/badger/y"
|
|
)
|
|
|
|
type subscriber struct {
|
|
prefixes [][]byte
|
|
sendCh chan<- *pb.KVList
|
|
subCloser *y.Closer
|
|
}
|
|
|
|
type publisher struct {
|
|
sync.Mutex
|
|
pubCh chan requests
|
|
subscribers map[uint64]subscriber
|
|
nextID uint64
|
|
indexer *trie.Trie
|
|
}
|
|
|
|
func newPublisher() *publisher {
|
|
return &publisher{
|
|
pubCh: make(chan requests, 1000),
|
|
subscribers: make(map[uint64]subscriber),
|
|
nextID: 0,
|
|
indexer: trie.NewTrie(),
|
|
}
|
|
}
|
|
|
|
func (p *publisher) listenForUpdates(c *y.Closer) {
|
|
defer func() {
|
|
p.cleanSubscribers()
|
|
c.Done()
|
|
}()
|
|
slurp := func(batch requests) {
|
|
for {
|
|
select {
|
|
case reqs := <-p.pubCh:
|
|
batch = append(batch, reqs...)
|
|
default:
|
|
p.publishUpdates(batch)
|
|
return
|
|
}
|
|
}
|
|
}
|
|
for {
|
|
select {
|
|
case <-c.HasBeenClosed():
|
|
return
|
|
case reqs := <-p.pubCh:
|
|
slurp(reqs)
|
|
}
|
|
}
|
|
}
|
|
|
|
func (p *publisher) publishUpdates(reqs requests) {
|
|
p.Lock()
|
|
defer func() {
|
|
p.Unlock()
|
|
// Release all the request.
|
|
reqs.DecrRef()
|
|
}()
|
|
batchedUpdates := make(map[uint64]*pb.KVList)
|
|
for _, req := range reqs {
|
|
for _, e := range req.Entries {
|
|
ids := p.indexer.Get(e.Key)
|
|
if len(ids) > 0 {
|
|
k := y.SafeCopy(nil, e.Key)
|
|
kv := &pb.KV{
|
|
Key: y.ParseKey(k),
|
|
Value: y.SafeCopy(nil, e.Value),
|
|
Meta: []byte{e.UserMeta},
|
|
ExpiresAt: e.ExpiresAt,
|
|
Version: y.ParseTs(k),
|
|
}
|
|
for id := range ids {
|
|
if _, ok := batchedUpdates[id]; !ok {
|
|
batchedUpdates[id] = &pb.KVList{}
|
|
}
|
|
batchedUpdates[id].Kv = append(batchedUpdates[id].Kv, kv)
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
for id, kvs := range batchedUpdates {
|
|
p.subscribers[id].sendCh <- kvs
|
|
}
|
|
}
|
|
|
|
func (p *publisher) newSubscriber(c *y.Closer, prefixes ...[]byte) (<-chan *pb.KVList, uint64) {
|
|
p.Lock()
|
|
defer p.Unlock()
|
|
ch := make(chan *pb.KVList, 1000)
|
|
id := p.nextID
|
|
// Increment next ID.
|
|
p.nextID++
|
|
p.subscribers[id] = subscriber{
|
|
prefixes: prefixes,
|
|
sendCh: ch,
|
|
subCloser: c,
|
|
}
|
|
for _, prefix := range prefixes {
|
|
p.indexer.Add(prefix, id)
|
|
}
|
|
return ch, id
|
|
}
|
|
|
|
// cleanSubscribers stops all the subscribers. Ideally, It should be called while closing DB.
|
|
func (p *publisher) cleanSubscribers() {
|
|
p.Lock()
|
|
defer p.Unlock()
|
|
for id, s := range p.subscribers {
|
|
for _, prefix := range s.prefixes {
|
|
p.indexer.Delete(prefix, id)
|
|
}
|
|
delete(p.subscribers, id)
|
|
s.subCloser.SignalAndWait()
|
|
}
|
|
}
|
|
|
|
func (p *publisher) deleteSubscriber(id uint64) {
|
|
p.Lock()
|
|
defer p.Unlock()
|
|
if s, ok := p.subscribers[id]; ok {
|
|
for _, prefix := range s.prefixes {
|
|
p.indexer.Delete(prefix, id)
|
|
}
|
|
}
|
|
delete(p.subscribers, id)
|
|
}
|
|
|
|
func (p *publisher) sendUpdates(reqs requests) {
|
|
if p.noOfSubscribers() != 0 {
|
|
reqs.IncrRef()
|
|
p.pubCh <- reqs
|
|
}
|
|
}
|
|
|
|
func (p *publisher) noOfSubscribers() int {
|
|
p.Lock()
|
|
defer p.Unlock()
|
|
return len(p.subscribers)
|
|
}
|