cmd: split 'surgery freelist' into separate files

Signed-off-by: Benjamin Wang <wachao@vmware.com>

cmd/bbolt/command_surgery.go

@@ -8,7 +8,6 @@ import (
 	"github.com/spf13/cobra"
 	"github.com/spf13/pflag"
 
-	bolt "go.etcd.io/bbolt"
 	"go.etcd.io/bbolt/internal/common"
 	"go.etcd.io/bbolt/internal/guts_cli"
 	"go.etcd.io/bbolt/internal/surgeon"
@@ -311,121 +310,6 @@ func surgeryClearPageElementFunc(srcDBPath string, cfg surgeryClearPageElementsO
 	return nil
 }
 
-func newSurgeryFreelistCommand() *cobra.Command {
-	cmd := &cobra.Command{
-		Use:   "freelist <subcommand>",
-		Short: "freelist related surgery commands",
-	}
-
-	cmd.AddCommand(newSurgeryFreelistAbandonCommand())
-	cmd.AddCommand(newSurgeryFreelistRebuildCommand())
-
-	return cmd
-}
-
-func newSurgeryFreelistAbandonCommand() *cobra.Command {
-	var o surgeryBaseOptions
-	abandonFreelistCmd := &cobra.Command{
-		Use:   "abandon <bbolt-file> [options]",
-		Short: "Abandon the freelist from both meta pages",
-		Args: func(cmd *cobra.Command, args []string) error {
-			if len(args) == 0 {
-				return errors.New("db file path not provided")
-			}
-			if len(args) > 1 {
-				return errors.New("too many arguments")
-			}
-			return nil
-		},
-		RunE: func(cmd *cobra.Command, args []string) error {
-			if err := o.Validate(); err != nil {
-				return err
-			}
-			return surgeryFreelistAbandonFunc(args[0], o)
-		},
-	}
-	o.AddFlags(abandonFreelistCmd.Flags())
-
-	return abandonFreelistCmd
-}
-
-func surgeryFreelistAbandonFunc(srcDBPath string, cfg surgeryBaseOptions) error {
-	if _, err := checkSourceDBPath(srcDBPath); err != nil {
-		return err
-	}
-
-	if err := common.CopyFile(srcDBPath, cfg.outputDBFilePath); err != nil {
-		return fmt.Errorf("[freelist abandon] copy file failed: %w", err)
-	}
-
-	if err := surgeon.ClearFreelist(cfg.outputDBFilePath); err != nil {
-		return fmt.Errorf("abandom-freelist command failed: %w", err)
-	}
-
-	fmt.Fprintf(os.Stdout, "The freelist was abandoned in both meta pages.\nIt may cause some delay on next startup because bbolt needs to scan the whole db to reconstruct the free list.\n")
-	return nil
-}
-
-func newSurgeryFreelistRebuildCommand() *cobra.Command {
-	var o surgeryBaseOptions
-	rebuildFreelistCmd := &cobra.Command{
-		Use:   "rebuild <bbolt-file> [options]",
-		Short: "Rebuild the freelist",
-		Args: func(cmd *cobra.Command, args []string) error {
-			if len(args) == 0 {
-				return errors.New("db file path not provided")
-			}
-			if len(args) > 1 {
-				return errors.New("too many arguments")
-			}
-			return nil
-		},
-		RunE: func(cmd *cobra.Command, args []string) error {
-			if err := o.Validate(); err != nil {
-				return err
-			}
-			return surgeryFreelistRebuildFunc(args[0], o)
-		},
-	}
-	o.AddFlags(rebuildFreelistCmd.Flags())
-
-	return rebuildFreelistCmd
-}
-
-func surgeryFreelistRebuildFunc(srcDBPath string, cfg surgeryBaseOptions) error {
-	// Ensure source file exists.
-	fi, err := checkSourceDBPath(srcDBPath)
-	if err != nil {
-		return err
-	}
-
-	// make sure the freelist isn't present in the file.
-	meta, err := readMetaPage(srcDBPath)
-	if err != nil {
-		return err
-	}
-	if meta.IsFreelistPersisted() {
-		return ErrSurgeryFreelistAlreadyExist
-	}
-
-	if err := common.CopyFile(srcDBPath, cfg.outputDBFilePath); err != nil {
-		return fmt.Errorf("[freelist rebuild] copy file failed: %w", err)
-	}
-
-	// bboltDB automatically reconstruct & sync freelist in write mode.
-	db, err := bolt.Open(cfg.outputDBFilePath, fi.Mode(), &bolt.Options{NoFreelistSync: false})
-	if err != nil {
-		return fmt.Errorf("[freelist rebuild] open db file failed: %w", err)
-	}
-	err = db.Close()
-	if err != nil {
-		return fmt.Errorf("[freelist rebuild] close db file failed: %w", err)
-	}
-
-	fmt.Fprintf(os.Stdout, "The freelist was successfully rebuilt.\n")
-	return nil
-}
-
 func readMetaPage(path string) (*common.Meta, error) {
 	_, activeMetaPageId, err := guts_cli.GetRootPage(path)
 	if err != nil {

cmd/bbolt/command_surgery_freelist.go

@@ -0,0 +1,128 @@
+package main
+
+import (
+	"errors"
+	"fmt"
+	"os"
+
+	"github.com/spf13/cobra"
+
+	bolt "go.etcd.io/bbolt"
+	"go.etcd.io/bbolt/internal/common"
+	"go.etcd.io/bbolt/internal/surgeon"
+)
+
+func newSurgeryFreelistCommand() *cobra.Command {
+	cmd := &cobra.Command{
+		Use:   "freelist <subcommand>",
+		Short: "freelist related surgery commands",
+	}
+
+	cmd.AddCommand(newSurgeryFreelistAbandonCommand())
+	cmd.AddCommand(newSurgeryFreelistRebuildCommand())
+
+	return cmd
+}
+
+func newSurgeryFreelistAbandonCommand() *cobra.Command {
+	var o surgeryBaseOptions
+	abandonFreelistCmd := &cobra.Command{
+		Use:   "abandon <bbolt-file> [options]",
+		Short: "Abandon the freelist from both meta pages",
+		Args: func(cmd *cobra.Command, args []string) error {
+			if len(args) == 0 {
+				return errors.New("db file path not provided")
+			}
+			if len(args) > 1 {
+				return errors.New("too many arguments")
+			}
+			return nil
+		},
+		RunE: func(cmd *cobra.Command, args []string) error {
+			if err := o.Validate(); err != nil {
+				return err
+			}
+			return surgeryFreelistAbandonFunc(args[0], o)
+		},
+	}
+	o.AddFlags(abandonFreelistCmd.Flags())
+
+	return abandonFreelistCmd
+}
+
+func surgeryFreelistAbandonFunc(srcDBPath string, cfg surgeryBaseOptions) error {
+	if _, err := checkSourceDBPath(srcDBPath); err != nil {
+		return err
+	}
+
+	if err := common.CopyFile(srcDBPath, cfg.outputDBFilePath); err != nil {
+		return fmt.Errorf("[freelist abandon] copy file failed: %w", err)
+	}
+
+	if err := surgeon.ClearFreelist(cfg.outputDBFilePath); err != nil {
+		return fmt.Errorf("abandom-freelist command failed: %w", err)
+	}
+
+	fmt.Fprintf(os.Stdout, "The freelist was abandoned in both meta pages.\nIt may cause some delay on next startup because bbolt needs to scan the whole db to reconstruct the free list.\n")
+	return nil
+}
+
+func newSurgeryFreelistRebuildCommand() *cobra.Command {
+	var o surgeryBaseOptions
+	rebuildFreelistCmd := &cobra.Command{
+		Use:   "rebuild <bbolt-file> [options]",
+		Short: "Rebuild the freelist",
+		Args: func(cmd *cobra.Command, args []string) error {
+			if len(args) == 0 {
+				return errors.New("db file path not provided")
+			}
+			if len(args) > 1 {
+				return errors.New("too many arguments")
+			}
+			return nil
+		},
+		RunE: func(cmd *cobra.Command, args []string) error {
+			if err := o.Validate(); err != nil {
+				return err
+			}
+			return surgeryFreelistRebuildFunc(args[0], o)
+		},
+	}
+	o.AddFlags(rebuildFreelistCmd.Flags())
+
+	return rebuildFreelistCmd
+}
+
+func surgeryFreelistRebuildFunc(srcDBPath string, cfg surgeryBaseOptions) error {
+	// Ensure source file exists.
+	fi, err := checkSourceDBPath(srcDBPath)
+	if err != nil {
+		return err
+	}
+
+	// make sure the freelist isn't present in the file.
+	meta, err := readMetaPage(srcDBPath)
+	if err != nil {
+		return err
+	}
+	if meta.IsFreelistPersisted() {
+		return ErrSurgeryFreelistAlreadyExist
+	}
+
+	if err := common.CopyFile(srcDBPath, cfg.outputDBFilePath); err != nil {
+		return fmt.Errorf("[freelist rebuild] copy file failed: %w", err)
+	}
+
+	// bboltDB automatically reconstruct & sync freelist in write mode.
+	db, err := bolt.Open(cfg.outputDBFilePath, fi.Mode(), &bolt.Options{NoFreelistSync: false})
+	if err != nil {
+		return fmt.Errorf("[freelist rebuild] open db file failed: %w", err)
+	}
+	err = db.Close()
+	if err != nil {
+		return fmt.Errorf("[freelist rebuild] close db file failed: %w", err)
+	}
+
+	fmt.Fprintf(os.Stdout, "The freelist was successfully rebuilt.\n")
+	return nil
+}

cmd/bbolt/command_surgery_freelist_test.go

@@ -0,0 +1,103 @@
+package main_test
+
+import (
+	"path/filepath"
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+
+	bolt "go.etcd.io/bbolt"
+	main "go.etcd.io/bbolt/cmd/bbolt"
+	"go.etcd.io/bbolt/internal/btesting"
+	"go.etcd.io/bbolt/internal/common"
+)
+
+func TestSurgery_Freelist_Abandon(t *testing.T) {
+	pageSize := 4096
+	db := btesting.MustCreateDBWithOption(t, &bolt.Options{PageSize: pageSize})
+	srcPath := db.Path()
+
+	defer requireDBNoChange(t, dbData(t, srcPath), srcPath)
+
+	rootCmd := main.NewRootCommand()
+	output := filepath.Join(t.TempDir(), "db")
+	rootCmd.SetArgs([]string{
+		"surgery", "freelist", "abandon", srcPath,
+		"--output", output,
+	})
+	err := rootCmd.Execute()
+	require.NoError(t, err)
+
+	meta0 := loadMetaPage(t, output, 0)
+	assert.Equal(t, common.PgidNoFreelist, meta0.Freelist())
+	meta1 := loadMetaPage(t, output, 1)
+	assert.Equal(t, common.PgidNoFreelist, meta1.Freelist())
+}
+
+func TestSurgery_Freelist_Rebuild(t *testing.T) {
+	testCases := []struct {
+		name          string
+		hasFreelist   bool
+		expectedError error
+	}{
+		{
+			name:          "normal operation",
+			hasFreelist:   false,
+			expectedError: nil,
+		},
+		{
+			name:          "already has freelist",
+			hasFreelist:   true,
+			expectedError: main.ErrSurgeryFreelistAlreadyExist,
+		},
+	}
+
+	for _, tc := range testCases {
+		tc := tc
+		t.Run(tc.name, func(t *testing.T) {
+			pageSize := 4096
+			db := btesting.MustCreateDBWithOption(t, &bolt.Options{
+				PageSize:       pageSize,
+				NoFreelistSync: !tc.hasFreelist,
+			})
+			srcPath := db.Path()
+
+			err := db.Update(func(tx *bolt.Tx) error {
+				// do nothing
+				return nil
+			})
+			require.NoError(t, err)
+
+			defer requireDBNoChange(t, dbData(t, srcPath), srcPath)
+
+			// Verify the freelist isn't synced in the beginning
+			meta := readMetaPage(t, srcPath)
+			if tc.hasFreelist {
+				if meta.Freelist() <= 1 || meta.Freelist() >= meta.Pgid() {
+					t.Fatalf("freelist (%d) isn't in the valid range (1, %d)", meta.Freelist(), meta.Pgid())
+				}
+			} else {
+				require.Equal(t, common.PgidNoFreelist, meta.Freelist())
+			}
+
+			// Execute `surgery freelist rebuild` command
+			rootCmd := main.NewRootCommand()
+			output := filepath.Join(t.TempDir(), "db")
+			rootCmd.SetArgs([]string{
+				"surgery", "freelist", "rebuild", srcPath,
+				"--output", output,
+			})
+			err = rootCmd.Execute()
+			require.Equal(t, tc.expectedError, err)
+
+			if tc.expectedError == nil {
+				// Verify the freelist has already been rebuilt.
+				meta = readMetaPage(t, output)
+				if meta.Freelist() <= 1 || meta.Freelist() >= meta.Pgid() {
+					t.Fatalf("freelist (%d) isn't in the valid range (1, %d)", meta.Freelist(), meta.Pgid())
+				}
+			}
+		})
+	}
+}

cmd/bbolt/command_surgery_test.go

@@ -551,130 +551,6 @@ func testSurgeryClearPageElementsWithOverflow(t *testing.T, startIdx, endIdx int
 	compareDataAfterClearingElement(t, srcPath, output, pageId, false, startIdx, endIdx)
 }
 
-func TestSurgery_Freelist_Abandon(t *testing.T) {
-	pageSize := 4096
-	db := btesting.MustCreateDBWithOption(t, &bolt.Options{PageSize: pageSize})
-	srcPath := db.Path()
-
-	defer requireDBNoChange(t, dbData(t, srcPath), srcPath)
-
-	rootCmd := main.NewRootCommand()
-	output := filepath.Join(t.TempDir(), "db")
-	rootCmd.SetArgs([]string{
-		"surgery", "freelist", "abandon", srcPath,
-		"--output", output,
-	})
-	err := rootCmd.Execute()
-	require.NoError(t, err)
-
-	meta0 := loadMetaPage(t, output, 0)
-	assert.Equal(t, common.PgidNoFreelist, meta0.Freelist())
-	meta1 := loadMetaPage(t, output, 1)
-	assert.Equal(t, common.PgidNoFreelist, meta1.Freelist())
-}
-
-func loadMetaPage(t *testing.T, dbPath string, pageID uint64) *common.Meta {
-	_, buf, err := guts_cli.ReadPage(dbPath, 0)
-	require.NoError(t, err)
-	return common.LoadPageMeta(buf)
-}
-
-func TestSurgery_Freelist_Rebuild(t *testing.T) {
-	testCases := []struct {
-		name          string
-		hasFreelist   bool
-		expectedError error
-	}{
-		{
-			name:          "normal operation",
-			hasFreelist:   false,
-			expectedError: nil,
-		},
-		{
-			name:          "already has freelist",
-			hasFreelist:   true,
-			expectedError: main.ErrSurgeryFreelistAlreadyExist,
-		},
-	}
-
-	for _, tc := range testCases {
-		tc := tc
-		t.Run(tc.name, func(t *testing.T) {
-			pageSize := 4096
-			db := btesting.MustCreateDBWithOption(t, &bolt.Options{
-				PageSize:       pageSize,
-				NoFreelistSync: !tc.hasFreelist,
-			})
-			srcPath := db.Path()
-
-			err := db.Update(func(tx *bolt.Tx) error {
-				// do nothing
-				return nil
-			})
-			require.NoError(t, err)
-
-			defer requireDBNoChange(t, dbData(t, srcPath), srcPath)
-
-			// Verify the freelist isn't synced in the beginning
-			meta := readMetaPage(t, srcPath)
-			if tc.hasFreelist {
-				if meta.Freelist() <= 1 || meta.Freelist() >= meta.Pgid() {
-					t.Fatalf("freelist (%d) isn't in the valid range (1, %d)", meta.Freelist(), meta.Pgid())
-				}
-			} else {
-				require.Equal(t, common.PgidNoFreelist, meta.Freelist())
-			}
-
-			// Execute `surgery freelist rebuild` command
-			rootCmd := main.NewRootCommand()
-			output := filepath.Join(t.TempDir(), "db")
-			rootCmd.SetArgs([]string{
-				"surgery", "freelist", "rebuild", srcPath,
-				"--output", output,
-			})
-			err = rootCmd.Execute()
-			require.Equal(t, tc.expectedError, err)
-
-			if tc.expectedError == nil {
-				// Verify the freelist has already been rebuilt.
-				meta = readMetaPage(t, output)
-				if meta.Freelist() <= 1 || meta.Freelist() >= meta.Pgid() {
-					t.Fatalf("freelist (%d) isn't in the valid range (1, %d)", meta.Freelist(), meta.Pgid())
-				}
-			}
-		})
-	}
-}
-
-func readMetaPage(t *testing.T, path string) *common.Meta {
-	_, activeMetaPageId, err := guts_cli.GetRootPage(path)
-	require.NoError(t, err)
-	_, buf, err := guts_cli.ReadPage(path, uint64(activeMetaPageId))
-	require.NoError(t, err)
-	return common.LoadPageMeta(buf)
-}
-
-func readPage(t *testing.T, path string, pageId int, pageSize int) []byte {
-	dbFile, err := os.Open(path)
-	require.NoError(t, err)
-	defer dbFile.Close()
-
-	fi, err := dbFile.Stat()
-	require.NoError(t, err)
-	require.GreaterOrEqual(t, fi.Size(), int64((pageId+1)*pageSize))
-
-	buf := make([]byte, pageSize)
-	byteRead, err := dbFile.ReadAt(buf, int64(pageId*pageSize))
-	require.NoError(t, err)
-	require.Equal(t, pageSize, byteRead)
-
-	return buf
-}
-
-func pageDataWithoutPageId(buf []byte) []byte {
-	return buf[8:]
-}
-
 func TestSurgeryRequiredFlags(t *testing.T) {
 	errMsgFmt := `required flag(s) "%s" not set`
 	testCases := []struct {

cmd/bbolt/utils_test.go

@@ -0,0 +1,46 @@
+package main_test
+
+import (
+	"os"
+	"testing"
+
+	"github.com/stretchr/testify/require"
+
+	"go.etcd.io/bbolt/internal/common"
+	"go.etcd.io/bbolt/internal/guts_cli"
+)
+
+func loadMetaPage(t *testing.T, dbPath string, pageID uint64) *common.Meta {
+	_, buf, err := guts_cli.ReadPage(dbPath, 0)
+	require.NoError(t, err)
+	return common.LoadPageMeta(buf)
+}
+
+func readMetaPage(t *testing.T, path string) *common.Meta {
+	_, activeMetaPageId, err := guts_cli.GetRootPage(path)
+	require.NoError(t, err)
+	_, buf, err := guts_cli.ReadPage(path, uint64(activeMetaPageId))
+	require.NoError(t, err)
+	return common.LoadPageMeta(buf)
+}
+
+func readPage(t *testing.T, path string, pageId int, pageSize int) []byte {
+	dbFile, err := os.Open(path)
+	require.NoError(t, err)
+	defer dbFile.Close()
+
+	fi, err := dbFile.Stat()
+	require.NoError(t, err)
+	require.GreaterOrEqual(t, fi.Size(), int64((pageId+1)*pageSize))
+
+	buf := make([]byte, pageSize)
+	byteRead, err := dbFile.ReadAt(buf, int64(pageId*pageSize))
+	require.NoError(t, err)
+	require.Equal(t, pageSize, byteRead)
+
+	return buf
+}
+
+func pageDataWithoutPageId(buf []byte) []byte {
+	return buf[8:]
+}