seaweedfs/weed/shell/command_fs_meta_load.go

package shell

import (
	"compress/gzip"
	"context"
	"flag"
	"fmt"
	"io"
	"os"
	"strings"
	"sync"
	"time"

	"github.com/seaweedfs/seaweedfs/weed/pb/filer_pb"
	"github.com/seaweedfs/seaweedfs/weed/util"
	"google.golang.org/protobuf/proto"
)

func init() {
	Commands = append(Commands, &commandFsMetaLoad{})
}

type commandFsMetaLoad struct {
	dirPrefix *string
}

func (c *commandFsMetaLoad) Name() string {
	return "fs.meta.load"
}

func (c *commandFsMetaLoad) Help() string {
	return `load saved filer meta data to restore the directory and file structure

	fs.meta.load <filer_host>-<port>-<time>.meta
	fs.meta.load -v=false <filer_host>-<port>-<time>.meta // skip printing out the verbose output
 	fs.meta.load -concurrency=1 <filer_host>-<port>-<time>.meta // number of parallel meta load to filer
	fs.meta.load -dirPrefix=/buckets/important <filer_host>.meta // load any dirs with prefix "important"

`
}

func (c *commandFsMetaLoad) HasTag(CommandTag) bool {
	return false
}

func (c *commandFsMetaLoad) Do(args []string, commandEnv *CommandEnv, writer io.Writer) (err error) {

	if len(args) == 0 {
		fmt.Fprintf(writer, "missing a metadata file\n")
		return nil
	}

	fileName := util.ResolvePath(args[len(args)-1])

	metaLoadCommand := flag.NewFlagSet(c.Name(), flag.ContinueOnError)
	c.dirPrefix = metaLoadCommand.String("dirPrefix", "", "load entries only with directories matching prefix")
	concurrency := metaLoadCommand.Int("concurrency", 1, "number of parallel meta load to filer")
	verbose := metaLoadCommand.Bool("v", true, "verbose mode")
	if err = metaLoadCommand.Parse(args[0 : len(args)-1]); err != nil {
		return nil
	}

	var dst io.Reader

	f, err := os.OpenFile(fileName, os.O_RDONLY, 0644)
	if err != nil {
		return fmt.Errorf("failed to open file %s: %v", fileName, err)
	}
	defer f.Close()

	dst = f

	if strings.HasSuffix(fileName, ".gz") || strings.HasSuffix(fileName, ".gzip") {
		var gr *gzip.Reader
		gr, err = gzip.NewReader(dst)
		if err != nil {
			return err
		}
		defer func() {
			err1 := gr.Close()
			if err == nil {
				err = err1
			}
		}()

		dst = gr
	}

	var dirCount, fileCount uint64

	err = commandEnv.WithFilerClient(false, func(client filer_pb.SeaweedFilerClient) error {
		var loadErr error
		dirCount, fileCount, loadErr = c.loadEntries(dst, client, *c.dirPrefix, *concurrency, *verbose, writer)
		return loadErr
	})

	if err == nil {
		fmt.Fprintf(writer, "\ntotal %d directories, %d files", dirCount, fileCount)
		fmt.Fprintf(writer, "\n%s is loaded.\n", fileName)
	}

	return err
}

func (c *commandFsMetaLoad) loadEntries(dst io.Reader, client filer_pb.SeaweedFilerClient, dirPrefix string, concurrency int, verbose bool, writer io.Writer) (dirCount, fileCount uint64, err error) {

	lastLogTime := time.Now()
	sizeBuf := make([]byte, 4)
	waitChan := make(chan struct{}, concurrency)
	defer close(waitChan)
	var wg sync.WaitGroup
	var mu sync.Mutex
	var loadErr error
	getErr := func() error {
		mu.Lock()
		defer mu.Unlock()
		return loadErr
	}
	setErr := func(e error) {
		mu.Lock()
		if loadErr == nil {
			loadErr = e
		}
		mu.Unlock()
	}

	for {
		if _, err := io.ReadFull(dst, sizeBuf); err != nil {
			wg.Wait()
			if loadErr := getErr(); loadErr != nil {
				return dirCount, fileCount, loadErr
			}
			if err == io.EOF {
				return dirCount, fileCount, nil
			}
			return dirCount, fileCount, err
		}

		size := util.BytesToUint32(sizeBuf)

		data := make([]byte, int(size))

		if _, err := io.ReadFull(dst, data); err != nil {
			wg.Wait()
			if loadErr := getErr(); loadErr != nil {
				return dirCount, fileCount, loadErr
			}
			return dirCount, fileCount, err
		}

		fullEntry := &filer_pb.FullEntry{}
		if err = proto.Unmarshal(data, fullEntry); err != nil {
			wg.Wait()
			if loadErr := getErr(); loadErr != nil {
				return dirCount, fileCount, loadErr
			}
			return dirCount, fileCount, err
		}

		// check collection name pattern
		entryFullName := string(util.FullPath(fullEntry.Dir).Child(fullEntry.Entry.Name))
		if dirPrefix != "" {
			if !strings.HasPrefix(fullEntry.Dir, dirPrefix) {
				if verbose {
					fmt.Fprintf(writer, "not match dir prefix %s\n", entryFullName)
				}
				continue
			}
		}

		if verbose || lastLogTime.Add(time.Second).Before(time.Now()) {
			if !verbose {
				lastLogTime = time.Now()
			}
			fmt.Fprintf(writer, "load %s\n", entryFullName)
		}

		fullEntry.Entry.Name = strings.ReplaceAll(fullEntry.Entry.Name, "/", "x")
		if fullEntry.Entry.IsDirectory {
			wg.Wait()
			if loadErr := getErr(); loadErr != nil {
				return dirCount, fileCount, loadErr
			}
			if errEntry := filer_pb.CreateEntry(context.Background(), client, &filer_pb.CreateEntryRequest{
				Directory: fullEntry.Dir,
				Entry:     fullEntry.Entry,
			}); errEntry != nil {
				return dirCount, fileCount, errEntry
			}
			dirCount++
		} else {
			wg.Add(1)
			waitChan <- struct{}{}
			go func(entry *filer_pb.FullEntry) {
				if errEntry := filer_pb.CreateEntry(context.Background(), client, &filer_pb.CreateEntryRequest{
					Directory: entry.Dir,
					Entry:     entry.Entry,
				}); errEntry != nil {
					setErr(errEntry)
				}
				defer wg.Done()
				<-waitChan
			}(fullEntry)
			if loadErr := getErr(); loadErr != nil {
				wg.Wait()
				return dirCount, fileCount, loadErr
			}
			fileCount++
		}
	}
}