milvus-io
/
milvus
mirror of https://github.com/milvus-io/milvus.git
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity

Refine rocksdb (#12743) 

Signed-off-by: xiaofan-luan <xiaofan.luan@zilliz.com>
pull/13683/head
Xiaofan 2021-12-17 23:44:42 +08:00 committed by GitHub
parent 2ef2228ad0
commit b6cca25d1f
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
13 changed files with 907 additions and 646 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

105
internal/kv/rocksdb/rocksdb_kv.go
View File

@ -20,6 +20,7 @@ import (
"errors" "errors"
"fmt" "fmt"
"github.com/milvus-io/milvus/internal/util/typeutil"
"github.com/tecbot/gorocksdb" "github.com/tecbot/gorocksdb"
) )
@ -37,8 +38,9 @@ const (
LRUCacheSize = 0 LRUCacheSize = 0
) )
// NewRocksdbKV returns a rockskv object // NewRocksdbKV returns a rockskv object, only used in test
func NewRocksdbKV(name string) (*RocksdbKV, error) { func NewRocksdbKV(name string) (*RocksdbKV, error) {
// TODO we should use multiple column family of rocks db rather than init multiple db instance
if name == "" { if name == "" {
return nil, errors.New("rocksdb name is nil") return nil, errors.New("rocksdb name is nil")
} }
@ -48,12 +50,21 @@ func NewRocksdbKV(name string) (*RocksdbKV, error) {
bbto.SetBlockCache(gorocksdb.NewLRUCache(LRUCacheSize)) bbto.SetBlockCache(gorocksdb.NewLRUCache(LRUCacheSize))
opts := gorocksdb.NewDefaultOptions() opts := gorocksdb.NewDefaultOptions()
opts.SetBlockBasedTableFactory(bbto) opts.SetBlockBasedTableFactory(bbto)
// by default there are only 1 thread for flush compaction, which may block each other.
// increase to a reasonable thread numbers
opts.IncreaseParallelism(2)
// enable back ground flush
opts.SetMaxBackgroundFlushes(1)
opts.SetCreateIfMissing(true) opts.SetCreateIfMissing(true)
return NewRocksdbKVWithOpts(name, opts)
}
// NewRocksdbKV returns a rockskv object
func NewRocksdbKVWithOpts(name string, opts *gorocksdb.Options) (*RocksdbKV, error) {
ro := gorocksdb.NewDefaultReadOptions() ro := gorocksdb.NewDefaultReadOptions()
ro.SetFillCache(false)
wo := gorocksdb.NewDefaultWriteOptions() wo := gorocksdb.NewDefaultWriteOptions()
// only has one columnn families
db, err := gorocksdb.OpenDb(opts, name) db, err := gorocksdb.OpenDb(opts, name)
if err != nil { if err != nil {
return nil, err return nil, err
@ -84,7 +95,9 @@ func (kv *RocksdbKV) Load(key string) (string, error) {
if kv.DB == nil { if kv.DB == nil {
return "", fmt.Errorf("rocksdb instance is nil when load %s", key) return "", fmt.Errorf("rocksdb instance is nil when load %s", key)
} }
if key == "" {
return "", errors.New("rocksdb kv does not support load empty key")
}
value, err := kv.DB.Get(kv.ReadOptions, []byte(key)) value, err := kv.DB.Get(kv.ReadOptions, []byte(key))
if err != nil { if err != nil {
return "", err return "", err
@ -94,27 +107,20 @@ func (kv *RocksdbKV) Load(key string) (string, error) {
} }
// LoadWithPrefix returns a batch values of keys with a prefix // LoadWithPrefix returns a batch values of keys with a prefix
func (kv *RocksdbKV) LoadWithPrefix(key string) ([]string, []string, error) { // if prefix is "", then load every thing from the database
if key == "" { func (kv *RocksdbKV) LoadWithPrefix(prefix string) ([]string, []string, error) {
return nil, nil, errors.New("key is nil in LoadWithPrefix")
}
if kv.DB == nil { if kv.DB == nil {
return nil, nil, fmt.Errorf("rocksdb instance is nil when load %s", key) return nil, nil, fmt.Errorf("rocksdb instance is nil when load %s", prefix)
} }
kv.ReadOptions.SetPrefixSameAsStart(true) kv.ReadOptions.SetPrefixSameAsStart(true)
kv.DB.Close() if prefix != "" {
kv.Opts.SetPrefixExtractor(gorocksdb.NewFixedPrefixTransform(len(key))) kv.ReadOptions.SetIterateUpperBound([]byte(typeutil.AddOne(prefix)))
var err error
kv.DB, err = gorocksdb.OpenDb(kv.Opts, kv.GetName())
if err != nil {
return nil, nil, err
} }
iter := kv.DB.NewIterator(kv.ReadOptions) iter := kv.DB.NewIterator(kv.ReadOptions)
defer iter.Close() defer iter.Close()
keys := make([]string, 0) keys := make([]string, 0)
values := make([]string, 0) values := make([]string, 0)
iter.Seek([]byte(key)) iter.Seek([]byte(prefix))
for ; iter.Valid(); iter.Next() { for ; iter.Valid(); iter.Next() {
key := iter.Key() key := iter.Key()
value := iter.Value() value := iter.Value()
@ -129,15 +135,6 @@ func (kv *RocksdbKV) LoadWithPrefix(key string) ([]string, []string, error) {
return keys, values, nil return keys, values, nil
} }
// ResetPrefixLength will close rocksdb object and open a new rocksdb with new prefix length
func (kv *RocksdbKV) ResetPrefixLength(len int) error {
kv.DB.Close()
kv.Opts.SetPrefixExtractor(gorocksdb.NewFixedPrefixTransform(len))
var err error
kv.DB, err = gorocksdb.OpenDb(kv.Opts, kv.GetName())
return err
}
// MultiLoad load a batch of values by keys // MultiLoad load a batch of values by keys
func (kv *RocksdbKV) MultiLoad(keys []string) ([]string, error) { func (kv *RocksdbKV) MultiLoad(keys []string) ([]string, error) {
if kv.DB == nil { if kv.DB == nil {
@ -160,6 +157,12 @@ func (kv *RocksdbKV) Save(key, value string) error {
if kv.DB == nil { if kv.DB == nil {
return errors.New("rocksdb instance is nil when do save") return errors.New("rocksdb instance is nil when do save")
} }
if key == "" {
return errors.New("rocksdb kv does not support empty key")
}
if value == "" {
return errors.New("rocksdb kv does not support empty value")
}
err := kv.DB.Put(kv.WriteOptions, []byte(key), []byte(value)) err := kv.DB.Put(kv.WriteOptions, []byte(key), []byte(value))
return err return err
} }
@ -179,34 +182,27 @@ func (kv *RocksdbKV) MultiSave(kvs map[string]string) error {
} }
// RemoveWithPrefix removes a batch of key-values with specified prefix // RemoveWithPrefix removes a batch of key-values with specified prefix
// If prefix is "", then all data in the rocksdb kv will be deleted
func (kv *RocksdbKV) RemoveWithPrefix(prefix string) error { func (kv *RocksdbKV) RemoveWithPrefix(prefix string) error {
if kv.DB == nil { if kv.DB == nil {
return errors.New("rocksdb instance is nil when do RemoveWithPrefix") return errors.New("rocksdb instance is nil when do RemoveWithPrefix")
} }
kv.ReadOptions.SetPrefixSameAsStart(true) if len(prefix) == 0 {
kv.DB.Close() // better to use drop column family, but as we use default column family, we just delete ["",lastKey+1)
kv.Opts.SetPrefixExtractor(gorocksdb.NewFixedPrefixTransform(len(prefix))) readOpts := gorocksdb.NewDefaultReadOptions()
var err error defer readOpts.Destroy()
kv.DB, err = gorocksdb.OpenDb(kv.Opts, kv.GetName()) iter := kv.DB.NewIterator(readOpts)
if err != nil { defer iter.Close()
return err // seek to the last key
} iter.SeekToLast()
if iter.Valid() {
iter := kv.DB.NewIterator(kv.ReadOptions) return kv.DeleteRange(prefix, typeutil.AddOne(string(iter.Key().Data())))
defer iter.Close()
iter.Seek([]byte(prefix))
for ; iter.Valid(); iter.Next() {
key := iter.Key()
err := kv.DB.Delete(kv.WriteOptions, key.Data())
key.Free()
if err != nil {
return nil
} }
// nothing in the range, skip
return nil
} }
if err := iter.Err(); err != nil { prefixEnd := typeutil.AddOne(prefix)
return err return kv.DeleteRange(prefix, prefixEnd)
}
return nil
} }
// Remove is used to remove a pair of key-value // Remove is used to remove a pair of key-value
@ -214,6 +210,9 @@ func (kv *RocksdbKV) Remove(key string) error {
if kv.DB == nil { if kv.DB == nil {
return errors.New("rocksdb instance is nil when do Remove") return errors.New("rocksdb instance is nil when do Remove")
} }
if key == "" {
return errors.New("rocksdb kv does not support empty key")
}
err := kv.DB.Delete(kv.WriteOptions, []byte(key)) err := kv.DB.Delete(kv.WriteOptions, []byte(key))
return err return err
} }
@ -254,15 +253,11 @@ func (kv *RocksdbKV) DeleteRange(startKey, endKey string) error {
if kv.DB == nil { if kv.DB == nil {
return errors.New("Rocksdb instance is nil when do DeleteRange") return errors.New("Rocksdb instance is nil when do DeleteRange")
} }
if startKey >= endKey {
return fmt.Errorf("rockskv delete range startkey must < endkey, startkey %s, endkey %s", startKey, endKey)
}
writeBatch := gorocksdb.NewWriteBatch() writeBatch := gorocksdb.NewWriteBatch()
defer writeBatch.Destroy() defer writeBatch.Destroy()
if len(startKey) == 0 {
iter := kv.DB.NewIterator(kv.ReadOptions)
defer iter.Close()
iter.SeekToFirst()
startKey = string(iter.Key().Data())
}
writeBatch.DeleteRange([]byte(startKey), []byte(endKey)) writeBatch.DeleteRange([]byte(startKey), []byte(endKey))
err := kv.DB.Write(kv.WriteOptions, writeBatch) err := kv.DB.Write(kv.WriteOptions, writeBatch)
return err return err

48
internal/kv/rocksdb/rocksdb_kv_test.go
View File

@ -108,10 +108,9 @@ func TestRocksdbKV_Prefix(t *testing.T) {
keys, vals, err := rocksdbKV.LoadWithPrefix("abc") keys, vals, err := rocksdbKV.LoadWithPrefix("abc")
assert.Nil(t, err) assert.Nil(t, err)
assert.Equal(t, len(keys), 1) assert.Equal(t, len(keys), 1)
assert.Equal(t, len(vals), 1) assert.Equal(t, len(vals), 1)
//fmt.Println(keys)
//fmt.Println(vals)
err = rocksdbKV.RemoveWithPrefix("abc") err = rocksdbKV.RemoveWithPrefix("abc")
assert.Nil(t, err) assert.Nil(t, err)
@ -124,6 +123,27 @@ func TestRocksdbKV_Prefix(t *testing.T) {
val, err = rocksdbKV.Load("abddqqq") val, err = rocksdbKV.Load("abddqqq")
assert.Nil(t, err) assert.Nil(t, err)
assert.Equal(t, val, "1234555") assert.Equal(t, val, "1234555")
// test remove ""
err = rocksdbKV.RemoveWithPrefix("")
assert.Nil(t, err)
// test remove from a empty cf
err = rocksdbKV.RemoveWithPrefix("")
assert.Nil(t, err)
val, err = rocksdbKV.Load("abddqqq")
assert.Nil(t, err)
assert.Equal(t, len(val), 0)
// test we can still save after drop
err = rocksdbKV.Save("abcd", "123")
assert.Nil(t, err)
val, err = rocksdbKV.Load("abcd")
assert.Nil(t, err)
assert.Equal(t, val, "123")
} }
func TestRocksdbKV_Goroutines(t *testing.T) { func TestRocksdbKV_Goroutines(t *testing.T) {
@ -151,7 +171,7 @@ func TestRocksdbKV_Goroutines(t *testing.T) {
wg.Wait() wg.Wait()
} }
func TestRocksdbKV_Dummy(t *testing.T) { func TestRocksdbKV_DummyDB(t *testing.T) {
name := "/tmp/rocksdb_dummy" name := "/tmp/rocksdb_dummy"
rocksdbkv, err := rocksdbkv.NewRocksdbKV(name) rocksdbkv, err := rocksdbkv.NewRocksdbKV(name)
assert.Nil(t, err) assert.Nil(t, err)
@ -184,3 +204,25 @@ func TestRocksdbKV_Dummy(t *testing.T) {
_, err = rocksdbkv.Load("dummy") _, err = rocksdbkv.Load("dummy")
assert.Error(t, err) assert.Error(t, err)
} }
func TestRocksdbKV_CornerCase(t *testing.T) {
name := "/tmp/rocksdb_corner"
rocksdbkv, err := rocksdbkv.NewRocksdbKV(name)
assert.Nil(t, err)
defer rocksdbkv.Close()
defer rocksdbkv.RemoveWithPrefix("")
_, err = rocksdbkv.Load("")
assert.Error(t, err)
keys, values, err := rocksdbkv.LoadWithPrefix("")
assert.NoError(t, err)
assert.Equal(t, len(keys), 0)
assert.Equal(t, len(values), 0)
err = rocksdbkv.Save("", "")
assert.Error(t, err)
err = rocksdbkv.Save("test", "")
assert.Error(t, err)
err = rocksdbkv.Remove("")
assert.Error(t, err)
err = rocksdbkv.DeleteRange("a", "a")
assert.Error(t, err)
}

6
internal/util/rocksmq/client/rocksmq/client_impl.go
View File

@ -71,7 +71,11 @@ func (c *client) Subscribe(options ConsumerOptions) (Consumer, error) {
if reflect.ValueOf(c.server).IsNil() { if reflect.ValueOf(c.server).IsNil() {
return nil, newError(0, "Rmq server is nil") return nil, newError(0, "Rmq server is nil")
} }
if exist, con := c.server.ExistConsumerGroup(options.Topic, options.SubscriptionName); exist { exist, con, err := c.server.ExistConsumerGroup(options.Topic, options.SubscriptionName)
if err != nil {
return nil, err
}
if exist {
log.Debug("ConsumerGroup already existed", zap.Any("topic", options.Topic), zap.Any("SubscriptionName", options.SubscriptionName)) log.Debug("ConsumerGroup already existed", zap.Any("topic", options.Topic), zap.Any("SubscriptionName", options.SubscriptionName))
consumer, err := getExistedConsumer(c, options, con.MsgMutex) consumer, err := getExistedConsumer(c, options, con.MsgMutex)
if err != nil { if err != nil {

20
internal/util/rocksmq/client/rocksmq/test_helper.go
View File

@ -16,10 +16,8 @@ import (
"os" "os"
"time" "time"
"github.com/milvus-io/milvus/internal/allocator"
rocksdbkv "github.com/milvus-io/milvus/internal/kv/rocksdb"
"github.com/milvus-io/milvus/internal/log" "github.com/milvus-io/milvus/internal/log"
rocksmq "github.com/milvus-io/milvus/internal/util/rocksmq/server/rocksmq" "github.com/milvus-io/milvus/internal/util/rocksmq/server/rocksmq"
server "github.com/milvus-io/milvus/internal/util/rocksmq/server/rocksmq" server "github.com/milvus-io/milvus/internal/util/rocksmq/server/rocksmq"
"go.uber.org/zap" "go.uber.org/zap"
@ -45,23 +43,9 @@ func newMockClient() *client {
return client return client
} }
func initIDAllocator(kvPath string) *allocator.GlobalIDAllocator {
rocksdbKV, err := rocksdbkv.NewRocksdbKV(kvPath)
if err != nil {
panic(err)
}
idAllocator := allocator.NewGlobalIDAllocator("rmq_id", rocksdbKV)
_ = idAllocator.Initialize()
return idAllocator
}
func newRocksMQ(rmqPath string) server.RocksMQ { func newRocksMQ(rmqPath string) server.RocksMQ {
kvPath := rmqPath + "_kv"
idAllocator := initIDAllocator(kvPath)
rocksdbPath := rmqPath + "_db" rocksdbPath := rmqPath + "_db"
rmq, _ := rocksmq.NewRocksMQ(rocksdbPath, nil)
rmq, _ := rocksmq.NewRocksMQ(rocksdbPath, idAllocator)
return rmq return rmq
} }

16
internal/util/rocksmq/server/rocksmq/global_rmq.go
View File

@ -19,7 +19,6 @@ import (
"sync/atomic" "sync/atomic"
"github.com/milvus-io/milvus/internal/allocator" "github.com/milvus-io/milvus/internal/allocator"
rocksdbkv "github.com/milvus-io/milvus/internal/kv/rocksdb"
"github.com/milvus-io/milvus/internal/log" "github.com/milvus-io/milvus/internal/log"
"github.com/milvus-io/milvus/internal/util/paramtable" "github.com/milvus-io/milvus/internal/util/paramtable"
@ -64,15 +63,6 @@ func InitRocksMQ() error {
} }
} }
kvname := rocksdbName + "_kv"
var rkv *rocksdbkv.RocksdbKV
rkv, finalErr = rocksdbkv.NewRocksdbKV(kvname)
if finalErr != nil {
return
}
idAllocator := allocator.NewGlobalIDAllocator("rmq_id", rkv)
_ = idAllocator.Initialize()
rawRmqPageSize, err := params.Load("rocksmq.rocksmqPageSize") rawRmqPageSize, err := params.Load("rocksmq.rocksmqPageSize")
if err == nil && rawRmqPageSize != "" { if err == nil && rawRmqPageSize != "" {
rmqPageSize, err := strconv.ParseInt(rawRmqPageSize, 10, 64) rmqPageSize, err := strconv.ParseInt(rawRmqPageSize, 10, 64)
@ -86,7 +76,7 @@ func InitRocksMQ() error {
if err == nil && rawRmqRetentionTimeInMinutes != "" { if err == nil && rawRmqRetentionTimeInMinutes != "" {
rawRmqRetentionTimeInMinutes, err := strconv.ParseInt(rawRmqRetentionTimeInMinutes, 10, 64) rawRmqRetentionTimeInMinutes, err := strconv.ParseInt(rawRmqRetentionTimeInMinutes, 10, 64)
if err == nil { if err == nil {
atomic.StoreInt64(&RocksmqRetentionTimeInMinutes, rawRmqRetentionTimeInMinutes) atomic.StoreInt64(&RocksmqRetentionTimeInSecs, rawRmqRetentionTimeInMinutes*60)
} else { } else {
log.Warn("rocksmq.retentionTimeInMinutes is invalid, using default value 3 days") log.Warn("rocksmq.retentionTimeInMinutes is invalid, using default value 3 days")
} }
@ -100,9 +90,9 @@ func InitRocksMQ() error {
log.Warn("rocksmq.retentionSizeInMB is invalid, using default value 0") log.Warn("rocksmq.retentionSizeInMB is invalid, using default value 0")
} }
} }
log.Debug("", zap.Any("RocksmqRetentionTimeInMinutes", RocksmqRetentionTimeInMinutes), log.Debug("", zap.Any("RocksmqRetentionTimeInMinutes", rawRmqRetentionTimeInMinutes),
zap.Any("RocksmqRetentionSizeInMB", RocksmqRetentionSizeInMB), zap.Any("RocksmqPageSize", RocksmqPageSize)) zap.Any("RocksmqRetentionSizeInMB", RocksmqRetentionSizeInMB), zap.Any("RocksmqPageSize", RocksmqPageSize))
Rmq, finalErr = NewRocksMQ(rocksdbName, idAllocator) Rmq, finalErr = NewRocksMQ(rocksdbName, nil)
}) })
return finalErr return finalErr
} }

18
internal/util/rocksmq/server/rocksmq/global_rmq_test.go
View File

@ -26,6 +26,7 @@ import (
func Test_InitRmq(t *testing.T) { func Test_InitRmq(t *testing.T) {
name := "/tmp/rmq_init" name := "/tmp/rmq_init"
defer os.RemoveAll("/tmp/rmq_init")
endpoints := os.Getenv("ETCD_ENDPOINTS") endpoints := os.Getenv("ETCD_ENDPOINTS")
if endpoints == "" { if endpoints == "" {
endpoints = "localhost:2379" endpoints = "localhost:2379"
@ -38,6 +39,8 @@ func Test_InitRmq(t *testing.T) {
idAllocator := allocator.NewGlobalIDAllocator("dummy", etcdKV) idAllocator := allocator.NewGlobalIDAllocator("dummy", etcdKV)
_ = idAllocator.Initialize() _ = idAllocator.Initialize()
defer os.RemoveAll(name + kvSuffix)
defer os.RemoveAll(name)
err = InitRmq(name, idAllocator) err = InitRmq(name, idAllocator)
defer Rmq.stopRetention() defer Rmq.stopRetention()
assert.NoError(t, err) assert.NoError(t, err)
@ -49,7 +52,7 @@ func Test_InitRocksMQ(t *testing.T) {
rmqPath := "/tmp/milvus/rdb_data_global" rmqPath := "/tmp/milvus/rdb_data_global"
err := os.Setenv("ROCKSMQ_PATH", rmqPath) err := os.Setenv("ROCKSMQ_PATH", rmqPath)
assert.Nil(t, err) assert.Nil(t, err)
defer os.RemoveAll(rmqPath) defer os.RemoveAll("/tmp/milvus")
err = InitRocksMQ() err = InitRocksMQ()
defer Rmq.stopRetention() defer Rmq.stopRetention()
assert.NoError(t, err) assert.NoError(t, err)
@ -73,10 +76,15 @@ func Test_InitRocksMQ(t *testing.T) {
func Test_InitRocksMQError(t *testing.T) { func Test_InitRocksMQError(t *testing.T) {
once = sync.Once{} once = sync.Once{}
dummyPath := "/tmp/milvus/dummy" dir := "/tmp/milvus/"
os.Create(dummyPath) dummyPath := dir + "dummy"
err := os.MkdirAll(dir, os.ModePerm)
assert.NoError(t, err)
f, err := os.Create(dummyPath)
defer f.Close()
assert.NoError(t, err)
os.Setenv("ROCKSMQ_PATH", dummyPath) os.Setenv("ROCKSMQ_PATH", dummyPath)
defer os.RemoveAll(dummyPath) defer os.RemoveAll(dir)
err := InitRocksMQ() err = InitRocksMQ()
assert.Error(t, err) assert.Error(t, err)
} }

7
internal/util/rocksmq/server/rocksmq/rocksmq.go
View File

@ -23,6 +23,7 @@ type Consumer struct {
Topic string Topic string
GroupName string GroupName string
MsgMutex chan struct{} MsgMutex chan struct{}
beginID UniqueID
} }
// ConsumerMessage that consumed from rocksdb // ConsumerMessage that consumed from rocksdb
@ -40,18 +41,18 @@ type RocksMQ interface {
DestroyConsumerGroup(topicName string, groupName string) error DestroyConsumerGroup(topicName string, groupName string) error
Close() Close()
RegisterConsumer(consumer *Consumer) RegisterConsumer(consumer *Consumer) error
Produce(topicName string, messages []ProducerMessage) ([]UniqueID, error) Produce(topicName string, messages []ProducerMessage) ([]UniqueID, error)
Consume(topicName string, groupName string, n int) ([]ConsumerMessage, error) Consume(topicName string, groupName string, n int) ([]ConsumerMessage, error)
Seek(topicName string, groupName string, msgID UniqueID) error Seek(topicName string, groupName string, msgID UniqueID) error
SeekToLatest(topicName, groupName string) error SeekToLatest(topicName, groupName string) error
ExistConsumerGroup(topicName string, groupName string) (bool, *Consumer) ExistConsumerGroup(topicName string, groupName string) (bool, *Consumer, error)
Notify(topicName, groupName string) Notify(topicName, groupName string)
CreateReader(topicName string, startMsgID UniqueID, messageIDInclusive bool, subscriptionRolePrefix string) (string, error) CreateReader(topicName string, startMsgID UniqueID, messageIDInclusive bool, subscriptionRolePrefix string) (string, error)
ReaderSeek(topicName string, readerName string, msgID UniqueID) ReaderSeek(topicName string, readerName string, msgID UniqueID) error
Next(ctx context.Context, topicName string, readerName string) (*ConsumerMessage, error) Next(ctx context.Context, topicName string, readerName string) (*ConsumerMessage, error)
HasNext(topicName string, readerName string) bool HasNext(topicName string, readerName string) bool
CloseReader(topicName string, readerName string) CloseReader(topicName string, readerName string)

538
internal/util/rocksmq/server/rocksmq/rocksmq_impl.go
View File

@ -17,6 +17,7 @@ import (
"fmt" "fmt"
"math" "math"
"path" "path"
"runtime"
"strconv" "strconv"
"strings" "strings"
"sync" "sync"
@ -40,29 +41,42 @@ type UniqueID = typeutil.UniqueID
// RmqState Rocksmq state // RmqState Rocksmq state
type RmqState = int64 type RmqState = int64
// RocksmqPageSize is the size of a message page, default 2GB // RocksmqPageSize is the size of a message page, default 256MB
var RocksmqPageSize int64 = 2 << 30 var RocksmqPageSize int64 = 256 << 20
// Const variable that will be used in rocksmqs // Const variable that will be used in rocksmqs
const ( const (
DefaultMessageID = "-1" DefaultMessageID = -1
FixedChannelNameLen = 320 FixedChannelNameLen = 320
RocksDBLRUCacheCapacity = 0
// TODO make it configable
RocksDBLRUCacheCapacity = 1 << 30
kvSuffix = "_meta_kv" kvSuffix = "_meta_kv"
MessageSizeTitle = "message_size/" // topic_begin_id/topicName
PageMsgSizeTitle = "page_message_size/" // topic begin id record a topic is valid, create when topic is created, cleaned up on destroy topic
TopicBeginIDTitle = "topic_begin_id/" TopicIDTitle = "topic_id/"
BeginIDTitle = "begin_id/"
AckedTsTitle = "acked_ts/" // message_size/topicName record the current page message size, once current message size > RocksMq size, reset this value and open a new page
AckedSizeTitle = "acked_size/" // TODO should be cached
LastRetTsTitle = "last_retention_ts/" MessageSizeTitle = "message_size/"
// page_message_size/topicName/pageId record the endId of each page, it will be purged either in retention or the destroy of topic
PageMsgSizeTitle = "page_message_size/"
// page_ts/topicName/pageId, record the page last ts, used for TTL functionality
PageTsTitle = "page_ts/"
// acked_ts/topicName/pageId, record the latest ack ts of each page, will be purged on retention or destroy of the topic
AckedTsTitle = "acked_ts/"
// only in memory
CurrentIDSuffix = "current_id"
CurrentIDSuffix = "current_id"
ReaderNamePrefix = "reader-" ReaderNamePrefix = "reader-"
RmqNotServingErrMsg = "rocksmq is not serving" RmqNotServingErrMsg = "Rocksmq is not serving"
) )
const ( const (
@ -74,6 +88,7 @@ const (
/** /**
* @brief fill with '_' to ensure channel name fixed length * @brief fill with '_' to ensure channel name fixed length
* TODO this is a waste of memory, remove it
*/ */
func fixChannelName(name string) (string, error) { func fixChannelName(name string) (string, error) {
if len(name) > FixedChannelNameLen { if len(name) > FixedChannelNameLen {
@ -109,26 +124,23 @@ func constructCurrentID(topicName, groupName string) string {
} }
/** /**
* Construct table name and fixed channel name to be a key with length of FixedChannelNameLen, * Combine metaname together with topic
* used for meta infos
*/ */
func constructKey(metaName, topic string) (string, error) { func constructKey(metaName, topic string) string {
// Check metaName/topic // Check metaName/topic
oldLen := len(metaName + topic) return metaName + topic
if oldLen > FixedChannelNameLen { }
return "", errors.New("topic name exceeds limit")
}
nameBytes := make([]byte, FixedChannelNameLen-oldLen) func parsePageID(key string) (int64, error) {
stringSlice := strings.Split(key, "/")
for i := 0; i < len(nameBytes); i++ { if len(stringSlice) != 3 {
nameBytes[i] = byte('*') return 0, fmt.Errorf("Invalid page id %s ", key)
} }
return metaName + topic + string(nameBytes), nil return strconv.ParseInt(stringSlice[2], 10, 64)
} }
func checkRetention() bool { func checkRetention() bool {
return RocksmqRetentionTimeInMinutes != -1 && RocksmqRetentionSizeInMB != -1 return RocksmqRetentionTimeInSecs != -1 || RocksmqRetentionSizeInMB != -1
} }
func getNowTs(idAllocator allocator.GIDAllocator) (int64, error) { func getNowTs(idAllocator allocator.GIDAllocator) (int64, error) {
@ -152,7 +164,7 @@ type rocksmq struct {
idAllocator allocator.GIDAllocator idAllocator allocator.GIDAllocator
storeMu *sync.Mutex storeMu *sync.Mutex
consumers sync.Map consumers sync.Map
ackedMu sync.Map consumersID sync.Map
retentionInfo *retentionInfo retentionInfo *retentionInfo
readers sync.Map readers sync.Map
@ -164,33 +176,64 @@ type rocksmq struct {
// 2. Init retention info, load retention info to memory // 2. Init retention info, load retention info to memory
// 3. Start retention goroutine // 3. Start retention goroutine
func NewRocksMQ(name string, idAllocator allocator.GIDAllocator) (*rocksmq, error) { func NewRocksMQ(name string, idAllocator allocator.GIDAllocator) (*rocksmq, error) {
// TODO we should use same rocksdb instance with different cfs
bbto := gorocksdb.NewDefaultBlockBasedTableOptions() bbto := gorocksdb.NewDefaultBlockBasedTableOptions()
bbto.SetCacheIndexAndFilterBlocks(true) bbto.SetCacheIndexAndFilterBlocks(true)
bbto.SetPinL0FilterAndIndexBlocksInCache(true)
bbto.SetBlockCache(gorocksdb.NewLRUCache(RocksDBLRUCacheCapacity)) bbto.SetBlockCache(gorocksdb.NewLRUCache(RocksDBLRUCacheCapacity))
opts := gorocksdb.NewDefaultOptions() optsKV := gorocksdb.NewDefaultOptions()
opts.SetBlockBasedTableFactory(bbto) optsKV.SetBlockBasedTableFactory(bbto)
opts.SetCreateIfMissing(true) optsKV.SetCreateIfMissing(true)
opts.SetPrefixExtractor(gorocksdb.NewFixedPrefixTransform(FixedChannelNameLen + 1)) // by default there are only 1 thread for flush compaction, which may block each other.
// opts.SetMaxOpenFiles(-1) // increase to a reasonable thread numbers
optsKV.IncreaseParallelism(runtime.NumCPU())
// enable back ground flush
optsKV.SetMaxBackgroundFlushes(1)
db, err := gorocksdb.OpenDb(opts, name) // finish rocks KV
kvName := name + kvSuffix
kv, err := rocksdbkv.NewRocksdbKVWithOpts(kvName, optsKV)
if err != nil { if err != nil {
return nil, err return nil, err
} }
kvName := name + kvSuffix // finish rocks mq store initialization, rocks mq store has to set the prefix extractor
kv, err := rocksdbkv.NewRocksdbKV(kvName) optsStore := gorocksdb.NewDefaultOptions()
// share block cache with kv
optsStore.SetBlockBasedTableFactory(bbto)
optsStore.SetCreateIfMissing(true)
// by default there are only 1 thread for flush compaction, which may block each other.
// increase to a reasonable thread numbers
optsStore.IncreaseParallelism(runtime.NumCPU())
// enable back ground flush
optsStore.SetMaxBackgroundFlushes(1)
// TODO remove fix channel name len logic
optsStore.SetPrefixExtractor(gorocksdb.NewFixedPrefixTransform(FixedChannelNameLen + 1))
db, err := gorocksdb.OpenDb(optsStore, name)
if err != nil { if err != nil {
return nil, err return nil, err
} }
var mqIDAllocator allocator.GIDAllocator
// if user didn't specify id allocator, init one with kv
if idAllocator == nil {
allocator := allocator.NewGlobalIDAllocator("rmq_id", kv)
err = allocator.Initialize()
if err != nil {
return nil, err
}
mqIDAllocator = allocator
} else {
mqIDAllocator = idAllocator
}
rmq := &rocksmq{ rmq := &rocksmq{
store: db, store: db,
kv: kv, kv: kv,
idAllocator: idAllocator, idAllocator: mqIDAllocator,
storeMu: &sync.Mutex{}, storeMu: &sync.Mutex{},
consumers: sync.Map{}, consumers: sync.Map{},
ackedMu: sync.Map{},
readers: sync.Map{}, readers: sync.Map{},
} }
@ -204,6 +247,24 @@ func NewRocksMQ(name string, idAllocator allocator.GIDAllocator) (*rocksmq, erro
rmq.retentionInfo.startRetentionInfo() rmq.retentionInfo.startRetentionInfo()
} }
atomic.StoreInt64(&rmq.state, RmqStateHealthy) atomic.StoreInt64(&rmq.state, RmqStateHealthy)
// TODO add this to monitor metrics
go func() {
for {
time.Sleep(5 * time.Minute)
log.Info("Rocksmq memory usage",
zap.String("rockskv kv cache", kv.DB.GetProperty("rocksdb.block-cache-usage")),
zap.String("rockskv memtable ", kv.DB.GetProperty("rocksdb.cur-size-all-mem-tables")),
zap.String("rockskv table readers", kv.DB.GetProperty("rocksdb.estimate-table-readers-mem")),
zap.String("rockskv pinned", kv.DB.GetProperty("rocksdb.block-cache-pinned-usage")),
zap.String("store kv cache", db.GetProperty("rocksdb.block-cache-usage")),
zap.String("store memtable ", db.GetProperty("rocksdb.cur-size-all-mem-tables")),
zap.String("store table readers", db.GetProperty("rocksdb.estimate-table-readers-mem")),
zap.String("store pinned", db.GetProperty("rocksdb.block-cache-pinned-usage")),
)
}
}()
return rmq, nil return rmq, nil
} }
@ -219,24 +280,19 @@ func (rmq *rocksmq) Close() {
atomic.StoreInt64(&rmq.state, RmqStateStopped) atomic.StoreInt64(&rmq.state, RmqStateStopped)
rmq.stopRetention() rmq.stopRetention()
rmq.consumers.Range(func(k, v interface{}) bool { rmq.consumers.Range(func(k, v interface{}) bool {
var topic string // TODO what happened if the server crashed? who handled the destroy consumer group? should we just handled it when rocksmq created?
// or we should not even make consumer info persistent?
for _, consumer := range v.([]*Consumer) { for _, consumer := range v.([]*Consumer) {
err := rmq.DestroyConsumerGroup(consumer.Topic, consumer.GroupName) err := rmq.destroyConsumerGroupInternal(consumer.Topic, consumer.GroupName)
if err != nil { if err != nil {
log.Warn("Failed to destroy consumer group in rocksmq!", zap.Any("topic", consumer.Topic), zap.Any("groupName", consumer.GroupName), zap.Any("error", err)) log.Warn("Failed to destroy consumer group in rocksmq!", zap.Any("topic", consumer.Topic), zap.Any("groupName", consumer.GroupName), zap.Any("error", err))
} }
topic = consumer.Topic
}
if topic != "" {
err := rmq.DestroyTopic(topic)
if err != nil {
log.Warn("Rocksmq DestroyTopic failed!", zap.Any("topic", topic), zap.Any("error", err))
}
} }
return true return true
}) })
rmq.storeMu.Lock() rmq.storeMu.Lock()
defer rmq.storeMu.Unlock() defer rmq.storeMu.Unlock()
rmq.kv.Close()
rmq.store.Close() rmq.store.Close()
} }
@ -246,56 +302,26 @@ func (rmq *rocksmq) stopRetention() {
} }
} }
func (rmq *rocksmq) checkKeyExist(key string) bool {
val, _ := rmq.kv.Load(key)
return val != ""
}
// CreateTopic writes initialized messages for topic in rocksdb // CreateTopic writes initialized messages for topic in rocksdb
func (rmq *rocksmq) CreateTopic(topicName string) error { func (rmq *rocksmq) CreateTopic(topicName string) error {
if rmq.isClosed() { if rmq.isClosed() {
return errors.New(RmqNotServingErrMsg) return errors.New(RmqNotServingErrMsg)
} }
start := time.Now() start := time.Now()
beginKey := topicName + "/begin_id"
endKey := topicName + "/end_id"
// Check if topic exist // TopicBeginIDTitle is the only identifier of a topic exist or not
if rmq.checkKeyExist(beginKey) || rmq.checkKeyExist(endKey) { topicIDKey := TopicIDTitle + topicName
log.Warn("RocksMQ: " + beginKey + " or " + endKey + " existed.") val, err := rmq.kv.Load(topicIDKey)
if err != nil {
return err
}
if val != "" {
return nil return nil
} }
// TODO change rmq kv save logic into a batch
err := rmq.kv.Save(beginKey, "0")
if err != nil {
return err
}
err = rmq.kv.Save(endKey, "0")
if err != nil {
return err
}
if _, ok := topicMu.Load(topicName); !ok { if _, ok := topicMu.Load(topicName); !ok {
topicMu.Store(topicName, new(sync.Mutex)) topicMu.Store(topicName, new(sync.Mutex))
} }
if _, ok := rmq.ackedMu.Load(topicName); !ok {
rmq.ackedMu.Store(topicName, new(sync.Mutex))
}
// Initialize retention infos
// Initialize acked size to 0 for topic
ackedSizeKey := AckedSizeTitle + topicName
err = rmq.kv.Save(ackedSizeKey, "0")
if err != nil {
return err
}
// Initialize topic begin id to defaultMessageID
topicBeginIDKey := TopicBeginIDTitle + topicName
err = rmq.kv.Save(topicBeginIDKey, DefaultMessageID)
if err != nil {
return err
}
// Initialize topic message size to 0 // Initialize topic message size to 0
msgSizeKey := MessageSizeTitle + topicName msgSizeKey := MessageSizeTitle + topicName
@ -304,9 +330,16 @@ func (rmq *rocksmq) CreateTopic(topicName string) error {
return err return err
} }
// Initialize topic id to its create Tme, we don't really use it for now
nowTs := strconv.FormatInt(time.Now().Unix(), 10)
err = rmq.kv.Save(topicIDKey, nowTs)
if err != nil {
return err
}
rmq.retentionInfo.mutex.Lock() rmq.retentionInfo.mutex.Lock()
defer rmq.retentionInfo.mutex.Unlock() defer rmq.retentionInfo.mutex.Unlock()
rmq.retentionInfo.topics.Store(topicName, time.Now().Unix()) rmq.retentionInfo.topicRetetionTime.Store(topicName, time.Now().Unix())
log.Debug("Rocksmq create topic successfully ", zap.String("topic", topicName), zap.Int64("elapsed", time.Since(start).Milliseconds())) log.Debug("Rocksmq create topic successfully ", zap.String("topic", topicName), zap.Int64("elapsed", time.Since(start).Milliseconds()))
return nil return nil
} }
@ -324,29 +357,55 @@ func (rmq *rocksmq) DestroyTopic(topicName string) error {
} }
lock.Lock() lock.Lock()
defer lock.Unlock() defer lock.Unlock()
beginKey := topicName + "/begin_id"
endKey := topicName + "/end_id"
var removedKeys []string
rmq.consumers.Delete(topicName) rmq.consumers.Delete(topicName)
ackedSizeKey := AckedSizeTitle + topicName // clean the topic data it self
topicBeginIDKey := TopicBeginIDTitle + topicName fixChanName, err := fixChannelName(topicName)
if err != nil {
return err
}
err = rmq.kv.RemoveWithPrefix(fixChanName)
if err != nil {
return err
}
// just for clean up old topics, for new topics this is not required // clean page size info
lastRetTsKey := LastRetTsTitle + topicName pageMsgSizeKey := constructKey(PageMsgSizeTitle, topicName)
err = rmq.kv.RemoveWithPrefix(pageMsgSizeKey)
if err != nil {
return err
}
// clean page ts info
pageMsgTsKey := constructKey(PageTsTitle, topicName)
err = rmq.kv.RemoveWithPrefix(pageMsgTsKey)
if err != nil {
return err
}
// cleaned acked ts info
ackedTsKey := constructKey(AckedTsTitle, topicName)
err = rmq.kv.RemoveWithPrefix(ackedTsKey)
if err != nil {
return err
}
// topic info
topicIDKey := TopicIDTitle + topicName
// message size of this topic
msgSizeKey := MessageSizeTitle + topicName msgSizeKey := MessageSizeTitle + topicName
var removedKeys []string
removedKeys = append(removedKeys, beginKey, endKey, ackedSizeKey, topicBeginIDKey, lastRetTsKey, msgSizeKey) removedKeys = append(removedKeys, topicIDKey, msgSizeKey)
// Batch remove, atomic operation // Batch remove, atomic operation
err := rmq.kv.MultiRemove(removedKeys) err = rmq.kv.MultiRemove(removedKeys)
if err != nil { if err != nil {
return err return err
} }
// clean up retention info // clean up retention info
topicMu.Delete(topicName) topicMu.Delete(topicName)
rmq.retentionInfo.topics.Delete(topicName) rmq.retentionInfo.topicRetetionTime.Delete(topicName)
// clean up reader // clean up reader
if val, ok := rmq.readers.LoadAndDelete(topicName); ok { if val, ok := rmq.readers.LoadAndDelete(topicName); ok {
@ -359,18 +418,19 @@ func (rmq *rocksmq) DestroyTopic(topicName string) error {
} }
// ExistConsumerGroup check if a consumer exists and return the existed consumer // ExistConsumerGroup check if a consumer exists and return the existed consumer
func (rmq *rocksmq) ExistConsumerGroup(topicName, groupName string) (bool, *Consumer) { func (rmq *rocksmq) ExistConsumerGroup(topicName, groupName string) (bool, *Consumer, error) {
key := constructCurrentID(topicName, groupName) key := constructCurrentID(topicName, groupName)
if rmq.checkKeyExist(key) { _, ok := rmq.consumersID.Load(key)
if ok {
if vals, ok := rmq.consumers.Load(topicName); ok { if vals, ok := rmq.consumers.Load(topicName); ok {
for _, v := range vals.([]*Consumer) { for _, v := range vals.([]*Consumer) {
if v.GroupName == groupName { if v.GroupName == groupName {
return true, v return true, v, nil
} }
} }
} }
} }
return false, nil return false, nil, nil
} }
// CreateConsumerGroup creates an nonexistent consumer group for topic // CreateConsumerGroup creates an nonexistent consumer group for topic
@ -380,14 +440,11 @@ func (rmq *rocksmq) CreateConsumerGroup(topicName, groupName string) error {
} }
start := time.Now() start := time.Now()
key := constructCurrentID(topicName, groupName) key := constructCurrentID(topicName, groupName)
if rmq.checkKeyExist(key) { _, ok := rmq.consumersID.Load(key)
log.Debug("RMQ CreateConsumerGroup key already exists", zap.String("key", key)) if ok {
return nil return fmt.Errorf("RMQ CreateConsumerGroup key already exists, key = %s", key)
}
err := rmq.kv.Save(key, DefaultMessageID)
if err != nil {
return err
} }
rmq.consumersID.Store(key, DefaultMessageID)
log.Debug("Rocksmq create consumer group successfully ", zap.String("topic", topicName), log.Debug("Rocksmq create consumer group successfully ", zap.String("topic", topicName),
zap.String("group", groupName), zap.String("group", groupName),
zap.Int64("elapsed", time.Since(start).Milliseconds())) zap.Int64("elapsed", time.Since(start).Milliseconds()))
@ -395,15 +452,15 @@ func (rmq *rocksmq) CreateConsumerGroup(topicName, groupName string) error {
} }
// RegisterConsumer registers a consumer in rocksmq consumers // RegisterConsumer registers a consumer in rocksmq consumers
func (rmq *rocksmq) RegisterConsumer(consumer *Consumer) { func (rmq *rocksmq) RegisterConsumer(consumer *Consumer) error {
if rmq.isClosed() { if rmq.isClosed() {
return return errors.New(RmqNotServingErrMsg)
} }
start := time.Now() start := time.Now()
if vals, ok := rmq.consumers.Load(consumer.Topic); ok { if vals, ok := rmq.consumers.Load(consumer.Topic); ok {
for _, v := range vals.([]*Consumer) { for _, v := range vals.([]*Consumer) {
if v.GroupName == consumer.GroupName { if v.GroupName == consumer.GroupName {
return return nil
} }
} }
consumers := vals.([]*Consumer) consumers := vals.([]*Consumer)
@ -415,10 +472,19 @@ func (rmq *rocksmq) RegisterConsumer(consumer *Consumer) {
rmq.consumers.Store(consumer.Topic, consumers) rmq.consumers.Store(consumer.Topic, consumers)
} }
log.Debug("Rocksmq register consumer successfully ", zap.String("topic", consumer.Topic), zap.Int64("elapsed", time.Since(start).Milliseconds())) log.Debug("Rocksmq register consumer successfully ", zap.String("topic", consumer.Topic), zap.Int64("elapsed", time.Since(start).Milliseconds()))
return nil
} }
// DestroyConsumerGroup removes a consumer group from rocksdb_kv // DestroyConsumerGroup removes a consumer group from rocksdb_kv
func (rmq *rocksmq) DestroyConsumerGroup(topicName, groupName string) error { func (rmq *rocksmq) DestroyConsumerGroup(topicName, groupName string) error {
if rmq.isClosed() {
return errors.New(RmqNotServingErrMsg)
}
return rmq.destroyConsumerGroupInternal(topicName, groupName)
}
// DestroyConsumerGroup removes a consumer group from rocksdb_kv
func (rmq *rocksmq) destroyConsumerGroupInternal(topicName, groupName string) error {
start := time.Now() start := time.Now()
ll, ok := topicMu.Load(topicName) ll, ok := topicMu.Load(topicName)
if !ok { if !ok {
@ -431,11 +497,7 @@ func (rmq *rocksmq) DestroyConsumerGroup(topicName, groupName string) error {
lock.Lock() lock.Lock()
defer lock.Unlock() defer lock.Unlock()
key := constructCurrentID(topicName, groupName) key := constructCurrentID(topicName, groupName)
rmq.consumersID.Delete(key)
err := rmq.kv.Remove(key)
if err != nil {
return err
}
if vals, ok := rmq.consumers.Load(topicName); ok { if vals, ok := rmq.consumers.Load(topicName); ok {
consumers := vals.([]*Consumer) consumers := vals.([]*Consumer)
for index, v := range consumers { for index, v := range consumers {
@ -484,7 +546,7 @@ func (rmq *rocksmq) Produce(topicName string, messages []ProducerMessage) ([]Uni
return []UniqueID{}, errors.New("Obtained id length is not equal that of message") return []UniqueID{}, errors.New("Obtained id length is not equal that of message")
} }
/* Step I: Insert data to store system */ // Insert data to store system
batch := gorocksdb.NewWriteBatch() batch := gorocksdb.NewWriteBatch()
defer batch.Destroy() defer batch.Destroy()
msgSizes := make(map[UniqueID]int64) msgSizes := make(map[UniqueID]int64)
@ -495,7 +557,6 @@ func (rmq *rocksmq) Produce(topicName string, messages []ProducerMessage) ([]Uni
if err != nil { if err != nil {
return []UniqueID{}, err return []UniqueID{}, err
} }
batch.Put([]byte(key), messages[i].Payload) batch.Put([]byte(key), messages[i].Payload)
msgIDs[i] = msgID msgIDs[i] = msgID
msgSizes[msgID] = int64(len(messages[i].Payload)) msgSizes[msgID] = int64(len(messages[i].Payload))
@ -509,29 +570,6 @@ func (rmq *rocksmq) Produce(topicName string, messages []ProducerMessage) ([]Uni
return []UniqueID{}, err return []UniqueID{}, err
} }
/* Step II: Update meta data to kv system */
kvChannelBeginID := topicName + "/begin_id"
beginIDValue, err := rmq.kv.Load(kvChannelBeginID)
if err != nil {
log.Debug("RocksMQ: load " + kvChannelBeginID + " failed")
return []UniqueID{}, err
}
kvValues := make(map[string]string)
if beginIDValue == "0" {
kvValues[kvChannelBeginID] = strconv.FormatInt(idStart, 10)
}
kvChannelEndID := topicName + "/end_id"
kvValues[kvChannelEndID] = strconv.FormatInt(idEnd, 10)
err = rmq.kv.MultiSave(kvValues)
if err != nil {
log.Debug("RocksMQ: multisave failed")
return []UniqueID{}, err
}
if vals, ok := rmq.consumers.Load(topicName); ok { if vals, ok := rmq.consumers.Load(topicName); ok {
for _, v := range vals.([]*Consumer) { for _, v := range vals.([]*Consumer) {
select { select {
@ -554,12 +592,12 @@ func (rmq *rocksmq) Produce(topicName string, messages []ProducerMessage) ([]Uni
} }
// Update message page info // Update message page info
// TODO(yukun): Should this be in a go routine
err = rmq.updatePageInfo(topicName, msgIDs, msgSizes) err = rmq.updatePageInfo(topicName, msgIDs, msgSizes)
if err != nil { if err != nil {
return []UniqueID{}, err return []UniqueID{}, err
} }
// TODO add this to monitor metrics
getProduceTime := time.Since(start).Milliseconds() getProduceTime := time.Since(start).Milliseconds()
if getLockTime > 200 || getProduceTime > 200 { if getLockTime > 200 || getProduceTime > 200 {
log.Warn("rocksmq produce too slowly", zap.String("topic", topicName), log.Warn("rocksmq produce too slowly", zap.String("topic", topicName),
@ -578,10 +616,10 @@ func (rmq *rocksmq) updatePageInfo(topicName string, msgIDs []UniqueID, msgSizes
if err != nil { if err != nil {
return err return err
} }
fixedPageSizeKey, err := constructKey(PageMsgSizeTitle, topicName) fixedPageSizeKey := constructKey(PageMsgSizeTitle, topicName)
if err != nil { fixedPageTsKey := constructKey(PageTsTitle, topicName)
return err nowTs := strconv.FormatInt(time.Now().Unix(), 10)
} mutateBuffer := make(map[string]string)
for _, id := range msgIDs { for _, id := range msgIDs {
msgSize := msgSizes[id] msgSize := msgSizes[id]
if curMsgSize+msgSize > RocksmqPageSize { if curMsgSize+msgSize > RocksmqPageSize {
@ -590,17 +628,16 @@ func (rmq *rocksmq) updatePageInfo(topicName string, msgIDs []UniqueID, msgSizes
pageEndID := id pageEndID := id
// Update page message size for current page. key is page end ID // Update page message size for current page. key is page end ID
pageMsgSizeKey := fixedPageSizeKey + "/" + strconv.FormatInt(pageEndID, 10) pageMsgSizeKey := fixedPageSizeKey + "/" + strconv.FormatInt(pageEndID, 10)
err := rmq.kv.Save(pageMsgSizeKey, strconv.FormatInt(newPageSize, 10)) mutateBuffer[pageMsgSizeKey] = strconv.FormatInt(newPageSize, 10)
if err != nil { pageTsKey := fixedPageTsKey + "/" + strconv.FormatInt(pageEndID, 10)
return err mutateBuffer[pageTsKey] = nowTs
}
curMsgSize = 0 curMsgSize = 0
} else { } else {
curMsgSize += msgSize curMsgSize += msgSize
} }
} }
// Update message size to current message size mutateBuffer[msgSizeKey] = strconv.FormatInt(curMsgSize, 10)
err = rmq.kv.Save(msgSizeKey, strconv.FormatInt(curMsgSize, 10)) err = rmq.kv.MultiSave(mutateBuffer)
return err return err
} }
@ -626,12 +663,8 @@ func (rmq *rocksmq) Consume(topicName string, groupName string, n int) ([]Consum
getLockTime := time.Since(start).Milliseconds() getLockTime := time.Since(start).Milliseconds()
metaKey := constructCurrentID(topicName, groupName) metaKey := constructCurrentID(topicName, groupName)
currentID, err := rmq.kv.Load(metaKey) currentID, ok := rmq.consumersID.Load(metaKey)
if err != nil { if !ok {
log.Debug("RocksMQ: load " + metaKey + " failed")
return nil, err
}
if currentID == "" {
return nil, fmt.Errorf("currentID of topicName=%s, groupName=%s not exist", topicName, groupName) return nil, fmt.Errorf("currentID of topicName=%s, groupName=%s not exist", topicName, groupName)
} }
@ -653,7 +686,7 @@ func (rmq *rocksmq) Consume(topicName string, groupName string, n int) ([]Consum
if currentID == DefaultMessageID { if currentID == DefaultMessageID {
dataKey = fixChanName + "/" dataKey = fixChanName + "/"
} else { } else {
dataKey = fixChanName + "/" + currentID dataKey = fixChanName + "/" + strconv.FormatInt(currentID.(int64), 10)
} }
iter.Seek([]byte(dataKey)) iter.Seek([]byte(dataKey))
@ -666,7 +699,7 @@ func (rmq *rocksmq) Consume(topicName string, groupName string, n int) ([]Consum
offset++ offset++
msgID, err := strconv.ParseInt(strKey[FixedChannelNameLen+1:], 10, 64) msgID, err := strconv.ParseInt(strKey[FixedChannelNameLen+1:], 10, 64)
if err != nil { if err != nil {
log.Debug("RocksMQ: parse int " + strKey[FixedChannelNameLen+1:] + " failed") log.Warn("RocksMQ: parse int " + strKey[FixedChannelNameLen+1:] + " failed")
val.Free() val.Free()
return nil, err return nil, err
} }
@ -684,6 +717,10 @@ func (rmq *rocksmq) Consume(topicName string, groupName string, n int) ([]Consum
consumerMessage = append(consumerMessage, msg) consumerMessage = append(consumerMessage, msg)
val.Free() val.Free()
} }
// if iterate fail
if err := iter.Err(); err != nil {
return nil, err
}
// When already consume to last mes, an empty slice will be returned // When already consume to last mes, an empty slice will be returned
if len(consumerMessage) == 0 { if len(consumerMessage) == 0 {
@ -696,12 +733,10 @@ func (rmq *rocksmq) Consume(topicName string, groupName string, n int) ([]Consum
consumedIDs = append(consumedIDs, msg.MsgID) consumedIDs = append(consumedIDs, msg.MsgID)
} }
newID := consumedIDs[len(consumedIDs)-1] newID := consumedIDs[len(consumedIDs)-1]
err = rmq.moveConsumePos(topicName, groupName, newID+1) rmq.moveConsumePos(topicName, groupName, newID+1)
if err != nil {
return nil, err
}
go rmq.updateAckedInfo(topicName, groupName, consumedIDs) rmq.updateAckedInfo(topicName, groupName, consumedIDs)
// TODO add this to monitor metrics
getConsumeTime := time.Since(start).Milliseconds() getConsumeTime := time.Since(start).Milliseconds()
if getLockTime > 200 || getConsumeTime > 200 { if getLockTime > 200 || getConsumeTime > 200 {
log.Warn("rocksmq consume too slowly", zap.String("topic", topicName), log.Warn("rocksmq consume too slowly", zap.String("topic", topicName),
@ -715,10 +750,12 @@ func (rmq *rocksmq) seek(topicName string, groupName string, msgID UniqueID) err
rmq.storeMu.Lock() rmq.storeMu.Lock()
defer rmq.storeMu.Unlock() defer rmq.storeMu.Unlock()
key := constructCurrentID(topicName, groupName) key := constructCurrentID(topicName, groupName)
if !rmq.checkKeyExist(key) { _, ok := rmq.consumersID.Load(key)
if !ok {
log.Warn("RocksMQ: channel " + key + " not exists") log.Warn("RocksMQ: channel " + key + " not exists")
return fmt.Errorf("consumerGroup %s, channel %s not exists", groupName, topicName) return fmt.Errorf("ConsumerGroup %s, channel %s not exists", groupName, topicName)
} }
storeKey, err := combKey(topicName, msgID) storeKey, err := combKey(topicName, msgID)
if err != nil { if err != nil {
log.Warn("RocksMQ: combKey(" + topicName + "," + strconv.FormatInt(msgID, 10) + ") failed") log.Warn("RocksMQ: combKey(" + topicName + "," + strconv.FormatInt(msgID, 10) + ") failed")
@ -727,28 +764,26 @@ func (rmq *rocksmq) seek(topicName string, groupName string, msgID UniqueID) err
opts := gorocksdb.NewDefaultReadOptions() opts := gorocksdb.NewDefaultReadOptions()
defer opts.Destroy() defer opts.Destroy()
val, err := rmq.store.Get(opts, []byte(storeKey)) val, err := rmq.store.Get(opts, []byte(storeKey))
defer val.Free()
if err != nil { if err != nil {
log.Warn("RocksMQ: get " + storeKey + " failed") log.Warn("RocksMQ: get " + storeKey + " failed")
return err return err
} }
defer val.Free()
if !val.Exists() { if !val.Exists() {
log.Warn("RocksMQ: trying to seek to no exist position, reset current id",
zap.String("topic", topicName), zap.String("group", groupName), zap.Int64("msgId", msgID))
rmq.moveConsumePos(topicName, groupName, DefaultMessageID)
//skip seek if key is not found, this is the behavior as pulsar //skip seek if key is not found, this is the behavior as pulsar
return nil return nil
} }
/* Step II: update current_id */
/* Step II: Save current_id in kv */ rmq.moveConsumePos(topicName, groupName, msgID)
return rmq.moveConsumePos(topicName, groupName, msgID) return nil
} }
func (rmq *rocksmq) moveConsumePos(topicName string, groupName string, msgID UniqueID) error { func (rmq *rocksmq) moveConsumePos(topicName string, groupName string, msgID UniqueID) {
key := constructCurrentID(topicName, groupName) key := constructCurrentID(topicName, groupName)
err := rmq.kv.Save(key, strconv.FormatInt(msgID, 10)) rmq.consumersID.Store(key, msgID)
if err != nil {
log.Warn("RocksMQ: save " + key + " failed")
return err
}
return nil
} }
// Seek updates the current id to the given msgID // Seek updates the current id to the given msgID
@ -779,8 +814,9 @@ func (rmq *rocksmq) SeekToLatest(topicName, groupName string) error {
rmq.storeMu.Lock() rmq.storeMu.Lock()
defer rmq.storeMu.Unlock() defer rmq.storeMu.Unlock()
key := constructCurrentID(topicName, groupName) key := constructCurrentID(topicName, groupName)
if !rmq.checkKeyExist(key) { _, ok := rmq.consumersID.Load(key)
log.Debug("RocksMQ: channel " + key + " not exists") if !ok {
log.Warn("RocksMQ: channel " + key + " not exists")
return fmt.Errorf("ConsumerGroup %s, channel %s not exists", groupName, topicName) return fmt.Errorf("ConsumerGroup %s, channel %s not exists", groupName, topicName)
} }
@ -794,6 +830,10 @@ func (rmq *rocksmq) SeekToLatest(topicName, groupName string) error {
// 0 is the ASC value of "/" + 1 // 0 is the ASC value of "/" + 1
iter.SeekForPrev([]byte(fixChanName + "0")) iter.SeekForPrev([]byte(fixChanName + "0"))
// if iterate fail
if err := iter.Err(); err != nil {
return err
}
// should find the last key we written into, start with fixChanName/ // should find the last key we written into, start with fixChanName/
// if not find, start from 0 // if not find, start from 0
if !iter.Valid() { if !iter.Valid() {
@ -813,7 +853,8 @@ func (rmq *rocksmq) SeekToLatest(topicName, groupName string) error {
return err return err
} }
// current msgID should not be included // current msgID should not be included
return rmq.moveConsumePos(topicName, groupName, msgID+1) rmq.moveConsumePos(topicName, groupName, msgID+1)
return nil
} }
// Notify sends a mutex in MsgMutex channel to tell consumers to consume // Notify sends a mutex in MsgMutex channel to tell consumers to consume
@ -837,48 +878,25 @@ func (rmq *rocksmq) updateAckedInfo(topicName, groupName string, ids []UniqueID)
if len(ids) == 0 { if len(ids) == 0 {
return nil return nil
} }
ll, ok := topicMu.Load(topicName)
if !ok {
return fmt.Errorf("topic name = %s not exist", topicName)
}
lock, ok := ll.(*sync.Mutex)
if !ok {
return fmt.Errorf("get mutex failed, topic name = %s", topicName)
}
lock.Lock()
defer lock.Unlock()
firstID := ids[0] firstID := ids[0]
lastID := ids[len(ids)-1] lastID := ids[len(ids)-1]
fixedBeginIDKey, err := constructKey(BeginIDTitle, topicName) // 1. Try to get the page id between first ID and last ID of ids
if err != nil { pageMsgPrefix := constructKey(PageMsgSizeTitle, topicName)
return err
}
// 1. Update begin_id for the consumer_group
beginIDKey := fixedBeginIDKey + "/" + groupName
err = rmq.kv.Save(beginIDKey, strconv.FormatInt(lastID, 10))
if err != nil {
return err
}
// 2. Try to get the page id between first ID and last ID of ids
pageMsgPrefix, err := constructKey(PageMsgSizeTitle, topicName)
if err != nil {
return err
}
readOpts := gorocksdb.NewDefaultReadOptions() readOpts := gorocksdb.NewDefaultReadOptions()
defer readOpts.Destroy() defer readOpts.Destroy()
readOpts.SetPrefixSameAsStart(true) pageMsgFirstKey := pageMsgPrefix + "/" + strconv.FormatInt(firstID, 10)
// set last key by lastID
pageMsgLastKey := pageMsgPrefix + "/" + strconv.FormatInt(lastID+1, 10)
readOpts.SetIterateUpperBound([]byte(pageMsgLastKey))
iter := rmq.kv.(*rocksdbkv.RocksdbKV).DB.NewIterator(readOpts) iter := rmq.kv.(*rocksdbkv.RocksdbKV).DB.NewIterator(readOpts)
defer iter.Close() defer iter.Close()
var pageIDs []UniqueID var pageIDs []UniqueID
pageMsgKey := pageMsgPrefix + "/" + strconv.FormatInt(firstID, 10)
for iter.Seek([]byte(pageMsgKey)); iter.Valid(); iter.Next() { for iter.Seek([]byte(pageMsgFirstKey)); iter.Valid(); iter.Next() {
key := iter.Key() key := iter.Key()
pageID, err := strconv.ParseInt(string(key.Data())[FixedChannelNameLen+1:], 10, 64) pageID, err := parsePageID(string(key.Data()))
if key != nil { if key != nil {
key.Free() key.Free()
} }
@ -891,80 +909,47 @@ func (rmq *rocksmq) updateAckedInfo(topicName, groupName string, ids []UniqueID)
break break
} }
} }
if err := iter.Err(); err != nil {
return err
}
if len(pageIDs) == 0 { if len(pageIDs) == 0 {
return nil return nil
} }
fixedAckedTsKey := constructKey(AckedTsTitle, topicName)
fixedAckedTsKey, err := constructKey(AckedTsTitle, topicName) // 2. Update acked ts and acked size for pageIDs
if err != nil {
return err
}
// 3. Update acked ts and acked size for pageIDs
if vals, ok := rmq.consumers.Load(topicName); ok { if vals, ok := rmq.consumers.Load(topicName); ok {
var minBeginID int64 = math.MaxInt64
consumers, ok := vals.([]*Consumer) consumers, ok := vals.([]*Consumer)
if !ok || len(consumers) == 0 { if !ok || len(consumers) == 0 {
return nil return nil
} }
for _, v := range consumers { // update consumer id
curBeginIDKey := path.Join(fixedBeginIDKey, v.GroupName) for _, consumer := range consumers {
curBeginIDVal, err := rmq.kv.Load(curBeginIDKey) if consumer.GroupName == groupName {
if err != nil { consumer.beginID = lastID
return err break
} }
curBeginID, err := strconv.ParseInt(curBeginIDVal, 10, 64) }
if err != nil {
return err // find min id of all consumer
} var minBeginID int64 = math.MaxInt64
if curBeginID < minBeginID { for _, consumer := range consumers {
minBeginID = curBeginID if consumer.beginID < minBeginID {
minBeginID = consumer.beginID
} }
} }
nowTs := strconv.FormatInt(time.Now().Unix(), 10) nowTs := strconv.FormatInt(time.Now().Unix(), 10)
ackedTsKvs := make(map[string]string) ackedTsKvs := make(map[string]string)
totalAckMsgSize := int64(0) // update ackedTs, if page is all acked, then ackedTs is set
fixedPageSizeKey, err := constructKey(PageMsgSizeTitle, topicName)
if err != nil {
return err
}
for _, pID := range pageIDs { for _, pID := range pageIDs {
if pID <= minBeginID { if pID <= minBeginID {
// Update acked info for message pID // Update acked info for message pID
pageAckedTsKey := path.Join(fixedAckedTsKey, strconv.FormatInt(pID, 10)) pageAckedTsKey := path.Join(fixedAckedTsKey, strconv.FormatInt(pID, 10))
ackedTsKvs[pageAckedTsKey] = nowTs ackedTsKvs[pageAckedTsKey] = nowTs
// get current page message size
pageMsgSizeKey := path.Join(fixedPageSizeKey, strconv.FormatInt(pID, 10))
pageMsgSizeVal, err := rmq.kv.Load(pageMsgSizeKey)
if err != nil {
return err
}
pageMsgSize, err := strconv.ParseInt(pageMsgSizeVal, 10, 64)
if err != nil {
return err
}
totalAckMsgSize += pageMsgSize
} }
} }
err = rmq.kv.MultiSave(ackedTsKvs) err := rmq.kv.MultiSave(ackedTsKvs)
if err != nil {
return err
}
ackedSizeKey := AckedSizeTitle + topicName
ackedSizeVal, err := rmq.kv.Load(ackedSizeKey)
if err != nil {
return err
}
ackedSize, err := strconv.ParseInt(ackedSizeVal, 10, 64)
if err != nil {
return err
}
ackedSize += totalAckMsgSize
err = rmq.kv.Save(ackedSizeKey, strconv.FormatInt(ackedSize, 10))
if err != nil { if err != nil {
return err return err
} }
@ -983,9 +968,21 @@ func (rmq *rocksmq) CreateReader(topicName string, startMsgID UniqueID, messageI
readOpts := gorocksdb.NewDefaultReadOptions() readOpts := gorocksdb.NewDefaultReadOptions()
readOpts.SetPrefixSameAsStart(true) readOpts.SetPrefixSameAsStart(true)
iter := rmq.store.NewIterator(readOpts) iter := rmq.store.NewIterator(readOpts)
fixChanName, err := fixChannelName(topicName)
if err != nil {
log.Debug("RocksMQ: fixChannelName " + topicName + " failed")
return "", err
}
dataKey := path.Join(fixChanName, strconv.FormatInt(startMsgID, 10))
iter.Seek([]byte(dataKey))
// if iterate fail
if err := iter.Err(); err != nil {
return "", err
}
nowTs, err := getNowTs(rmq.idAllocator) nowTs, err := getNowTs(rmq.idAllocator)
if err != nil { if err != nil {
return "", errors.New("can't get current ts from rocksmq idAllocator") return "", errors.New("Can't get current ts from rocksmq idAllocator")
} }
readerName := subscriptionRolePrefix + ReaderNamePrefix + strconv.FormatInt(nowTs, 10) readerName := subscriptionRolePrefix + ReaderNamePrefix + strconv.FormatInt(nowTs, 10)
@ -999,7 +996,6 @@ func (rmq *rocksmq) CreateReader(topicName string, startMsgID UniqueID, messageI
messageIDInclusive: messageIDInclusive, messageIDInclusive: messageIDInclusive,
readerMutex: make(chan struct{}, 1), readerMutex: make(chan struct{}, 1),
} }
reader.Seek(startMsgID)
if vals, ok := rmq.readers.Load(topicName); ok { if vals, ok := rmq.readers.Load(topicName); ok {
readers := vals.([]*rocksmqReader) readers := vals.([]*rocksmqReader)
readers = append(readers, reader) readers = append(readers, reader)
@ -1024,17 +1020,17 @@ func (rmq *rocksmq) getReader(topicName, readerName string) *rocksmqReader {
} }
// ReaderSeek seek a reader to the pointed position // ReaderSeek seek a reader to the pointed position
func (rmq *rocksmq) ReaderSeek(topicName string, readerName string, msgID UniqueID) { func (rmq *rocksmq) ReaderSeek(topicName string, readerName string, msgID UniqueID) error {
if rmq.isClosed() { if rmq.isClosed() {
return return errors.New(RmqNotServingErrMsg)
} }
reader := rmq.getReader(topicName, readerName) reader := rmq.getReader(topicName, readerName)
if reader == nil { if reader == nil {
log.Warn("reader not exist", zap.String("topic", topicName), zap.String("readerName", readerName)) log.Warn("reader not exist", zap.String("topic", topicName), zap.String("readerName", readerName))
return return fmt.Errorf("reader not exist, topic %s, reader %s", topicName, readerName)
} }
reader.Seek(msgID) reader.Seek(msgID)
return nil
} }
// Next get the next message of reader // Next get the next message of reader

42
internal/util/rocksmq/server/rocksmq/rocksmq_impl_test.go
View File

@ -78,9 +78,8 @@ func TestRocksmq_RegisterConsumer(t *testing.T) {
idAllocator := InitIDAllocator(kvPath) idAllocator := InitIDAllocator(kvPath)
rocksdbPath := rmqPath + dbPathSuffix + suffix rocksdbPath := rmqPath + dbPathSuffix + suffix
defer os.RemoveAll(rocksdbPath + kvSuffix)
defer os.RemoveAll(rocksdbPath) defer os.RemoveAll(rocksdbPath)
metaPath := rmqPath + metaPathSuffix + suffix
defer os.RemoveAll(metaPath)
rmq, err := NewRocksMQ(rocksdbPath, idAllocator) rmq, err := NewRocksMQ(rocksdbPath, idAllocator)
assert.NoError(t, err) assert.NoError(t, err)
@ -88,9 +87,14 @@ func TestRocksmq_RegisterConsumer(t *testing.T) {
topicName := "topic_register" topicName := "topic_register"
groupName := "group_register" groupName := "group_register"
_ = rmq.DestroyConsumerGroup(topicName, groupName)
err = rmq.CreateTopic(topicName)
assert.NoError(t, err)
defer rmq.DestroyTopic(topicName)
err = rmq.CreateConsumerGroup(topicName, groupName) err = rmq.CreateConsumerGroup(topicName, groupName)
assert.Nil(t, err) assert.Nil(t, err)
defer rmq.DestroyConsumerGroup(topicName, groupName)
consumer := &Consumer{ consumer := &Consumer{
Topic: topicName, Topic: topicName,
@ -98,10 +102,10 @@ func TestRocksmq_RegisterConsumer(t *testing.T) {
MsgMutex: make(chan struct{}), MsgMutex: make(chan struct{}),
} }
rmq.RegisterConsumer(consumer) rmq.RegisterConsumer(consumer)
exist, _ := rmq.ExistConsumerGroup(topicName, groupName) exist, _, _ := rmq.ExistConsumerGroup(topicName, groupName)
assert.Equal(t, exist, true) assert.Equal(t, exist, true)
dummyGrpName := "group_dummy" dummyGrpName := "group_dummy"
exist, _ = rmq.ExistConsumerGroup(topicName, dummyGrpName) exist, _, _ = rmq.ExistConsumerGroup(topicName, dummyGrpName)
assert.Equal(t, exist, false) assert.Equal(t, exist, false)
msgA := "a_message" msgA := "a_message"
@ -111,7 +115,7 @@ func TestRocksmq_RegisterConsumer(t *testing.T) {
_ = idAllocator.UpdateID() _ = idAllocator.UpdateID()
_, err = rmq.Produce(topicName, pMsgs) _, err = rmq.Produce(topicName, pMsgs)
assert.Error(t, err) assert.Nil(t, err)
rmq.Notify(topicName, groupName) rmq.Notify(topicName, groupName)
@ -129,25 +133,18 @@ func TestRocksmq_RegisterConsumer(t *testing.T) {
MsgMutex: make(chan struct{}), MsgMutex: make(chan struct{}),
} }
rmq.RegisterConsumer(consumer2) rmq.RegisterConsumer(consumer2)
topicMu.Delete(topicName)
topicMu.Store(topicName, topicName)
assert.Error(t, rmq.DestroyConsumerGroup(topicName, groupName))
err = rmq.DestroyConsumerGroup(topicName, groupName)
assert.Error(t, err)
} }
func TestRocksmq(t *testing.T) { func TestRocksmq_Basic(t *testing.T) {
suffix := "_rmq" suffix := "_rmq"
kvPath := rmqPath + kvPathSuffix + suffix kvPath := rmqPath + kvPathSuffix + suffix
defer os.RemoveAll(kvPath) defer os.RemoveAll(kvPath)
idAllocator := InitIDAllocator(kvPath) idAllocator := InitIDAllocator(kvPath)
rocksdbPath := rmqPath + dbPathSuffix + suffix rocksdbPath := rmqPath + dbPathSuffix + suffix
defer os.RemoveAll(rocksdbPath + kvSuffix)
defer os.RemoveAll(rocksdbPath) defer os.RemoveAll(rocksdbPath)
metaPath := rmqPath + metaPathSuffix + suffix
defer os.RemoveAll(metaPath)
rmq, err := NewRocksMQ(rocksdbPath, idAllocator) rmq, err := NewRocksMQ(rocksdbPath, idAllocator)
assert.Nil(t, err) assert.Nil(t, err)
defer rmq.Close() defer rmq.Close()
@ -181,7 +178,7 @@ func TestRocksmq(t *testing.T) {
assert.Nil(t, err) assert.Nil(t, err)
// double create consumer group // double create consumer group
err = rmq.CreateConsumerGroup(channelName, groupName) err = rmq.CreateConsumerGroup(channelName, groupName)
assert.Nil(t, err) assert.Error(t, err)
cMsgs, err := rmq.Consume(channelName, groupName, 1) cMsgs, err := rmq.Consume(channelName, groupName, 1)
assert.Nil(t, err) assert.Nil(t, err)
assert.Equal(t, len(cMsgs), 1) assert.Equal(t, len(cMsgs), 1)
@ -201,9 +198,8 @@ func TestRocksmq_Dummy(t *testing.T) {
idAllocator := InitIDAllocator(kvPath) idAllocator := InitIDAllocator(kvPath)
rocksdbPath := rmqPath + dbPathSuffix + suffix rocksdbPath := rmqPath + dbPathSuffix + suffix
defer os.RemoveAll(rocksdbPath + kvSuffix)
defer os.RemoveAll(rocksdbPath) defer os.RemoveAll(rocksdbPath)
metaPath := rmqPath + metaPathSuffix + suffix
defer os.RemoveAll(metaPath)
rmq, err := NewRocksMQ(rocksdbPath, idAllocator) rmq, err := NewRocksMQ(rocksdbPath, idAllocator)
assert.Nil(t, err) assert.Nil(t, err)
@ -216,8 +212,9 @@ func TestRocksmq_Dummy(t *testing.T) {
err = rmq.CreateTopic(channelName) err = rmq.CreateTopic(channelName)
assert.Nil(t, err) assert.Nil(t, err)
defer rmq.DestroyTopic(channelName) defer rmq.DestroyTopic(channelName)
// create topic twice should be ignored
err = rmq.CreateTopic(channelName) err = rmq.CreateTopic(channelName)
assert.NoError(t, err) assert.Nil(t, err)
channelName1 := "channel_dummy" channelName1 := "channel_dummy"
topicMu.Store(channelName1, new(sync.Mutex)) topicMu.Store(channelName1, new(sync.Mutex))
@ -258,7 +255,6 @@ func TestRocksmq_Dummy(t *testing.T) {
_, err = rmq.Consume(channelName, groupName1, 1) _, err = rmq.Consume(channelName, groupName1, 1)
assert.Error(t, err) assert.Error(t, err)
} }
func TestRocksmq_Seek(t *testing.T) { func TestRocksmq_Seek(t *testing.T) {
@ -268,9 +264,8 @@ func TestRocksmq_Seek(t *testing.T) {
idAllocator := InitIDAllocator(kvPath) idAllocator := InitIDAllocator(kvPath)
rocksdbPath := rmqPath + dbPathSuffix + suffix rocksdbPath := rmqPath + dbPathSuffix + suffix
defer os.RemoveAll(rocksdbPath + kvSuffix)
defer os.RemoveAll(rocksdbPath) defer os.RemoveAll(rocksdbPath)
metaPath := rmqPath + metaPathSuffix + suffix
defer os.RemoveAll(metaPath)
rmq, err := NewRocksMQ(rocksdbPath, idAllocator) rmq, err := NewRocksMQ(rocksdbPath, idAllocator)
assert.Nil(t, err) assert.Nil(t, err)
@ -824,7 +819,6 @@ func TestReader_CornerCase(t *testing.T) {
extraMsgs[0] = ProducerMessage{Payload: []byte(msg)} extraMsgs[0] = ProducerMessage{Payload: []byte(msg)}
extraIds, _ = rmq.Produce(channelName, extraMsgs) extraIds, _ = rmq.Produce(channelName, extraMsgs)
// assert.NoError(t, er) // assert.NoError(t, er)
fmt.Println(extraIds[0])
assert.Equal(t, 1, len(extraIds)) assert.Equal(t, 1, len(extraIds))
}() }()

319
internal/util/rocksmq/server/rocksmq/rocksmq_retention.go
View File

@ -20,30 +20,30 @@ import (
rocksdbkv "github.com/milvus-io/milvus/internal/kv/rocksdb" rocksdbkv "github.com/milvus-io/milvus/internal/kv/rocksdb"
"github.com/milvus-io/milvus/internal/log" "github.com/milvus-io/milvus/internal/log"
"github.com/milvus-io/milvus/internal/util/typeutil"
"github.com/tecbot/gorocksdb" "github.com/tecbot/gorocksdb"
"go.uber.org/zap" "go.uber.org/zap"
) )
// RocksmqRetentionTimeInMinutes is the time of retention // RocksmqRetentionTimeInMinutes is the time of retention
var RocksmqRetentionTimeInMinutes int64 = 10080 var RocksmqRetentionTimeInSecs int64 = 10080 * 60
// RocksmqRetentionSizeInMB is the size of retention // RocksmqRetentionSizeInMB is the size of retention
var RocksmqRetentionSizeInMB int64 = 8192 var RocksmqRetentionSizeInMB int64 = 8192
// Const value that used to convert unit // Const value that used to convert unit
const ( const (
MB = 1024 * 1024 MB = 1024 * 1024
MINUTE = 60
) )
// TickerTimeInSeconds is the time of expired check, default 10 minutes // TickerTimeInSeconds is the time of expired check, default 10 minutes
var TickerTimeInSeconds int64 = 10 * MINUTE var TickerTimeInSeconds int64 = 60
type retentionInfo struct { type retentionInfo struct {
// key is topic name, value is last retention type // key is topic name, value is last retention time
topics sync.Map topicRetetionTime sync.Map
mutex sync.RWMutex mutex sync.RWMutex
kv *rocksdbkv.RocksdbKV kv *rocksdbkv.RocksdbKV
db *gorocksdb.DB db *gorocksdb.DB
@ -55,21 +55,21 @@ type retentionInfo struct {
func initRetentionInfo(kv *rocksdbkv.RocksdbKV, db *gorocksdb.DB) (*retentionInfo, error) { func initRetentionInfo(kv *rocksdbkv.RocksdbKV, db *gorocksdb.DB) (*retentionInfo, error) {
ri := &retentionInfo{ ri := &retentionInfo{
topics: sync.Map{}, topicRetetionTime: sync.Map{},
mutex: sync.RWMutex{}, mutex: sync.RWMutex{},
kv: kv, kv: kv,
db: db, db: db,
closeCh: make(chan struct{}), closeCh: make(chan struct{}),
closeWg: sync.WaitGroup{}, closeWg: sync.WaitGroup{},
} }
// Get topic from topic begin id // Get topic from topic begin id
beginIDKeys, _, err := ri.kv.LoadWithPrefix(TopicBeginIDTitle) topicKeys, _, err := ri.kv.LoadWithPrefix(TopicIDTitle)
if err != nil { if err != nil {
return nil, err return nil, err
} }
for _, key := range beginIDKeys { for _, key := range topicKeys {
topic := key[len(TopicBeginIDTitle):] topic := key[len(TopicIDTitle):]
ri.topics.Store(topic, time.Now().Unix()) ri.topicRetetionTime.Store(topic, time.Now().Unix())
topicMu.Store(topic, new(sync.Mutex)) topicMu.Store(topic, new(sync.Mutex))
} }
return ri, nil return ri, nil
@ -79,7 +79,6 @@ func initRetentionInfo(kv *rocksdbkv.RocksdbKV, db *gorocksdb.DB) (*retentionInf
// Because loadRetentionInfo may need some time, so do this asynchronously. Finally start retention goroutine. // Because loadRetentionInfo may need some time, so do this asynchronously. Finally start retention goroutine.
func (ri *retentionInfo) startRetentionInfo() { func (ri *retentionInfo) startRetentionInfo() {
// var wg sync.WaitGroup // var wg sync.WaitGroup
ri.kv.ResetPrefixLength(FixedChannelNameLen)
ri.closeWg.Add(1) ri.closeWg.Add(1)
go ri.retention() go ri.retention()
} }
@ -98,9 +97,9 @@ func (ri *retentionInfo) retention() error {
return nil return nil
case t := <-ticker.C: case t := <-ticker.C:
timeNow := t.Unix() timeNow := t.Unix()
checkTime := atomic.LoadInt64(&RocksmqRetentionTimeInMinutes) * MINUTE / 10 checkTime := atomic.LoadInt64(&RocksmqRetentionTimeInSecs) / 10
ri.mutex.RLock() ri.mutex.RLock()
ri.topics.Range(func(k, v interface{}) bool { ri.topicRetetionTime.Range(func(k, v interface{}) bool {
topic, _ := k.(string) topic, _ := k.(string)
lastRetentionTs, ok := v.(int64) lastRetentionTs, ok := v.(int64)
if !ok { if !ok {
@ -112,6 +111,7 @@ func (ri *retentionInfo) retention() error {
if err != nil { if err != nil {
log.Warn("Retention expired clean failed", zap.Any("error", err)) log.Warn("Retention expired clean failed", zap.Any("error", err))
} }
ri.topicRetetionTime.Store(topic, timeNow)
} }
return true return true
}) })
@ -136,132 +136,75 @@ func (ri *retentionInfo) Stop() {
func (ri *retentionInfo) expiredCleanUp(topic string) error { func (ri *retentionInfo) expiredCleanUp(topic string) error {
log.Debug("Timeticker triggers an expiredCleanUp task for topic: " + topic) log.Debug("Timeticker triggers an expiredCleanUp task for topic: " + topic)
var deletedAckedSize int64 var deletedAckedSize int64
var startID UniqueID var pageCleaned UniqueID
var pageStartID UniqueID
var pageEndID UniqueID var pageEndID UniqueID
var err error var err error
fixedAckedTsKey, _ := constructKey(AckedTsTitle, topic) fixedAckedTsKey := constructKey(AckedTsTitle, topic)
// calculate total acked size, simply add all page info
totalAckedSize, err := ri.calculateTopicAckedSize(topic)
if err != nil {
return err
}
// Quick Path, No page to check
if totalAckedSize == 0 {
log.Debug("All messages are not expired, skip retention because no ack", zap.Any("topic", topic))
return nil
}
pageReadOpts := gorocksdb.NewDefaultReadOptions() pageReadOpts := gorocksdb.NewDefaultReadOptions()
defer pageReadOpts.Destroy() defer pageReadOpts.Destroy()
pageReadOpts.SetPrefixSameAsStart(true) pageMsgPrefix := constructKey(PageMsgSizeTitle, topic)
// ensure the iterator won't iterate to other topics
pageReadOpts.SetIterateUpperBound([]byte(typeutil.AddOne(pageMsgPrefix)))
pageIter := ri.kv.DB.NewIterator(pageReadOpts) pageIter := ri.kv.DB.NewIterator(pageReadOpts)
defer pageIter.Close() defer pageIter.Close()
pageMsgPrefix, _ := constructKey(PageMsgSizeTitle, topic)
pageIter.Seek([]byte(pageMsgPrefix)) pageIter.Seek([]byte(pageMsgPrefix))
if pageIter.Valid() { for ; pageIter.Valid(); pageIter.Next() {
pageStartID, err = strconv.ParseInt(string(pageIter.Key().Data())[FixedChannelNameLen+1:], 10, 64) pKey := pageIter.Key()
pageID, err := parsePageID(string(pKey.Data()))
if pKey != nil {
pKey.Free()
}
if err != nil { if err != nil {
return err return err
} }
ackedTsKey := fixedAckedTsKey + "/" + strconv.FormatInt(pageID, 10)
for ; pageIter.Valid(); pageIter.Next() { ackedTsVal, err := ri.kv.Load(ackedTsKey)
pKey := pageIter.Key() if err != nil {
pageID, err := strconv.ParseInt(string(pKey.Data())[FixedChannelNameLen+1:], 10, 64) return err
if pKey != nil { }
pKey.Free() // not acked page, TODO add TTL info there
if ackedTsVal == "" {
break
}
ackedTs, err := strconv.ParseInt(ackedTsVal, 10, 64)
if err != nil {
return err
}
if msgTimeExpiredCheck(ackedTs) {
pageEndID = pageID
pValue := pageIter.Value()
size, err := strconv.ParseInt(string(pValue.Data()), 10, 64)
if pValue != nil {
pValue.Free()
} }
if err != nil { if err != nil {
return err return err
} }
deletedAckedSize += size
ackedTsKey := fixedAckedTsKey + "/" + strconv.FormatInt(pageID, 10) pageCleaned++
ackedTsVal, err := ri.kv.Load(ackedTsKey) } else {
if err != nil { break
return err
}
if ackedTsVal == "" {
break
}
ackedTs, err := strconv.ParseInt(ackedTsVal, 10, 64)
if err != nil {
return err
}
if msgTimeExpiredCheck(ackedTs) {
pageEndID = pageID
pValue := pageIter.Value()
size, err := strconv.ParseInt(string(pValue.Data()), 10, 64)
if pValue != nil {
pValue.Free()
}
if err != nil {
return err
}
deletedAckedSize += size
} else {
break
}
} }
} }
if err := pageIter.Err(); err != nil {
// TODO(yukun): Remove ackedTs expiredCheck one by one
// ackedReadOpts := gorocksdb.NewDefaultReadOptions()
// defer ackedReadOpts.Destroy()
// ackedReadOpts.SetPrefixSameAsStart(true)
// ackedIter := ri.kv.DB.NewIterator(ackedReadOpts)
// defer ackedIter.Close()
// if err != nil {
// return err
// }
// ackedIter.Seek([]byte(fixedAckedTsKey))
// if !ackedIter.Valid() {
// return nil
// }
// startID, err = strconv.ParseInt(string(ackedIter.Key().Data())[FixedChannelNameLen+1:], 10, 64)
// if err != nil {
// return err
// }
// if endID > startID {
// newPos := fixedAckedTsKey + "/" + strconv.FormatInt(endID, 10)
// ackedIter.Seek([]byte(newPos))
// }
// for ; ackedIter.Valid(); ackedIter.Next() {
// aKey := ackedIter.Key()
// aValue := ackedIter.Value()
// ackedTs, err := strconv.ParseInt(string(aValue.Data()), 10, 64)
// if aValue != nil {
// aValue.Free()
// }
// if err != nil {
// if aKey != nil {
// aKey.Free()
// }
// return err
// }
// if msgTimeExpiredCheck(ackedTs) {
// endID, err = strconv.ParseInt(string(aKey.Data())[FixedChannelNameLen+1:], 10, 64)
// if aKey != nil {
// aKey.Free()
// }
// if err != nil {
// return err
// }
// } else {
// if aKey != nil {
// aKey.Free()
// }
// break
// }
// }
if pageEndID == 0 {
log.Debug("All messages are not time expired")
}
log.Debug("Expired check by retention time", zap.Any("topic", topic), zap.Any("pageEndID", pageEndID), zap.Any("deletedAckedSize", deletedAckedSize))
ackedSizeKey := AckedSizeTitle + topic
totalAckedSizeVal, err := ri.kv.Load(ackedSizeKey)
if err != nil {
return err
}
totalAckedSize, err := strconv.ParseInt(totalAckedSizeVal, 10, 64)
if err != nil {
return err return err
} }
log.Debug("Expired check by retention time", zap.Any("topic", topic),
zap.Any("pageEndID", pageEndID), zap.Any("deletedAckedSize", deletedAckedSize), zap.Any("pageCleaned", pageCleaned))
for ; pageIter.Valid(); pageIter.Next() { for ; pageIter.Valid(); pageIter.Next() {
pValue := pageIter.Value() pValue := pageIter.Value()
size, err := strconv.ParseInt(string(pValue.Data()), 10, 64) size, err := strconv.ParseInt(string(pValue.Data()), 10, 64)
@ -278,39 +221,97 @@ func (ri *retentionInfo) expiredCleanUp(topic string) error {
} }
curDeleteSize := deletedAckedSize + size curDeleteSize := deletedAckedSize + size
if msgSizeExpiredCheck(curDeleteSize, totalAckedSize) { if msgSizeExpiredCheck(curDeleteSize, totalAckedSize) {
pageEndID, err = strconv.ParseInt(pKeyStr[FixedChannelNameLen+1:], 10, 64) pageEndID, err = parsePageID(pKeyStr)
if err != nil { if err != nil {
return err return err
} }
deletedAckedSize += size deletedAckedSize += size
pageCleaned++
} else { } else {
break break
} }
} }
if err := pageIter.Err(); err != nil {
return err
}
if pageEndID == 0 { if pageEndID == 0 {
log.Debug("All messages are not expired") log.Debug("All messages are not expired, skip retention", zap.Any("topic", topic))
return nil return nil
} }
log.Debug("ExpiredCleanUp: ", zap.Any("topic", topic), zap.Any("pageEndID", pageEndID), zap.Any("deletedAckedSize", deletedAckedSize)) log.Debug("Expired check by message size: ", zap.Any("topic", topic),
zap.Any("pageEndID", pageEndID), zap.Any("deletedAckedSize", deletedAckedSize), zap.Any("pageCleaned", pageCleaned))
return ri.cleanData(topic, pageEndID)
}
func (ri *retentionInfo) calculateTopicAckedSize(topic string) (int64, error) {
fixedAckedTsKey := constructKey(AckedTsTitle, topic)
pageReadOpts := gorocksdb.NewDefaultReadOptions()
defer pageReadOpts.Destroy()
pageMsgPrefix := constructKey(PageMsgSizeTitle, topic)
// ensure the iterator won't iterate to other topics
pageReadOpts.SetIterateUpperBound([]byte(typeutil.AddOne(pageMsgPrefix)))
pageIter := ri.kv.DB.NewIterator(pageReadOpts)
defer pageIter.Close()
pageIter.Seek([]byte(pageMsgPrefix))
var ackedSize int64
for ; pageIter.Valid(); pageIter.Next() {
key := pageIter.Key()
pageID, err := parsePageID(string(key.Data()))
if key != nil {
key.Free()
}
if err != nil {
return -1, err
}
// check if page is acked
ackedTsKey := fixedAckedTsKey + "/" + strconv.FormatInt(pageID, 10)
ackedTsVal, err := ri.kv.Load(ackedTsKey)
if err != nil {
return -1, err
}
// not acked yet, break
// TODO, Add TTL logic here, mark it as acked if not
if ackedTsVal == "" {
break
}
// Get page size
val := pageIter.Value()
size, err := strconv.ParseInt(string(val.Data()), 10, 64)
if val != nil {
val.Free()
}
if err != nil {
return -1, err
}
ackedSize += size
}
if err := pageIter.Err(); err != nil {
return -1, err
}
return ackedSize, nil
}
func (ri *retentionInfo) cleanData(topic string, pageEndID UniqueID) error {
writeBatch := gorocksdb.NewWriteBatch() writeBatch := gorocksdb.NewWriteBatch()
defer writeBatch.Destroy() defer writeBatch.Destroy()
pageStartIDKey := pageMsgPrefix + "/" + strconv.FormatInt(pageStartID, 10) pageMsgPrefix := constructKey(PageMsgSizeTitle, topic)
fixedAckedTsKey := constructKey(AckedTsTitle, topic)
pageStartIDKey := pageMsgPrefix + "/"
pageEndIDKey := pageMsgPrefix + "/" + strconv.FormatInt(pageEndID+1, 10) pageEndIDKey := pageMsgPrefix + "/" + strconv.FormatInt(pageEndID+1, 10)
if pageStartID == pageEndID { writeBatch.DeleteRange([]byte(pageStartIDKey), []byte(pageEndIDKey))
if pageStartID != 0 {
writeBatch.Delete([]byte(pageStartIDKey))
}
} else if pageStartID < pageEndID {
writeBatch.DeleteRange([]byte(pageStartIDKey), []byte(pageEndIDKey))
}
ackedStartIDKey := fixedAckedTsKey + "/" + strconv.Itoa(int(startID)) pageTsPrefix := constructKey(PageTsTitle, topic)
ackedEndIDKey := fixedAckedTsKey + "/" + strconv.Itoa(int(pageEndID+1)) pageTsStartIDKey := pageTsPrefix + "/"
if startID > pageEndID { pageTsEndIDKey := pageTsPrefix + "/" + strconv.FormatInt(pageEndID+1, 10)
return nil writeBatch.DeleteRange([]byte(pageTsStartIDKey), []byte(pageTsEndIDKey))
}
ackedStartIDKey := fixedAckedTsKey + "/"
ackedEndIDKey := fixedAckedTsKey + "/" + strconv.FormatInt(pageEndID+1, 10)
writeBatch.DeleteRange([]byte(ackedStartIDKey), []byte(ackedEndIDKey)) writeBatch.DeleteRange([]byte(ackedStartIDKey), []byte(ackedEndIDKey))
ll, ok := topicMu.Load(topic) ll, ok := topicMu.Load(topic)
@ -323,26 +324,18 @@ func (ri *retentionInfo) expiredCleanUp(topic string) error {
} }
lock.Lock() lock.Lock()
defer lock.Unlock() defer lock.Unlock()
currentAckedSizeVal, err := ri.kv.Load(ackedSizeKey)
if err != nil {
return err
}
currentAckedSize, err := strconv.ParseInt(currentAckedSizeVal, 10, 64)
if err != nil {
return err
}
newAckedSize := currentAckedSize - deletedAckedSize
writeBatch.Put([]byte(ackedSizeKey), []byte(strconv.FormatInt(newAckedSize, 10)))
err = DeleteMessages(ri.db, topic, startID, pageEndID) err := DeleteMessages(ri.db, topic, 0, pageEndID)
if err != nil { if err != nil {
return err return err
} }
writeOpts := gorocksdb.NewDefaultWriteOptions() writeOpts := gorocksdb.NewDefaultWriteOptions()
defer writeOpts.Destroy() defer writeOpts.Destroy()
ri.kv.DB.Write(writeOpts, writeBatch) err = ri.kv.DB.Write(writeOpts, writeBatch)
if err != nil {
return err
}
return nil return nil
} }
@ -359,14 +352,9 @@ func DeleteMessages(db *gorocksdb.DB, topic string, startID, endID UniqueID) err
log.Debug("RocksMQ: combKey(" + topic + "," + strconv.FormatInt(endID, 10) + ")") log.Debug("RocksMQ: combKey(" + topic + "," + strconv.FormatInt(endID, 10) + ")")
return err return err
} }
writeBatch := gorocksdb.NewWriteBatch() writeBatch := gorocksdb.NewWriteBatch()
defer writeBatch.Destroy() defer writeBatch.Destroy()
if startID == endID { writeBatch.DeleteRange([]byte(startKey), []byte(endKey))
writeBatch.Delete([]byte(startKey))
} else {
writeBatch.DeleteRange([]byte(startKey), []byte(endKey))
}
opts := gorocksdb.NewDefaultWriteOptions() opts := gorocksdb.NewDefaultWriteOptions()
defer opts.Destroy() defer opts.Destroy()
err = db.Write(opts, writeBatch) err = db.Write(opts, writeBatch)
@ -375,14 +363,19 @@ func DeleteMessages(db *gorocksdb.DB, topic string, startID, endID UniqueID) err
} }
log.Debug("Delete message for topic: "+topic, zap.Any("startID", startID), zap.Any("endID", endID)) log.Debug("Delete message for topic: "+topic, zap.Any("startID", startID), zap.Any("endID", endID))
return nil return nil
} }
func msgTimeExpiredCheck(ackedTs int64) bool { func msgTimeExpiredCheck(ackedTs int64) bool {
return ackedTs+atomic.LoadInt64(&RocksmqRetentionTimeInMinutes)*MINUTE < time.Now().Unix() if RocksmqRetentionTimeInSecs < 0 {
return false
}
return ackedTs+atomic.LoadInt64(&RocksmqRetentionTimeInSecs) < time.Now().Unix()
} }
func msgSizeExpiredCheck(deletedAckedSize, ackedSize int64) bool { func msgSizeExpiredCheck(deletedAckedSize, ackedSize int64) bool {
if RocksmqRetentionSizeInMB < 0 {
return false
}
return ackedSize-deletedAckedSize > atomic.LoadInt64(&RocksmqRetentionSizeInMB)*MB return ackedSize-deletedAckedSize > atomic.LoadInt64(&RocksmqRetentionSizeInMB)*MB
} }

390
internal/util/rocksmq/server/rocksmq/rocksmq_retention_test.go
View File

@ -12,7 +12,6 @@
package rocksmq package rocksmq
import ( import (
"math/rand"
"os" "os"
"strconv" "strconv"
"sync/atomic" "sync/atomic"
@ -28,43 +27,30 @@ import (
var retentionPath = "/tmp/rmq_retention/" var retentionPath = "/tmp/rmq_retention/"
func TestMain(m *testing.M) { func TestMain(m *testing.M) {
err := os.MkdirAll(retentionPath, os.ModePerm)
if err != nil {
log.Error("MkdirALl error for path", zap.Any("path", retentionPath))
return
}
atomic.StoreInt64(&TickerTimeInSeconds, 6)
code := m.Run() code := m.Run()
os.Exit(code) os.Exit(code)
} }
func genRandonName() string { // Test write data and wait for retention
len := 6 func TestRmqRetention_Basic(t *testing.T) {
r := rand.New(rand.NewSource(time.Now().UnixNano())) err := os.MkdirAll(retentionPath, os.ModePerm)
bytes := make([]byte, len) if err != nil {
for i := 0; i < len; i++ { log.Error("MkdirAll error for path", zap.Any("path", retentionPath))
b := r.Intn(26) + 65 return
bytes[i] = byte(b)
} }
return string(bytes) defer os.RemoveAll(retentionPath)
}
func TestRmqRetention(t *testing.T) {
atomic.StoreInt64(&RocksmqRetentionSizeInMB, 0) atomic.StoreInt64(&RocksmqRetentionSizeInMB, 0)
atomic.StoreInt64(&RocksmqRetentionTimeInMinutes, 0) atomic.StoreInt64(&RocksmqRetentionTimeInSecs, 0)
atomic.StoreInt64(&RocksmqPageSize, 10) atomic.StoreInt64(&RocksmqPageSize, 10)
atomic.StoreInt64(&TickerTimeInSeconds, 2) atomic.StoreInt64(&TickerTimeInSeconds, 2)
defer atomic.StoreInt64(&TickerTimeInSeconds, 6)
kvPath := retentionPath + kvPathSuffix
defer os.RemoveAll(kvPath)
idAllocator := InitIDAllocator(kvPath)
rocksdbPath := retentionPath + dbPathSuffix rocksdbPath := retentionPath + dbPathSuffix
defer os.RemoveAll(rocksdbPath) defer os.RemoveAll(rocksdbPath)
metaPath := retentionPath + metaPathSuffix metaPath := retentionPath + metaPathSuffix
defer os.RemoveAll(metaPath) defer os.RemoveAll(metaPath)
rmq, err := NewRocksMQ(rocksdbPath, idAllocator) rmq, err := NewRocksMQ(rocksdbPath, nil)
defer rmq.Close()
assert.Nil(t, err) assert.Nil(t, err)
defer rmq.stopRetention() defer rmq.stopRetention()
@ -111,13 +97,147 @@ func TestRmqRetention(t *testing.T) {
newRes, err := rmq.Consume(topicName, groupName, 1) newRes, err := rmq.Consume(topicName, groupName, 1)
assert.Nil(t, err) assert.Nil(t, err)
assert.Equal(t, len(newRes), 0) assert.Equal(t, len(newRes), 0)
//////////////////////////////////////////////////
// test valid value case // test acked size acked ts and other meta are updated as expect
rmq.retentionInfo.topics.Store(topicName, "dummy") msgSizeKey := MessageSizeTitle + topicName
msgSizeVal, err := rmq.kv.Load(msgSizeKey)
assert.NoError(t, err)
assert.Equal(t, msgSizeVal, "0")
pageMsgSizeKey := constructKey(PageMsgSizeTitle, topicName)
keys, values, err := rmq.kv.LoadWithPrefix(pageMsgSizeKey)
assert.NoError(t, err)
assert.Equal(t, len(keys), 0)
assert.Equal(t, len(values), 0)
pageTsSizeKey := constructKey(PageTsTitle, topicName)
keys, values, err = rmq.kv.LoadWithPrefix(pageTsSizeKey)
assert.NoError(t, err)
assert.Equal(t, len(keys), 0)
assert.Equal(t, len(values), 0)
aclTsSizeKey := constructKey(AckedTsTitle, topicName)
keys, values, err = rmq.kv.LoadWithPrefix(aclTsSizeKey)
assert.NoError(t, err)
assert.Equal(t, len(keys), 0)
assert.Equal(t, len(values), 0)
}
// Not acked message should not be purged
func TestRmqRetention_NotConsumed(t *testing.T) {
err := os.MkdirAll(retentionPath, os.ModePerm)
if err != nil {
log.Error("MkdirAll error for path", zap.Any("path", retentionPath))
return
}
defer os.RemoveAll(retentionPath)
atomic.StoreInt64(&RocksmqRetentionSizeInMB, 0)
atomic.StoreInt64(&RocksmqRetentionTimeInSecs, 0)
atomic.StoreInt64(&RocksmqPageSize, 10)
atomic.StoreInt64(&TickerTimeInSeconds, 2)
rocksdbPath := retentionPath + dbPathSuffix
defer os.RemoveAll(rocksdbPath)
metaPath := retentionPath + metaPathSuffix
defer os.RemoveAll(metaPath)
rmq, err := NewRocksMQ(rocksdbPath, nil)
defer rmq.Close()
assert.Nil(t, err)
defer rmq.stopRetention()
topicName := "topic_a"
err = rmq.CreateTopic(topicName)
assert.Nil(t, err)
defer rmq.DestroyTopic(topicName)
msgNum := 100
pMsgs := make([]ProducerMessage, msgNum)
for i := 0; i < msgNum; i++ {
msg := "message_" + strconv.Itoa(i)
pMsg := ProducerMessage{Payload: []byte(msg)}
pMsgs[i] = pMsg
}
ids, err := rmq.Produce(topicName, pMsgs)
assert.Nil(t, err)
assert.Equal(t, len(pMsgs), len(ids))
groupName := "test_group"
_ = rmq.DestroyConsumerGroup(topicName, groupName)
err = rmq.CreateConsumerGroup(topicName, groupName)
consumer := &Consumer{
Topic: topicName,
GroupName: groupName,
}
rmq.RegisterConsumer(consumer)
assert.Nil(t, err)
cMsgs := make([]ConsumerMessage, 0)
for i := 0; i < 5; i++ {
cMsg, err := rmq.Consume(topicName, groupName, 1)
assert.Nil(t, err)
cMsgs = append(cMsgs, cMsg[0])
}
assert.Equal(t, len(cMsgs), 5)
id := cMsgs[0].MsgID
aclTsSizeKey := constructKey(AckedTsTitle, topicName)
keys, values, err := rmq.kv.LoadWithPrefix(aclTsSizeKey)
assert.NoError(t, err)
assert.Equal(t, len(keys), 2)
assert.Equal(t, len(values), 2)
// wait for retention
checkTimeInterval := 2
time.Sleep(time.Duration(checkTimeInterval+1) * time.Second) time.Sleep(time.Duration(checkTimeInterval+1) * time.Second)
// Seek to a previous consumed message, the message should be clean up
err = rmq.Seek(topicName, groupName, cMsgs[1].MsgID)
assert.Nil(t, err)
newRes, err := rmq.Consume(topicName, groupName, 1)
assert.Nil(t, err)
assert.Equal(t, len(newRes), 1)
assert.Equal(t, newRes[0].MsgID, id+4)
// test acked size acked ts and other meta are updated as expect
msgSizeKey := MessageSizeTitle + topicName
msgSizeVal, err := rmq.kv.Load(msgSizeKey)
assert.NoError(t, err)
assert.Equal(t, msgSizeVal, "0")
// should only clean 2 pages
pageMsgSizeKey := constructKey(PageMsgSizeTitle, topicName)
keys, values, err = rmq.kv.LoadWithPrefix(pageMsgSizeKey)
assert.NoError(t, err)
assert.Equal(t, len(keys), 48)
assert.Equal(t, len(values), 48)
pageTsSizeKey := constructKey(PageTsTitle, topicName)
keys, values, err = rmq.kv.LoadWithPrefix(pageTsSizeKey)
assert.NoError(t, err)
assert.Equal(t, len(keys), 48)
assert.Equal(t, len(values), 48)
aclTsSizeKey = constructKey(AckedTsTitle, topicName)
keys, values, err = rmq.kv.LoadWithPrefix(aclTsSizeKey)
assert.NoError(t, err)
assert.Equal(t, len(keys), 0)
assert.Equal(t, len(values), 0)
}
// Test multiple topic
func TestRmqRetention_MultipleTopic(t *testing.T) {
} }
func TestRetentionInfo_InitRetentionInfo(t *testing.T) { func TestRetentionInfo_InitRetentionInfo(t *testing.T) {
err := os.MkdirAll(retentionPath, os.ModePerm)
if err != nil {
log.Error("MkdirALl error for path", zap.Any("path", retentionPath))
return
}
defer os.RemoveAll(retentionPath)
suffix := "init" suffix := "init"
kvPath := retentionPath + kvPathSuffix + suffix kvPath := retentionPath + kvPathSuffix + suffix
defer os.RemoveAll(kvPath) defer os.RemoveAll(kvPath)
@ -133,32 +253,24 @@ func TestRetentionInfo_InitRetentionInfo(t *testing.T) {
assert.Nil(t, err) assert.Nil(t, err)
assert.NotNil(t, rmq) assert.NotNil(t, rmq)
rmq.retentionInfo.kv.DB = nil
_, err = initRetentionInfo(rmq.retentionInfo.kv, rmq.retentionInfo.db)
assert.Error(t, err)
}
func TestRmqRetention_Complex(t *testing.T) {
atomic.StoreInt64(&RocksmqRetentionSizeInMB, 0)
atomic.StoreInt64(&RocksmqRetentionTimeInMinutes, 1)
atomic.StoreInt64(&RocksmqPageSize, 10)
kvPath := retentionPath + "kv_com"
defer os.RemoveAll(kvPath)
idAllocator := InitIDAllocator(kvPath)
rocksdbPath := retentionPath + "db_com"
defer os.RemoveAll(rocksdbPath)
metaPath := retentionPath + "meta_kv_com"
defer os.RemoveAll(metaPath)
rmq, err := NewRocksMQ(rocksdbPath, idAllocator)
assert.Nil(t, err)
defer rmq.stopRetention()
topicName := "topic_a" topicName := "topic_a"
err = rmq.CreateTopic(topicName) err = rmq.CreateTopic(topicName)
assert.Nil(t, err) assert.Nil(t, err)
defer rmq.DestroyTopic(topicName)
rmq.Close()
rmq, err = NewRocksMQ(rocksdbPath, idAllocator)
assert.Nil(t, err)
assert.NotNil(t, rmq)
assert.Equal(t, rmq.isClosed(), false)
// write some data, restart and check.
topicName = "topic_a"
err = rmq.CreateTopic(topicName)
assert.Nil(t, err)
topicName = "topic_b"
err = rmq.CreateTopic(topicName)
assert.Nil(t, err)
msgNum := 100 msgNum := 100
pMsgs := make([]ProducerMessage, msgNum) pMsgs := make([]ProducerMessage, msgNum)
@ -171,42 +283,22 @@ func TestRmqRetention_Complex(t *testing.T) {
assert.Nil(t, err) assert.Nil(t, err)
assert.Equal(t, len(pMsgs), len(ids)) assert.Equal(t, len(pMsgs), len(ids))
groupName := "test_group" rmq.Close()
_ = rmq.DestroyConsumerGroup(topicName, groupName)
err = rmq.CreateConsumerGroup(topicName, groupName)
consumer := &Consumer{
Topic: topicName,
GroupName: groupName,
}
rmq.RegisterConsumer(consumer)
assert.Nil(t, err)
cMsgs := make([]ConsumerMessage, 0)
for i := 0; i < msgNum; i++ {
cMsg, err := rmq.Consume(topicName, groupName, 1)
assert.Nil(t, err)
cMsgs = append(cMsgs, cMsg[0])
}
assert.Equal(t, len(cMsgs), msgNum)
checkTimeInterval := atomic.LoadInt64(&RocksmqRetentionTimeInMinutes) * MINUTE / 10
time.Sleep(time.Duration(checkTimeInterval*2) * time.Second)
// Seek to a previous consumed message, the message should be clean up
log.Debug("cMsg", zap.Any("id", cMsgs[10].MsgID))
err = rmq.Seek(topicName, groupName, cMsgs[10].MsgID)
assert.Nil(t, err)
newRes, err := rmq.Consume(topicName, groupName, 1)
assert.Nil(t, err)
//TODO(yukun)
log.Debug("Consume result", zap.Any("result len", len(newRes)))
// assert.NotEqual(t, newRes[0].MsgID, cMsgs[11].MsgID)
} }
func TestRmqRetention_PageTimeExpire(t *testing.T) { func TestRmqRetention_PageTimeExpire(t *testing.T) {
atomic.StoreInt64(&RocksmqRetentionSizeInMB, 0) err := os.MkdirAll(retentionPath, os.ModePerm)
atomic.StoreInt64(&RocksmqRetentionTimeInMinutes, 0) if err != nil {
log.Error("MkdirALl error for path", zap.Any("path", retentionPath))
return
}
defer os.RemoveAll(retentionPath)
// no retention by size
atomic.StoreInt64(&RocksmqRetentionSizeInMB, -1)
// retention by secs
atomic.StoreInt64(&RocksmqRetentionTimeInSecs, 5)
atomic.StoreInt64(&RocksmqPageSize, 10) atomic.StoreInt64(&RocksmqPageSize, 10)
atomic.StoreInt64(&TickerTimeInSeconds, 1)
kvPath := retentionPath + "kv_com1" kvPath := retentionPath + "kv_com1"
os.RemoveAll(kvPath) os.RemoveAll(kvPath)
idAllocator := InitIDAllocator(kvPath) idAllocator := InitIDAllocator(kvPath)
@ -218,7 +310,7 @@ func TestRmqRetention_PageTimeExpire(t *testing.T) {
rmq, err := NewRocksMQ(rocksdbPath, idAllocator) rmq, err := NewRocksMQ(rocksdbPath, idAllocator)
assert.Nil(t, err) assert.Nil(t, err)
defer rmq.stopRetention() defer rmq.Close()
topicName := "topic_a" topicName := "topic_a"
err = rmq.CreateTopic(topicName) err = rmq.CreateTopic(topicName)
@ -236,6 +328,126 @@ func TestRmqRetention_PageTimeExpire(t *testing.T) {
assert.Nil(t, err) assert.Nil(t, err)
assert.Equal(t, len(pMsgs), len(ids)) assert.Equal(t, len(pMsgs), len(ids))
groupName := "test_group"
_ = rmq.DestroyConsumerGroup(topicName, groupName)
err = rmq.CreateConsumerGroup(topicName, groupName)
assert.NoError(t, err)
consumer := &Consumer{
Topic: topicName,
GroupName: groupName,
}
rmq.RegisterConsumer(consumer)
cMsgs := make([]ConsumerMessage, 0)
for i := 0; i < msgNum; i++ {
cMsg, err := rmq.Consume(topicName, groupName, 1)
assert.Nil(t, err)
cMsgs = append(cMsgs, cMsg[0])
}
assert.Equal(t, len(cMsgs), msgNum)
assert.Equal(t, cMsgs[0].MsgID, ids[0])
time.Sleep(time.Duration(3) * time.Second)
// insert another 100 messages which should not be cleand up
pMsgs2 := make([]ProducerMessage, msgNum)
for i := 0; i < msgNum; i++ {
msg := "message_" + strconv.Itoa(i+100)
pMsg := ProducerMessage{Payload: []byte(msg)}
pMsgs2[i] = pMsg
}
ids2, err := rmq.Produce(topicName, pMsgs2)
assert.Nil(t, err)
assert.Equal(t, len(pMsgs2), len(ids2))
assert.Nil(t, err)
cMsgs = make([]ConsumerMessage, 0)
for i := 0; i < msgNum; i++ {
cMsg, err := rmq.Consume(topicName, groupName, 1)
assert.Nil(t, err)
cMsgs = append(cMsgs, cMsg[0])
}
assert.Equal(t, len(cMsgs), msgNum)
assert.Equal(t, cMsgs[0].MsgID, ids2[0])
time.Sleep(time.Duration(3) * time.Second)
err = rmq.Seek(topicName, groupName, ids[10])
assert.Nil(t, err)
newRes, err := rmq.Consume(topicName, groupName, 1)
assert.Nil(t, err)
assert.Equal(t, len(newRes), 1)
// point to first not consumed messages
assert.Equal(t, newRes[0].MsgID, ids2[0])
// test acked size acked ts and other meta are updated as expect
msgSizeKey := MessageSizeTitle + topicName
msgSizeVal, err := rmq.kv.Load(msgSizeKey)
assert.NoError(t, err)
assert.Equal(t, msgSizeVal, "0")
// 100 page left, each entity is a page
pageMsgSizeKey := constructKey(PageMsgSizeTitle, topicName)
keys, values, err := rmq.kv.LoadWithPrefix(pageMsgSizeKey)
assert.NoError(t, err)
assert.Equal(t, len(keys), 100)
assert.Equal(t, len(values), 100)
pageTsSizeKey := constructKey(PageTsTitle, topicName)
keys, values, err = rmq.kv.LoadWithPrefix(pageTsSizeKey)
assert.NoError(t, err)
assert.Equal(t, len(keys), 100)
assert.Equal(t, len(values), 100)
aclTsSizeKey := constructKey(AckedTsTitle, topicName)
keys, values, err = rmq.kv.LoadWithPrefix(aclTsSizeKey)
assert.NoError(t, err)
assert.Equal(t, len(keys), 100)
assert.Equal(t, len(values), 100)
}
func TestRmqRetention_PageSizeExpire(t *testing.T) {
err := os.MkdirAll(retentionPath, os.ModePerm)
if err != nil {
log.Error("MkdirALl error for path", zap.Any("path", retentionPath))
return
}
defer os.RemoveAll(retentionPath)
atomic.StoreInt64(&RocksmqRetentionSizeInMB, 1)
atomic.StoreInt64(&RocksmqRetentionTimeInSecs, -1)
atomic.StoreInt64(&RocksmqPageSize, 10)
atomic.StoreInt64(&TickerTimeInSeconds, 2)
kvPath := retentionPath + "kv_com2"
os.RemoveAll(kvPath)
idAllocator := InitIDAllocator(kvPath)
rocksdbPath := retentionPath + "db_com2"
os.RemoveAll(rocksdbPath)
metaPath := retentionPath + "meta_kv_com2"
os.RemoveAll(metaPath)
rmq, err := NewRocksMQ(rocksdbPath, idAllocator)
assert.Nil(t, err)
defer rmq.Close()
topicName := "topic_a"
err = rmq.CreateTopic(topicName)
assert.Nil(t, err)
defer rmq.DestroyTopic(topicName)
// need to be larger than 1M
msgNum := 100000
pMsgs := make([]ProducerMessage, msgNum)
for i := 0; i < msgNum; i++ {
msg := "message_" + strconv.Itoa(i)
pMsg := ProducerMessage{Payload: []byte(msg)}
pMsgs[i] = pMsg
}
ids, err := rmq.Produce(topicName, pMsgs)
assert.Nil(t, err)
assert.Equal(t, len(pMsgs), len(ids))
groupName := "test_group" groupName := "test_group"
_ = rmq.DestroyConsumerGroup(topicName, groupName) _ = rmq.DestroyConsumerGroup(topicName, groupName)
err = rmq.CreateConsumerGroup(topicName, groupName) err = rmq.CreateConsumerGroup(topicName, groupName)
@ -255,14 +467,12 @@ func TestRmqRetention_PageTimeExpire(t *testing.T) {
} }
assert.Equal(t, len(cMsgs), msgNum) assert.Equal(t, len(cMsgs), msgNum)
checkTimeInterval := 7 time.Sleep(time.Duration(2) * time.Second)
time.Sleep(time.Duration(checkTimeInterval) * time.Second) err = rmq.Seek(topicName, groupName, ids[0])
// Seek to a previous consumed message, the message should be clean up
log.Debug("cMsg", zap.Any("id", cMsgs[10].MsgID))
err = rmq.Seek(topicName, groupName, cMsgs[len(cMsgs)/2].MsgID)
assert.Nil(t, err) assert.Nil(t, err)
newRes, err := rmq.Consume(topicName, groupName, 1) newRes, err := rmq.Consume(topicName, groupName, 1)
assert.Nil(t, err) assert.Nil(t, err)
assert.Equal(t, len(newRes), 0) assert.Equal(t, len(newRes), 1)
// assert.NotEqual(t, newRes[0].MsgID, cMsgs[11].MsgID) // make sure clean up happens
assert.True(t, newRes[0].MsgID > ids[0])
} }

15
internal/util/typeutil/string_util.go Normal file
View File

@ -0,0 +1,15 @@
package typeutil
// Add one to string, add one on empty string return empty
func AddOne(data string) string {
if len(data) == 0 {
return data
}
var datab = []byte(data)
if datab[len(datab)-1] != 255 {
datab[len(datab)-1]++
} else {
datab = append(datab, byte(0))
}
return string(datab)
}

29
internal/util/typeutil/string_util_test.go Normal file
View File

@ -0,0 +1,29 @@
package typeutil
import (
"testing"
"github.com/stretchr/testify/assert"
)
func TestAddOne(t *testing.T) {
input := ""
output := AddOne(input)
assert.Equal(t, output, "")
input = "a"
output = AddOne(input)
assert.Equal(t, output, "b")
input = "aaa="
output = AddOne(input)
assert.Equal(t, output, "aaa>")
// test the increate case
binary := []byte{1, 20, 255}
input = string(binary)
output = AddOne(input)
assert.Equal(t, len(output), 4)
resultb := []byte(output)
assert.Equal(t, resultb, []byte{1, 20, 255, 0})
}