diff --git a/weed/shell/command_fs_meta_load.go b/weed/shell/command_fs_meta_load.go index 93426ce516..a8ff28a383 100644 --- a/weed/shell/command_fs_meta_load.go +++ b/weed/shell/command_fs_meta_load.go @@ -11,10 +11,9 @@ import ( "sync" "time" - "google.golang.org/protobuf/proto" - "github.com/seaweedfs/seaweedfs/weed/pb/filer_pb" "github.com/seaweedfs/seaweedfs/weed/util" + "google.golang.org/protobuf/proto" ) func init() { @@ -88,83 +87,11 @@ func (c *commandFsMetaLoad) Do(args []string, commandEnv *CommandEnv, writer io. } var dirCount, fileCount uint64 - lastLogTime := time.Now() err = commandEnv.WithFilerClient(false, func(client filer_pb.SeaweedFilerClient) error { - - sizeBuf := make([]byte, 4) - waitChan := make(chan struct{}, *concurrency) - defer close(waitChan) - var wg sync.WaitGroup - - for { - if _, err := io.ReadFull(dst, sizeBuf); err != nil { - if err == io.EOF { - return nil - } - return err - } - - size := util.BytesToUint32(sizeBuf) - - data := make([]byte, int(size)) - - if _, err := io.ReadFull(dst, data); err != nil { - return err - } - - fullEntry := &filer_pb.FullEntry{} - if err = proto.Unmarshal(data, fullEntry); err != nil { - return err - } - - // check collection name pattern - entryFullName := string(util.FullPath(fullEntry.Dir).Child(fullEntry.Entry.Name)) - if *c.dirPrefix != "" { - if !strings.HasPrefix(fullEntry.Dir, *c.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 errEntry := filer_pb.CreateEntry(context.Background(), client, &filer_pb.CreateEntryRequest{ - Directory: fullEntry.Dir, - Entry: fullEntry.Entry, - }); errEntry != nil { - return 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 { - err = errEntry - } - defer wg.Done() - <-waitChan - }(fullEntry) - if err != nil { - return err - } - fileCount++ - } - } + var loadErr error + dirCount, fileCount, loadErr = c.loadEntries(dst, client, *c.dirPrefix, *concurrency, *verbose, writer) + return loadErr }) if err == nil { @@ -174,3 +101,111 @@ func (c *commandFsMetaLoad) Do(args []string, commandEnv *CommandEnv, writer io. 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++ + } + } +} diff --git a/weed/shell/command_fs_meta_load_test.go b/weed/shell/command_fs_meta_load_test.go new file mode 100644 index 0000000000..a3c3f36e66 --- /dev/null +++ b/weed/shell/command_fs_meta_load_test.go @@ -0,0 +1,247 @@ +package shell + +import ( + "bytes" + "context" + "errors" + "io" + "sync" + "testing" + "time" + + "google.golang.org/grpc" + "google.golang.org/protobuf/proto" + + "github.com/seaweedfs/seaweedfs/weed/pb/filer_pb" + "github.com/seaweedfs/seaweedfs/weed/util" +) + +type stubMetaLoadClient struct { + filer_pb.SeaweedFilerClient + createEntry func(ctx context.Context, req *filer_pb.CreateEntryRequest) (*filer_pb.CreateEntryResponse, error) +} + +func (s *stubMetaLoadClient) CreateEntry(ctx context.Context, req *filer_pb.CreateEntryRequest, opts ...grpc.CallOption) (*filer_pb.CreateEntryResponse, error) { + return s.createEntry(ctx, req) +} + +func writeMetaEntry(t *testing.T, buf *bytes.Buffer, dir, name string, isDir bool) { + t.Helper() + data, err := proto.Marshal(&filer_pb.FullEntry{ + Dir: dir, + Entry: &filer_pb.Entry{Name: name, IsDirectory: isDir}, + }) + if err != nil { + t.Fatalf("marshal entry %s/%s: %v", dir, name, err) + } + sizeBuf := make([]byte, 4) + util.Uint32toBytes(sizeBuf, uint32(len(data))) + buf.Write(sizeBuf) + buf.Write(data) +} + +func runLoadEntries(c *commandFsMetaLoad, buf *bytes.Buffer, client filer_pb.SeaweedFilerClient) (dirCount, fileCount uint64, err error) { + return c.loadEntries(buf, client, "", 4, false, io.Discard) +} + +func TestFsMetaLoadCreatesAllEntries(t *testing.T) { + var buf bytes.Buffer + writeMetaEntry(t, &buf, "/dir", "sub", true) + writeMetaEntry(t, &buf, "/dir", "file1", false) + writeMetaEntry(t, &buf, "/dir", "file2", false) + + var mu sync.Mutex + var created []string + client := &stubMetaLoadClient{createEntry: func(ctx context.Context, req *filer_pb.CreateEntryRequest) (*filer_pb.CreateEntryResponse, error) { + mu.Lock() + created = append(created, req.Directory+"/"+req.Entry.Name) + mu.Unlock() + return &filer_pb.CreateEntryResponse{}, nil + }} + + c := &commandFsMetaLoad{} + dirCount, fileCount, err := runLoadEntries(c, &buf, client) + if err != nil { + t.Fatalf("loadEntries: %v", err) + } + if dirCount != 1 || fileCount != 2 { + t.Fatalf("got %d dirs %d files, want 1 dir 2 files", dirCount, fileCount) + } + if len(created) != 3 { + t.Fatalf("got %d CreateEntry calls, want 3", len(created)) + } +} + +func TestFsMetaLoadWaitsForInflightCreateAtEOF(t *testing.T) { + var buf bytes.Buffer + writeMetaEntry(t, &buf, "/dir", "file1", false) + + started := make(chan struct{}) + release := make(chan struct{}) + client := &stubMetaLoadClient{createEntry: func(ctx context.Context, req *filer_pb.CreateEntryRequest) (*filer_pb.CreateEntryResponse, error) { + close(started) + <-release + return &filer_pb.CreateEntryResponse{}, nil + }} + + c := &commandFsMetaLoad{} + done := make(chan error, 1) + go func() { + _, _, err := runLoadEntries(c, &buf, client) + done <- err + }() + + select { + case <-started: + case <-time.After(5 * time.Second): + t.Fatal("CreateEntry was not called") + } + + select { + case err := <-done: + t.Fatalf("loadEntries returned %v while CreateEntry was still in flight", err) + case <-time.After(100 * time.Millisecond): + } + + close(release) + select { + case err := <-done: + if err != nil { + t.Fatalf("loadEntries: %v", err) + } + case <-time.After(5 * time.Second): + t.Fatal("loadEntries did not return after CreateEntry completed") + } +} + +func TestFsMetaLoadWaitsForInflightCreateOnCorruptData(t *testing.T) { + var buf bytes.Buffer + writeMetaEntry(t, &buf, "/dir", "file1", false) + writeCorruptTail(&buf) + + started := make(chan struct{}) + release := make(chan struct{}) + client := &stubMetaLoadClient{createEntry: func(ctx context.Context, req *filer_pb.CreateEntryRequest) (*filer_pb.CreateEntryResponse, error) { + close(started) + <-release + return &filer_pb.CreateEntryResponse{}, nil + }} + + c := &commandFsMetaLoad{} + done := make(chan error, 1) + go func() { + _, _, err := runLoadEntries(c, &buf, client) + done <- err + }() + + select { + case <-started: + case <-time.After(5 * time.Second): + t.Fatal("CreateEntry was not called") + } + + select { + case err := <-done: + t.Fatalf("loadEntries returned %v while CreateEntry was still in flight", err) + case <-time.After(100 * time.Millisecond): + } + + close(release) + select { + case err := <-done: + if err == nil { + t.Fatal("loadEntries returned nil on truncated input") + } + case <-time.After(5 * time.Second): + t.Fatal("loadEntries did not return after CreateEntry completed") + } +} + +func TestFsMetaLoadReturnsWorkerError(t *testing.T) { + var buf bytes.Buffer + writeMetaEntry(t, &buf, "/dir", "file1", false) + + workerErr := errors.New("create entry rejected") + client := &stubMetaLoadClient{createEntry: func(ctx context.Context, req *filer_pb.CreateEntryRequest) (*filer_pb.CreateEntryResponse, error) { + return nil, workerErr + }} + + c := &commandFsMetaLoad{} + _, _, err := runLoadEntries(c, &buf, client) + if !errors.Is(err, workerErr) { + t.Fatalf("got %v, want wrapped %v", err, workerErr) + } +} + +func TestFsMetaLoadPrefersWorkerErrorOverReadError(t *testing.T) { + var buf bytes.Buffer + writeMetaEntry(t, &buf, "/dir", "file1", false) + writeCorruptTail(&buf) + + workerErr := errors.New("create entry rejected") + started := make(chan struct{}) + release := make(chan struct{}) + client := &stubMetaLoadClient{createEntry: func(ctx context.Context, req *filer_pb.CreateEntryRequest) (*filer_pb.CreateEntryResponse, error) { + close(started) + <-release + return nil, workerErr + }} + + c := &commandFsMetaLoad{} + done := make(chan error, 1) + go func() { + _, _, err := runLoadEntries(c, &buf, client) + done <- err + }() + + select { + case <-started: + case <-time.After(5 * time.Second): + t.Fatal("CreateEntry was not called") + } + close(release) + + select { + case err := <-done: + if !errors.Is(err, workerErr) { + t.Fatalf("got %v, want wrapped %v", err, workerErr) + } + case <-time.After(5 * time.Second): + t.Fatal("loadEntries did not return") + } +} + +func TestFsMetaLoadReturnsErrorOnCorruptData(t *testing.T) { + client := &stubMetaLoadClient{createEntry: func(ctx context.Context, req *filer_pb.CreateEntryRequest) (*filer_pb.CreateEntryResponse, error) { + return &filer_pb.CreateEntryResponse{}, nil + }} + c := &commandFsMetaLoad{} + + t.Run("truncated entry data", func(t *testing.T) { + var buf bytes.Buffer + writeMetaEntry(t, &buf, "/dir", "file1", false) + writeCorruptTail(&buf) + if _, _, err := runLoadEntries(c, &buf, client); err == nil { + t.Fatal("loadEntries returned nil on truncated input") + } + }) + + t.Run("malformed proto", func(t *testing.T) { + var buf bytes.Buffer + writeMetaEntry(t, &buf, "/dir", "file1", false) + sizeBuf := make([]byte, 4) + util.Uint32toBytes(sizeBuf, 4) + buf.Write(sizeBuf) + buf.Write([]byte{0xff, 0xff, 0xff, 0xff}) + if _, _, err := runLoadEntries(c, &buf, client); err == nil { + t.Fatal("loadEntries returned nil on malformed proto") + } + }) +} + +func writeCorruptTail(buf *bytes.Buffer) { + sizeBuf := make([]byte, 4) + util.Uint32toBytes(sizeBuf, 100) + buf.Write(sizeBuf) + buf.Write([]byte{1, 2, 3}) +}