Merge pull request #539 from ahrtr/surgery_meta_20230717

cmd: add meta page related surgery commands

cmd/bbolt/command_surgery.go

@@ -28,6 +28,7 @@ func newSurgeryCobraCommand() *cobra.Command {
 	surgeryCmd.AddCommand(newSurgeryClearPageCommand())
 	surgeryCmd.AddCommand(newSurgeryClearPageElementsCommand())
 	surgeryCmd.AddCommand(newSurgeryFreelistCommand())
+	surgeryCmd.AddCommand(newSurgeryMetaCommand())
 
 	return surgeryCmd
 }
@@ -311,15 +312,23 @@ func surgeryClearPageElementFunc(srcDBPath string, cfg surgeryClearPageElementsO
 }
 
 func readMetaPage(path string) (*common.Meta, error) {
-	_, activeMetaPageId, err := guts_cli.GetRootPage(path)
+	pageSize, _, err := guts_cli.ReadPageAndHWMSize(path)
 	if err != nil {
-		return nil, fmt.Errorf("read root page failed: %w", err)
+		return nil, fmt.Errorf("read Page size failed: %w", err)
 	}
-	_, buf, err := guts_cli.ReadPage(path, uint64(activeMetaPageId))
-	if err != nil {
-		return nil, fmt.Errorf("read active mage page failed: %w", err)
+
+	m := make([]*common.Meta, 2)
+	for i := 0; i < 2; i++ {
+		m[i], _, err = ReadMetaPageAt(path, uint32(i), uint32(pageSize))
+		if err != nil {
+			return nil, fmt.Errorf("read meta page %d failed: %w", i, err)
+		}
 	}
-	return common.LoadPageMeta(buf), nil
+
+	if m[0].Txid() > m[1].Txid() {
+		return m[0], nil
+	}
+	return m[1], nil
 }
 
 func checkSourceDBPath(srcPath string) (os.FileInfo, error) {

cmd/bbolt/command_surgery_meta.go

@@ -0,0 +1,292 @@
+package main
+
+import (
+	"errors"
+	"fmt"
+	"io"
+	"os"
+	"strconv"
+	"strings"
+
+	"github.com/spf13/cobra"
+	"github.com/spf13/pflag"
+
+	"go.etcd.io/bbolt/internal/common"
+)
+
+const (
+	metaFieldPageSize = "pageSize"
+	metaFieldRoot     = "root"
+	metaFieldFreelist = "freelist"
+	metaFieldPgid     = "pgid"
+)
+
+func newSurgeryMetaCommand() *cobra.Command {
+	cmd := &cobra.Command{
+		Use:   "meta <subcommand>",
+		Short: "meta page related surgery commands",
+	}
+
+	cmd.AddCommand(newSurgeryMetaValidateCommand())
+	cmd.AddCommand(newSurgeryMetaUpdateCommand())
+
+	return cmd
+}
+
+func newSurgeryMetaValidateCommand() *cobra.Command {
+	metaValidateCmd := &cobra.Command{
+		Use:   "validate <bbolt-file> [options]",
+		Short: "Validate 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 {
+			return surgeryMetaValidateFunc(args[0])
+		},
+	}
+	return metaValidateCmd
+}
+
+func surgeryMetaValidateFunc(srcDBPath string) error {
+	if _, err := checkSourceDBPath(srcDBPath); err != nil {
+		return err
+	}
+
+	var pageSize uint32
+
+	for i := 0; i <= 1; i++ {
+		m, _, err := ReadMetaPageAt(srcDBPath, uint32(i), pageSize)
+		if err != nil {
+			return fmt.Errorf("read meta page %d failed: %w", i, err)
+		}
+		if mValidateErr := m.Validate(); mValidateErr != nil {
+			fmt.Fprintf(os.Stdout, "WARNING: The meta page %d isn't valid: %v!\n", i, mValidateErr)
+		} else {
+			fmt.Fprintf(os.Stdout, "The meta page %d is valid!\n", i)
+		}
+
+		pageSize = m.PageSize()
+	}
+
+	return nil
+}
+
+type surgeryMetaUpdateOptions struct {
+	surgeryBaseOptions
+	fields     []string
+	metaPageId uint32
+}
+
+var allowedMetaUpdateFields = map[string]struct{}{
+	metaFieldPageSize: {},
+	metaFieldRoot:     {},
+	metaFieldFreelist: {},
+	metaFieldPgid:     {},
+}
+
+// AddFlags sets the flags for `meta update` command.
+// Example: --fields root:16,freelist:8 --fields pgid:128 --fields txid:1234
+// Result: []string{"root:16", "freelist:8", "pgid:128", "txid:1234"}
+func (o *surgeryMetaUpdateOptions) AddFlags(fs *pflag.FlagSet) {
+	o.surgeryBaseOptions.AddFlags(fs)
+	fs.StringSliceVarP(&o.fields, "fields", "", o.fields, "comma separated list of fields (supported fields: pageSize, root, freelist, pgid and txid) to be updated, and each item is a colon-separated key-value pair")
+	fs.Uint32VarP(&o.metaPageId, "meta-page", "", o.metaPageId, "the meta page ID to operate on, valid values are 0 and 1")
+}
+
+func (o *surgeryMetaUpdateOptions) Validate() error {
+	if err := o.surgeryBaseOptions.Validate(); err != nil {
+		return err
+	}
+
+	if o.metaPageId > 1 {
+		return fmt.Errorf("invalid meta page id: %d", o.metaPageId)
+	}
+
+	for _, field := range o.fields {
+		kv := strings.Split(field, ":")
+		if len(kv) != 2 {
+			return fmt.Errorf("invalid key-value pair: %s", field)
+		}
+
+		if _, ok := allowedMetaUpdateFields[kv[0]]; !ok {
+			return fmt.Errorf("field %q isn't allowed to be updated", kv[0])
+		}
+
+		if _, err := strconv.ParseUint(kv[1], 10, 64); err != nil {
+			return fmt.Errorf("invalid value %q for field %q", kv[1], kv[0])
+		}
+	}
+
+	return nil
+}
+
+func newSurgeryMetaUpdateCommand() *cobra.Command {
+	var o surgeryMetaUpdateOptions
+	metaUpdateCmd := &cobra.Command{
+		Use:   "update <bbolt-file> [options]",
+		Short: "Update fields in 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 surgeryMetaUpdateFunc(args[0], o)
+		},
+	}
+	o.AddFlags(metaUpdateCmd.Flags())
+	return metaUpdateCmd
+}
+
+func surgeryMetaUpdateFunc(srcDBPath string, cfg surgeryMetaUpdateOptions) error {
+	if _, err := checkSourceDBPath(srcDBPath); err != nil {
+		return err
+	}
+
+	if err := common.CopyFile(srcDBPath, cfg.outputDBFilePath); err != nil {
+		return fmt.Errorf("[meta update] copy file failed: %w", err)
+	}
+
+	// read the page size from the first meta page if we want to edit the second meta page.
+	var pageSize uint32
+	if cfg.metaPageId == 1 {
+		m0, _, err := ReadMetaPageAt(cfg.outputDBFilePath, 0, pageSize)
+		if err != nil {
+			return fmt.Errorf("read the first meta page failed: %w", err)
+		}
+		pageSize = m0.PageSize()
+	}
+
+	// update the specified meta page
+	m, buf, err := ReadMetaPageAt(cfg.outputDBFilePath, cfg.metaPageId, pageSize)
+	if err != nil {
+		return fmt.Errorf("read meta page %d failed: %w", cfg.metaPageId, err)
+	}
+	mChanged := updateMetaField(m, parseFields(cfg.fields))
+	if mChanged {
+		if err := writeMetaPageAt(cfg.outputDBFilePath, buf, cfg.metaPageId, pageSize); err != nil {
+			return fmt.Errorf("[meta update] write meta page %d failed: %w", cfg.metaPageId, err)
+		}
+	}
+
+	if cfg.metaPageId == 1 && pageSize != m.PageSize() {
+		fmt.Fprintf(os.Stdout, "WARNING: The page size (%d) in the first meta page doesn't match the second meta page (%d)\n", pageSize, m.PageSize())
+	}
+
+	// Display results
+	if !mChanged {
+		fmt.Fprintln(os.Stdout, "Nothing changed!")
+	}
+
+	if mChanged {
+		fmt.Fprintf(os.Stdout, "The meta page %d has been updated!\n", cfg.metaPageId)
+	}
+
+	return nil
+}
+
+func parseFields(fields []string) map[string]uint64 {
+	fieldsMap := make(map[string]uint64)
+	for _, field := range fields {
+		kv := strings.SplitN(field, ":", 2)
+		val, _ := strconv.ParseUint(kv[1], 10, 64)
+		fieldsMap[kv[0]] = val
+	}
+	return fieldsMap
+}
+
+func updateMetaField(m *common.Meta, fields map[string]uint64) bool {
+	changed := false
+	for key, val := range fields {
+		switch key {
+		case metaFieldPageSize:
+			m.SetPageSize(uint32(val))
+		case metaFieldRoot:
+			m.SetRootBucket(common.NewInBucket(common.Pgid(val), 0))
+		case metaFieldFreelist:
+			m.SetFreelist(common.Pgid(val))
+		case metaFieldPgid:
+			m.SetPgid(common.Pgid(val))
+		}
+
+		changed = true
+	}
+
+	if m.Magic() != common.Magic {
+		m.SetMagic(common.Magic)
+		changed = true
+	}
+	if m.Version() != common.Version {
+		m.SetVersion(common.Version)
+		changed = true
+	}
+	if m.Flags() != common.MetaPageFlag {
+		m.SetFlags(common.MetaPageFlag)
+		changed = true
+	}
+
+	newChecksum := m.Sum64()
+	if m.Checksum() != newChecksum {
+		m.SetChecksum(newChecksum)
+		changed = true
+	}
+
+	return changed
+}
+
+func ReadMetaPageAt(dbPath string, metaPageId uint32, pageSize uint32) (*common.Meta, []byte, error) {
+	if metaPageId > 1 {
+		return nil, nil, fmt.Errorf("invalid metaPageId: %d", metaPageId)
+	}
+
+	f, err := os.OpenFile(dbPath, os.O_RDONLY, 0444)
+	if err != nil {
+		return nil, nil, err
+	}
+	defer f.Close()
+
+	// The meta page is just 64 bytes, and definitely less than 1024 bytes,
+	// so it's fine to only read 1024 bytes. Note we don't care about the
+	// pageSize when reading the first meta page, because we always read the
+	// file starting from offset 0. Actually the passed pageSize is 0 when
+	// reading the first meta page in the `surgery meta update` command.
+	buf := make([]byte, 1024)
+	n, err := f.ReadAt(buf, int64(metaPageId*pageSize))
+	if n == len(buf) && (err == nil || err == io.EOF) {
+		return common.LoadPageMeta(buf), buf, nil
+	}
+
+	return nil, nil, err
+}
+
+func writeMetaPageAt(dbPath string, buf []byte, metaPageId uint32, pageSize uint32) error {
+	if metaPageId > 1 {
+		return fmt.Errorf("invalid metaPageId: %d", metaPageId)
+	}
+
+	f, err := os.OpenFile(dbPath, os.O_RDWR, 0666)
+	if err != nil {
+		return err
+	}
+	defer f.Close()
+
+	n, err := f.WriteAt(buf, int64(metaPageId*pageSize))
+	if n == len(buf) && (err == nil || err == io.EOF) {
+		return nil
+	}
+
+	return err
+}

cmd/bbolt/command_surgery_meta_test.go

@@ -0,0 +1,126 @@
+package main_test
+
+import (
+	"fmt"
+	"path/filepath"
+	"strings"
+	"testing"
+
+	"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_Meta_Validate(t *testing.T) {
+	pageSize := 4096
+	db := btesting.MustCreateDBWithOption(t, &bolt.Options{PageSize: pageSize})
+	srcPath := db.Path()
+
+	defer requireDBNoChange(t, dbData(t, db.Path()), db.Path())
+
+	// validate the meta pages
+	rootCmd := main.NewRootCommand()
+	rootCmd.SetArgs([]string{
+		"surgery", "meta", "validate", srcPath,
+	})
+	err := rootCmd.Execute()
+	require.NoError(t, err)
+
+	// TODD: add one more case that the validation may fail. We need to
+	// make the command output configurable, so that test cases can set
+	// a customized io.Writer.
+}
+
+func TestSurgery_Meta_Update(t *testing.T) {
+	testCases := []struct {
+		name     string
+		root     common.Pgid
+		freelist common.Pgid
+		pgid     common.Pgid
+	}{
+		{
+			name: "root changed",
+			root: 50,
+		},
+		{
+			name:     "freelist changed",
+			freelist: 40,
+		},
+		{
+			name: "pgid changed",
+			pgid: 600,
+		},
+		{
+			name:     "both root and freelist changed",
+			root:     45,
+			freelist: 46,
+		},
+		{
+			name:     "both pgid and freelist changed",
+			pgid:     256,
+			freelist: 47,
+		},
+		{
+			name:     "all fields changed",
+			root:     43,
+			freelist: 62,
+			pgid:     256,
+		},
+	}
+
+	for _, tc := range testCases {
+		for i := 0; i <= 1; i++ {
+			tc := tc
+			metaPageId := uint32(i)
+
+			t.Run(tc.name, func(t *testing.T) {
+				pageSize := 4096
+				db := btesting.MustCreateDBWithOption(t, &bolt.Options{PageSize: pageSize})
+				srcPath := db.Path()
+
+				defer requireDBNoChange(t, dbData(t, db.Path()), db.Path())
+
+				var fields []string
+				if tc.root != 0 {
+					fields = append(fields, fmt.Sprintf("root:%d", tc.root))
+				}
+				if tc.freelist != 0 {
+					fields = append(fields, fmt.Sprintf("freelist:%d", tc.freelist))
+				}
+				if tc.pgid != 0 {
+					fields = append(fields, fmt.Sprintf("pgid:%d", tc.pgid))
+				}
+
+				rootCmd := main.NewRootCommand()
+				output := filepath.Join(t.TempDir(), "db")
+				rootCmd.SetArgs([]string{
+					"surgery", "meta", "update", srcPath,
+					"--output", output,
+					"--meta-page", fmt.Sprintf("%d", metaPageId),
+					"--fields", strings.Join(fields, ","),
+				})
+				err := rootCmd.Execute()
+				require.NoError(t, err)
+
+				m, _, err := main.ReadMetaPageAt(output, metaPageId, 4096)
+				require.NoError(t, err)
+
+				require.Equal(t, common.Magic, m.Magic())
+				require.Equal(t, common.Version, m.Version())
+
+				if tc.root != 0 {
+					require.Equal(t, tc.root, m.RootBucket().RootPage())
+				}
+				if tc.freelist != 0 {
+					require.Equal(t, tc.freelist, m.Freelist())
+				}
+				if tc.pgid != 0 {
+					require.Equal(t, tc.pgid, m.Pgid())
+				}
+			})
+		}
+	}
+}

internal/common/meta.go

@@ -72,6 +72,10 @@ func (m *Meta) SetMagic(v uint32) {
 	m.magic = v
 }
 
+func (m *Meta) Version() uint32 {
+	return m.version
+}
+
 func (m *Meta) SetVersion(v uint32) {
 	m.version = v
 }
@@ -136,6 +140,10 @@ func (m *Meta) DecTxid() {
 	m.txid -= 1
 }
 
+func (m *Meta) Checksum() uint64 {
+	return m.checksum
+}
+
 func (m *Meta) SetChecksum(v uint64) {
 	m.checksum = v
 }

internal/common/types.go

@@ -10,7 +10,7 @@ import (
 const MaxMmapStep = 1 << 30 // 1GB
 
 // Version represents the data file format version.
-const Version = 2
+const Version uint32 = 2
 
 // Magic represents a marker value to indicate that a file is a Bolt DB.
 const Magic uint32 = 0xED0CDAED