Add storage-analysis command (#183)

Signed-off-by: Congqi Xia <[email protected]>

models/fs_stat.go

@@ -0,0 +1,6 @@
+package models
+
+// FsStat model for file system statistics
+type FsStat struct {
+	Size int64
+}

states/app_state.go

@@ -108,7 +108,7 @@ func (app *ApplicationState) IsEnding() bool {
 }
 
 func (app *ApplicationState) ConnectMinioCommand(ctx context.Context, p *storage.ConnectMinioParam) error {
-	state, err := storage.ConnectMinio(ctx, p)
+	state, err := storage.ConnectMinio(ctx, p, app.core)
 	if err != nil {
 		return err
 	}

states/segment_anaylsis.go

@@ -0,0 +1,68 @@
+package states
+
+import (
+	"context"
+	"fmt"
+
+	"github.com/cockroachdb/errors"
+	"github.com/milvus-io/birdwatcher/framework"
+	"github.com/milvus-io/birdwatcher/models"
+	"github.com/milvus-io/birdwatcher/states/etcd/common"
+	etcdversion "github.com/milvus-io/birdwatcher/states/etcd/version"
+	"github.com/milvus-io/birdwatcher/states/storage"
+)
+
+type StorageAnalysisParam struct {
+	framework.ParamBase `use:"storage-analysis" desc:"segment storage analysis" require:"etcd,minio"`
+	CollectionID        int64 `name:"collection" default:"0" desc:"collection id to analysis"`
+	Detail              bool  `name:"detail" default:"false" desc:"print detailed binlog size info"`
+}
+
+func (app *ApplicationState) StorageAnalysisCommand(ctx context.Context, p *StorageAnalysisParam) error {
+	state, ok := app.states[etcdTag]
+	if !ok {
+		return errors.New("Etcd instance not connected")
+	}
+	etcd, ok := state.(*InstanceState)
+	if !ok {
+		return errors.New("Etcd instance not connected")
+	}
+	state, ok = app.states[minioTag]
+	if !ok {
+		return errors.New("Minio instance not connected")
+	}
+	minio, ok := state.(*storage.MinioState)
+	if !ok {
+		return errors.New("Minio instance not connected")
+	}
+
+	segments, err := common.ListSegmentsVersion(ctx, etcd.client, etcd.basePath, etcdversion.GetVersion(), func(s *models.Segment) bool {
+		return p.CollectionID == 0 || s.CollectionID == p.CollectionID
+	})
+
+	if err != nil {
+		return err
+	}
+
+	for _, segment := range segments {
+		fmt.Printf("segment %d\n", segment.ID)
+		for _, fieldBinlog := range segment.GetBinlogs() {
+			fmt.Println("fieldID:", fieldBinlog.FieldID)
+			var size int64
+			for _, binlog := range fieldBinlog.Binlogs {
+				info, err := minio.Stat(ctx, binlog.LogPath)
+				if err != nil {
+					fmt.Println("failed to stats", err.Error())
+					continue
+				}
+				if p.Detail {
+					fmt.Printf("Binlog %s size %s\n", binlog.LogPath, hrSize(info.Size))
+				}
+				size += info.Size
+			}
+			fmt.Printf("Total binlog size: %s\n", hrSize(size))
+		}
+	}
+
+	return nil
+}

states/storage/minio.go

@@ -7,7 +7,6 @@ import (
 	"github.com/milvus-io/birdwatcher/framework"
 	"github.com/minio/minio-go/v7"
 	"github.com/minio/minio-go/v7/pkg/credentials"
-	"github.com/spf13/cobra"
 )
 
 type MinioState struct {
@@ -19,11 +18,10 @@ type MinioState struct {
 }
 
 func (s *MinioState) SetupCommands() {
-	cmd := &cobra.Command{}
+	cmd := s.GetCmd()
 
 	s.MergeFunctionCommands(cmd, s)
-	s.CmdState.RootCmd = cmd
-	s.SetupFn = s.SetupCommands
+	s.UpdateState(cmd, s, s.SetupCommands)
 }
 
 type ConnectMinioParam struct {
@@ -37,7 +35,7 @@ type ConnectMinioParam struct {
 	UseSSL              bool   `name:"ssl" default:"" desc:"use SSL"`
 }
 
-func ConnectMinio(ctx context.Context, p *ConnectMinioParam) (*MinioState, error) {
+func ConnectMinio(ctx context.Context, p *ConnectMinioParam, parent *framework.CmdState) (*MinioState, error) {
 	var cred *credentials.Credentials
 	if p.UseIAM {
 		cred = credentials.NewIAM(p.IAMEndpoint)
@@ -66,6 +64,6 @@ func ConnectMinio(ctx context.Context, p *ConnectMinioParam) (*MinioState, error
 	return &MinioState{
 		client:   minioClient,
 		bucket:   p.Bucket,
-		CmdState: framework.NewCmdState("Minio(Addr)"),
+		CmdState: parent.Spawn("Minio(Addr)"),
 	}, nil
 }

states/storage/ops.go

@@ -0,0 +1,20 @@
+package storage
+
+import (
+	"context"
+
+	"github.com/milvus-io/birdwatcher/models"
+	"github.com/minio/minio-go/v7"
+)
+
+func (s *MinioState) Stat(ctx context.Context, path string) (*models.FsStat, error) {
+	info, err := s.client.StatObject(ctx, s.bucket, path, minio.StatObjectOptions{})
+	if err != nil {
+		return nil, err
+	}
+
+	result := &models.FsStat{
+		Size: info.Size,
+	}
+	return result, nil
+}