feat: read,write data with buffered

This commit is contained in:
loveuer 2024-03-26 21:05:37 +08:00
parent 486ffba682
commit 3ecc0eff5a

View File

@ -3,9 +3,7 @@ package cmd
import ( import (
"context" "context"
"encoding/json" "encoding/json"
"errors"
"fmt" "fmt"
"io"
"net/url" "net/url"
"os" "os"
@ -87,37 +85,84 @@ func run(cmd *cobra.Command, args []string) error {
func executeData(ctx context.Context, input, output interfaces.DumpIO) error { func executeData(ctx context.Context, input, output interfaces.DumpIO) error {
var ( var (
err error err error
lines []*interfaces.ESSource ch = make(chan []*interfaces.ESSource, 1)
succeed int errCh = make(chan error)
total int
) )
for { // write goroutine
go func(c context.Context) {
var (
lines []*interfaces.ESSource
)
if lines, err = input.ReadData(ctx, f_limit); err != nil { defer func() {
if errors.Is(err, io.EOF) { close(ch)
return nil }()
for {
select {
case <-c.Done():
return
default:
if lines, err = input.ReadData(c, f_limit); err != nil {
errCh <- err
return
}
if len(lines) == 0 {
ch <- lines
return
}
ch <- lines
} }
return err
} }
}(ctx)
if len(lines) == 0 { // read goroutine
return nil go func(c context.Context) {
var (
succeed int
total int
)
for {
select {
case <-c.Done():
return
case docs, ok := <-ch:
if !ok {
return
}
if len(docs) == 0 {
errCh <- nil
return
}
if succeed, err = output.WriteData(c, docs); err != nil {
errCh <- err
return
}
if succeed != len(docs) {
errCh <- fmt.Errorf("cmd.run: got lines=%d, only succeed=%d", len(docs), succeed)
return
}
total += succeed
logrus.Infof("Dump: succeed=%d total=%d docs succeed!!!", succeed, total)
}
} }
}(ctx)
if succeed, err = output.WriteData(ctx, lines); err != nil { select {
return err case err = <-errCh:
} return err
case <-ctx.Done():
if succeed != len(lines) { return context.Canceled
return fmt.Errorf("cmd.run: got lines=%d, only succeed=%d", len(lines), succeed)
}
total += succeed
logrus.Infof("Dump: succeed=%d total=%d docs succeed!!!", succeed, total)
} }
} }