Remove data directory for test (milvus-io#14363)

Signed-off-by: xiaofan-luan <xiaofan.luan@zilliz.com>

internal/kv/rocksdb/rocksdb_kv_test.go

@@ -17,6 +17,7 @@
 package rocksdbkv_test
 
 import (
+	"os"
 	"strconv"
 	"sync"
 	"testing"
@@ -31,7 +32,7 @@ func TestRocksdbKV(t *testing.T) {
 	if err != nil {
 		panic(err)
 	}
-
+	defer os.RemoveAll(name)
 	defer rocksdbKV.Close()
 	// Need to call RemoveWithPrefix
 	defer rocksdbKV.RemoveWithPrefix("")
@@ -92,7 +93,7 @@ func TestRocksdbKV_Prefix(t *testing.T) {
 	if err != nil {
 		panic(err)
 	}
-
+	defer os.RemoveAll(name)
 	defer rocksdbKV.Close()
 	// Need to call RemoveWithPrefix
 	defer rocksdbKV.RemoveWithPrefix("")
@@ -150,6 +151,7 @@ func TestRocksdbKV_Goroutines(t *testing.T) {
 	name := "/tmp/rocksdb"
 	rocksdbkv, err := rocksdbkv.NewRocksdbKV(name)
 	assert.Nil(t, err)
+	defer os.RemoveAll(name)
 	defer rocksdbkv.Close()
 	defer rocksdbkv.RemoveWithPrefix("")
 
@@ -175,6 +177,7 @@ func TestRocksdbKV_DummyDB(t *testing.T) {
 	name := "/tmp/rocksdb_dummy"
 	rocksdbkv, err := rocksdbkv.NewRocksdbKV(name)
 	assert.Nil(t, err)
+	defer os.RemoveAll(name)
 	defer rocksdbkv.Close()
 	defer rocksdbkv.RemoveWithPrefix("")
 
@@ -209,6 +212,7 @@ func TestRocksdbKV_CornerCase(t *testing.T) {
 	name := "/tmp/rocksdb_corner"
 	rocksdbkv, err := rocksdbkv.NewRocksdbKV(name)
 	assert.Nil(t, err)
+	defer os.RemoveAll(name)
 	defer rocksdbkv.Close()
 	defer rocksdbkv.RemoveWithPrefix("")
 	_, err = rocksdbkv.Load("")

internal/storage/print_binlog_test.go

@@ -73,6 +73,7 @@ func TestPrintBinlogFilesInt64(t *testing.T) {
 	w.Close()
 
 	fd, err := ioutil.TempFile("", "binlog_int64.db")
+	defer os.RemoveAll(fd.Name())
 	assert.Nil(t, err)
 	num, err := fd.Write(buf)
 	assert.Nil(t, err)
@@ -322,6 +323,9 @@ func TestPrintBinlogFiles(t *testing.T) {
 	binlogFiles = append(binlogFiles, "test")
 
 	PrintBinlogFiles(binlogFiles)
+	for _, file := range binlogFiles {
+		_ = os.RemoveAll(file)
+	}
 }
 
 func TestPrintDDFiles(t *testing.T) {
@@ -433,6 +437,10 @@ func TestPrintDDFiles(t *testing.T) {
 	assert.Equal(t, resultRequests, ddRequests)
 
 	PrintBinlogFiles(binlogFiles)
+
+	for _, file := range binlogFiles {
+		_ = os.RemoveAll(file)
+	}
 }
 
 func TestPrintIndexFile(t *testing.T) {

internal/util/mqclient/rmq_consumer.go

@@ -13,6 +13,7 @@ package mqclient
 
 import (
 	"sync"
+	"sync/atomic"
 
 	"github.com/milvus-io/milvus/internal/util/rocksmq/client/rocksmq"
 )
@@ -23,7 +24,7 @@ type RmqConsumer struct {
 	msgChannel chan Message
 	closeCh    chan struct{}
 	once       sync.Once
-	skip       bool
+	skip       int32
 }
 
 // Subscription returns the subscription name of this consumer
@@ -44,10 +45,11 @@ func (rc *RmqConsumer) Chan() <-chan Message {
 							close(rc.msgChannel)
 							return
 						}
-						if !rc.skip {
+						skip := atomic.LoadInt32(&rc.skip)
+						if skip != 1 {
 							rc.msgChannel <- &rmqMessage{msg: msg}
 						} else {
-							rc.skip = false
+							atomic.StoreInt32(&rc.skip, 0)
 						}
 					case <-rc.closeCh:
 						close(rc.msgChannel)
@@ -64,7 +66,9 @@ func (rc *RmqConsumer) Chan() <-chan Message {
 func (rc *RmqConsumer) Seek(id MessageID, inclusive bool) error {
 	msgID := id.(*rmqID).messageID
 	// skip the first message when consume
-	rc.skip = !inclusive
+	if !inclusive {
+		atomic.StoreInt32(&rc.skip, 1)
+	}
 	return rc.c.Seek(msgID)
 }
 

internal/util/rocksmq/client/rocksmq/client_impl_test.go

@@ -21,12 +21,6 @@ import (
 
 var rmqPath = "/tmp/rocksmq_client"
 
-func TestMain(m *testing.M) {
-	os.MkdirAll(rmqPath, os.ModePerm)
-	defer os.RemoveAll(rmqPath)
-	os.Exit(m.Run())
-}
-
 func TestClient(t *testing.T) {
 	client, err := NewClient(ClientOptions{})
 	assert.NotNil(t, client)
@@ -52,8 +46,9 @@ func TestClient_CreateProducer(t *testing.T) {
 	assert.Nil(t, producer)
 
 	/////////////////////////////////////////////////
-	rmqPath := rmqPath + "/test_client1"
-	rmq := newRocksMQ(rmqPath)
+	os.MkdirAll(rmqPath, os.ModePerm)
+	rmqPathTest := rmqPath + "/test_client1"
+	rmq := newRocksMQ(t, rmqPathTest)
 	defer removePath(rmqPath)
 	client1, err := NewClient(ClientOptions{
 		Server: rmq,
@@ -91,8 +86,9 @@ func TestClient_Subscribe(t *testing.T) {
 	assert.Nil(t, consumer)
 
 	/////////////////////////////////////////////////
-	rmqPath := rmqPath + "/test_client2"
-	rmq := newRocksMQ(rmqPath)
+	os.MkdirAll(rmqPath, os.ModePerm)
+	rmqPathTest := rmqPath + "/test_client2"
+	rmq := newRocksMQ(t, rmqPathTest)
 	defer removePath(rmqPath)
 	client1, err := NewClient(ClientOptions{
 		Server: rmq,
@@ -131,8 +127,9 @@ func TestClient_Subscribe(t *testing.T) {
 }
 
 func TestClient_SeekLatest(t *testing.T) {
-	rmqPath := rmqPath + "/seekLatest"
-	rmq := newRocksMQ(rmqPath)
+	os.MkdirAll(rmqPath, os.ModePerm)
+	rmqPathTest := rmqPath + "/seekLatest"
+	rmq := newRocksMQ(t, rmqPathTest)
 	defer removePath(rmqPath)
 	client, err := NewClient(ClientOptions{
 		Server: rmq,
@@ -201,8 +198,9 @@ func TestClient_SeekLatest(t *testing.T) {
 }
 
 func TestClient_consume(t *testing.T) {
-	rmqPath := rmqPath + "/test_client3"
-	rmq := newRocksMQ(rmqPath)
+	os.MkdirAll(rmqPath, os.ModePerm)
+	rmqPathTest := rmqPath + "/test_client3"
+	rmq := newRocksMQ(t, rmqPathTest)
 	defer removePath(rmqPath)
 	client, err := NewClient(ClientOptions{
 		Server: rmq,

internal/util/rocksmq/client/rocksmq/consumer_impl_test.go

@@ -12,6 +12,7 @@
 package rocksmq
 
 import (
+	"os"
 	"testing"
 
 	"github.com/stretchr/testify/assert"
@@ -47,8 +48,9 @@ func TestConsumer_newConsumer(t *testing.T) {
 	assert.Equal(t, InvalidConfiguration, err.(*Error).Result())
 
 	/////////////////////////////////////////////////
-	rmqPath := rmqPath + "/test_consumer1"
-	rmq := newRocksMQ(rmqPath)
+	os.MkdirAll(rmqPath, os.ModePerm)
+	rmqPathTest := rmqPath + "/test_consumer1"
+	rmq := newRocksMQ(t, rmqPathTest)
 	defer removePath(rmqPath)
 	client, err := newClient(ClientOptions{
 		Server: rmq,
@@ -117,8 +119,9 @@ func TestConsumer_Subscription(t *testing.T) {
 }
 
 func TestConsumer_Seek(t *testing.T) {
-	rmqPath := rmqPath + "/test_consumer2"
-	rmq := newRocksMQ(rmqPath)
+	os.MkdirAll(rmqPath, os.ModePerm)
+	rmqPathTest := rmqPath + "/test_consumer2"
+	rmq := newRocksMQ(t, rmqPathTest)
 	defer removePath(rmqPath)
 	client, err := newClient(ClientOptions{
 		Server: rmq,

internal/util/rocksmq/client/rocksmq/reader_impl_test.go

@@ -13,6 +13,7 @@ package rocksmq
 
 import (
 	"context"
+	"os"
 	"strconv"
 	"testing"
 
@@ -40,8 +41,9 @@ func Test_NewReader(t *testing.T) {
 }
 
 func TestReader_Next(t *testing.T) {
-	rmqPath := rmqPath + "/test_reader"
-	rmq := newRocksMQ(rmqPath)
+	os.MkdirAll(rmqPath, os.ModePerm)
+	rmqPathTest := rmqPath + "/test_reader"
+	rmq := newRocksMQ(t, rmqPathTest)
 	defer removePath(rmqPath)
 	client, err := newClient(ClientOptions{
 		Server: rmq,

internal/util/rocksmq/client/rocksmq/test_helper.go

@@ -14,11 +14,13 @@ package rocksmq
 import (
 	"fmt"
 	"os"
+	"testing"
 	"time"
 
 	"github.com/milvus-io/milvus/internal/log"
 	"github.com/milvus-io/milvus/internal/util/rocksmq/server/rocksmq"
 	server "github.com/milvus-io/milvus/internal/util/rocksmq/server/rocksmq"
+	"github.com/stretchr/testify/assert"
 
 	"go.uber.org/zap"
 )
@@ -43,26 +45,23 @@ func newMockClient() *client {
 	return client
 }
 
-func newRocksMQ(rmqPath string) server.RocksMQ {
-	rocksdbPath := rmqPath + "_db"
-	rmq, _ := rocksmq.NewRocksMQ(rocksdbPath, nil)
+func newRocksMQ(t *testing.T, rmqPath string) server.RocksMQ {
+	rocksdbPath := rmqPath
+	rmq, err := rocksmq.NewRocksMQ(rocksdbPath, nil)
+	assert.NoError(t, err)
 	return rmq
 }
 
 func removePath(rmqPath string) {
-	kvPath := rmqPath + "_kv"
-	err := os.RemoveAll(kvPath)
+	// remove path rocksmq created
+	rocksdbPath := rmqPath
+	err := os.RemoveAll(rocksdbPath)
 	if err != nil {
-		log.Error("Failed to call os.removeAll.", zap.Any("path", kvPath))
-	}
-	rocksdbPath := rmqPath + "_db"
-	err = os.RemoveAll(rocksdbPath)
-	if err != nil {
-		log.Error("Failed to call os.removeAll.", zap.Any("path", kvPath))
+		log.Error("Failed to call os.removeAll.", zap.Any("path", rocksdbPath))
 	}
 	metaPath := rmqPath + "_meta_kv"
 	err = os.RemoveAll(metaPath)
 	if err != nil {
-		log.Error("Failed to call os.removeAll.", zap.Any("path", kvPath))
+		log.Error("Failed to call os.removeAll.", zap.Any("path", metaPath))
 	}
 }

internal/util/rocksmq/server/rocksmq/rocksmq_impl_test.go

@@ -34,7 +34,6 @@ import (
 var Params paramtable.BaseTable
 var rmqPath = "/tmp/rocksmq"
 var kvPathSuffix = "_kv"
-var dbPathSuffix = "_db"
 var metaPathSuffix = "_meta"
 
 func InitIDAllocator(kvPath string) *allocator.GlobalIDAllocator {
@@ -70,7 +69,7 @@ func TestRocksmq_RegisterConsumer(t *testing.T) {
 	defer os.RemoveAll(kvPath)
 	idAllocator := InitIDAllocator(kvPath)
 
-	rocksdbPath := rmqPath + dbPathSuffix + suffix
+	rocksdbPath := rmqPath + suffix
 	defer os.RemoveAll(rocksdbPath + kvSuffix)
 	defer os.RemoveAll(rocksdbPath)
 
@@ -135,7 +134,7 @@ func TestRocksmq_Basic(t *testing.T) {
 	defer os.RemoveAll(kvPath)
 	idAllocator := InitIDAllocator(kvPath)
 
-	rocksdbPath := rmqPath + dbPathSuffix + suffix
+	rocksdbPath := rmqPath + suffix
 	defer os.RemoveAll(rocksdbPath + kvSuffix)
 	defer os.RemoveAll(rocksdbPath)
 	rmq, err := NewRocksMQ(rocksdbPath, idAllocator)
@@ -190,7 +189,7 @@ func TestRocksmq_Dummy(t *testing.T) {
 	defer os.RemoveAll(kvPath)
 	idAllocator := InitIDAllocator(kvPath)
 
-	rocksdbPath := rmqPath + dbPathSuffix + suffix
+	rocksdbPath := rmqPath + suffix
 	defer os.RemoveAll(rocksdbPath + kvSuffix)
 	defer os.RemoveAll(rocksdbPath)
 
@@ -260,7 +259,7 @@ func TestRocksmq_Seek(t *testing.T) {
 	defer os.RemoveAll(kvPath)
 	idAllocator := InitIDAllocator(kvPath)
 
-	rocksdbPath := rmqPath + dbPathSuffix + suffix
+	rocksdbPath := rmqPath + suffix
 	defer os.RemoveAll(rocksdbPath + kvSuffix)
 	defer os.RemoveAll(rocksdbPath)
 

internal/util/rocksmq/server/rocksmq/rocksmq_retention_test.go

@@ -44,7 +44,7 @@ func TestRmqRetention_Basic(t *testing.T) {
 	atomic.StoreInt64(&RocksmqPageSize, 10)
 	atomic.StoreInt64(&TickerTimeInSeconds, 2)
 
-	rocksdbPath := retentionPath + dbPathSuffix
+	rocksdbPath := retentionPath
 	defer os.RemoveAll(rocksdbPath)
 	metaPath := retentionPath + metaPathSuffix
 	defer os.RemoveAll(metaPath)
@@ -136,7 +136,7 @@ func TestRmqRetention_NotConsumed(t *testing.T) {
 	atomic.StoreInt64(&RocksmqPageSize, 10)
 	atomic.StoreInt64(&TickerTimeInSeconds, 2)
 
-	rocksdbPath := retentionPath + dbPathSuffix
+	rocksdbPath := retentionPath
 	defer os.RemoveAll(rocksdbPath)
 	metaPath := retentionPath + metaPathSuffix
 	defer os.RemoveAll(metaPath)
@@ -400,7 +400,7 @@ func TestRetentionInfo_InitRetentionInfo(t *testing.T) {
 	defer os.RemoveAll(kvPath)
 	idAllocator := InitIDAllocator(kvPath)
 
-	rocksdbPath := retentionPath + dbPathSuffix + suffix
+	rocksdbPath := retentionPath + suffix
 	defer os.RemoveAll(rocksdbPath)
 	metaPath := retentionPath + metaPathSuffix + suffix