From b6d73cfea6e0fbcff91cbbee02aa7819c5809606 Mon Sep 17 00:00:00 2001 From: qydysky Date: Sat, 19 Aug 2023 01:29:20 +0800 Subject: [PATCH] add --- file/FileWR.go | 53 +++------------------- file/FileWR_test.go | 5 ++- io/io.go | 106 ++++++++++++++++++++++++++++++++++++++++++++ io/io_test.go | 14 ++++++ 4 files changed, 128 insertions(+), 50 deletions(-) diff --git a/file/FileWR.go b/file/FileWR.go index 89e18cd..debe6f6 100644 --- a/file/FileWR.go +++ b/file/FileWR.go @@ -10,8 +10,8 @@ import ( "path/filepath" "strings" "sync" - "time" + pio "github.com/qydysky/part/io" encoder "golang.org/x/text/encoding" ) @@ -52,7 +52,7 @@ func New(filePath string, curIndex int64, autoClose bool) *File { } } -func (t *File) CopyTo(to *File, byteInSec int64, tryLock bool) error { +func (t *File) CopyTo(to *File, copyIOConfig pio.CopyConfig, tryLock bool) error { t.getRWCloser() if t.Config.AutoClose { defer t.Close() @@ -77,10 +77,10 @@ func (t *File) CopyTo(to *File, byteInSec int64, tryLock bool) error { } defer to.l.Unlock() - return transferIO(t.read(), to.write(), byteInSec, -1) + return pio.Copy(t.read(), to.write(), copyIOConfig) } -func (t *File) CopyToIoWriter(to io.Writer, byteInSec int64, tryLock bool) error { +func (t *File) CopyToIoWriter(to io.Writer, copyIOConfig pio.CopyConfig) error { t.getRWCloser() if t.Config.AutoClose { defer t.Close() @@ -91,21 +91,7 @@ func (t *File) CopyToIoWriter(to io.Writer, byteInSec int64, tryLock bool) error } defer t.l.RUnlock() - return transferIO(t.read(), to, byteInSec, -1) -} - -func (t *File) CopyToIoWriterUntil(to io.Writer, byteInSec, totalSec int64, tryLock bool) error { - t.getRWCloser() - if t.Config.AutoClose { - defer t.Close() - } - - if !t.l.TryRLock() { - return ErrFailToLock - } - defer t.l.RUnlock() - - return transferIO(t.read(), to, byteInSec, totalSec) + return pio.Copy(t.read(), to, copyIOConfig) } func (t *File) Write(data []byte, tryLock bool) (int, error) { @@ -428,35 +414,6 @@ func newPath(path string, mode fs.FileMode) { } } -func transferIO(r io.Reader, w io.Writer, byteInSec, totalSec int64) (e error) { - if byteInSec > 0 { - ticker := time.NewTicker(time.Second) - defer ticker.Stop() - - for buf := make([]byte, byteInSec); totalSec < 0 || totalSec > 0; totalSec -= 1 { - if n, err := r.Read(buf); n != 0 { - if _, werr := w.Write(buf[:n]); werr != nil { - return err - } - } else if err != nil { - if !errors.Is(err, io.EOF) { - return err - } else { - return nil - } - } - <-ticker.C - } - } else if _, err := io.Copy(w, r); err != nil { - if !errors.Is(err, io.EOF) { - return err - } else { - return nil - } - } - return nil -} - func (t *File) write() io.Writer { if t.Config.AutoClose || t.wr == nil { t.wr = io.Writer(t.file) diff --git a/file/FileWR_test.go b/file/FileWR_test.go index 4f010ae..04217d2 100644 --- a/file/FileWR_test.go +++ b/file/FileWR_test.go @@ -9,6 +9,7 @@ import ( "testing" "time" + part "github.com/qydysky/part/io" "golang.org/x/text/encoding/simplifiedchinese" "golang.org/x/text/encoding/unicode" ) @@ -170,7 +171,7 @@ func TestCopy(t *testing.T) { } tf := New("t.txt", 0, true) - if e := sf.CopyTo(tf, 1, true); e != nil { + if e := sf.CopyTo(tf, part.CopyConfig{BytePerSec: 1}, true); e != nil { t.Fatal(e) } @@ -226,7 +227,7 @@ func TestEncoderDecoder(t *testing.T) { tf := New("UTF8.txt", 0, true) tf.Config.Coder = unicode.UTF8 - if e := sf.CopyTo(tf, 5, true); e != nil { + if e := sf.CopyTo(tf, part.CopyConfig{BytePerSec: 5}, true); e != nil { t.Fatal(e) } diff --git a/io/io.go b/io/io.go index 83b8abe..dbb4daa 100644 --- a/io/io.go +++ b/io/io.go @@ -284,3 +284,109 @@ func WithCtxCopy(ctx context.Context, callTree string, copybuf []byte, to time.D } } } + +type CopyConfig struct { + BytePerLoop, MaxLoop, MaxByte, BytePerSec uint64 +} + +// close by yourself +// +// watch out uint64(c.MaxLoop*c.BytePerLoop) overflow +func Copy(r io.Reader, w io.Writer, c CopyConfig) (e error) { + var ( + ticker *time.Ticker + leftN uint64 + ) + if c.BytePerSec > 0 { + if c.BytePerLoop == 0 || c.BytePerLoop > c.BytePerSec { + c.BytePerLoop = c.BytePerSec + } + ticker = time.NewTicker(time.Second) + defer ticker.Stop() + } + if c.BytePerLoop == 0 { + c.BytePerLoop = 1 << 17 + if c.MaxLoop == 0 { + if c.MaxByte != 0 { + leftN = c.MaxByte % c.BytePerLoop + c.MaxLoop = c.MaxByte/c.BytePerLoop + 1 + } + } else { + if c.MaxByte != 0 { + c.MaxByte = min(c.MaxByte, c.MaxLoop*c.BytePerLoop) + leftN = c.MaxByte % c.BytePerLoop + c.MaxLoop = c.MaxByte/c.BytePerLoop + 1 + } + } + } else if c.BytePerLoop > 1<<17 { + if c.MaxLoop == 0 { + if c.MaxByte != 0 { + c.BytePerLoop = 1 << 17 + leftN = c.MaxByte % c.BytePerLoop + c.MaxLoop = c.MaxByte/c.BytePerLoop + 1 + } else { + c.BytePerLoop = 1 << 17 + } + } else { + if c.MaxByte != 0 { + c.MaxByte = min(c.MaxByte, c.MaxLoop*c.BytePerLoop) + c.BytePerLoop = 1 << 17 + leftN = c.MaxByte % c.BytePerLoop + c.MaxLoop = c.MaxByte/c.BytePerLoop + 1 + } else { + c.MaxByte = c.MaxLoop * c.BytePerLoop + c.BytePerLoop = 1 << 17 + leftN = c.MaxByte % c.BytePerLoop + c.MaxLoop = c.MaxByte/c.BytePerLoop + 1 + } + } + } else { + if c.MaxLoop == 0 { + if c.MaxByte != 0 { + leftN = c.MaxByte % c.BytePerLoop + c.MaxLoop = c.MaxByte/c.BytePerLoop + 1 + } + } else { + if c.MaxByte != 0 { + c.MaxByte = min(c.MaxByte, c.MaxLoop*c.BytePerLoop) + leftN = c.MaxByte % c.BytePerLoop + c.MaxLoop = c.MaxByte/c.BytePerLoop + 1 + } else { + c.MaxByte = c.MaxLoop * c.BytePerLoop + leftN = c.MaxByte % c.BytePerLoop + c.MaxLoop = c.MaxByte/c.BytePerLoop + 1 + } + } + } + buf := make([]byte, c.BytePerLoop) + readC := uint64(0) + for { + if n, err := r.Read(buf); n != 0 { + if _, werr := w.Write(buf[:n]); werr != nil { + return err + } + if c.BytePerSec != 0 { + readC += uint64(n) + } + } else if err != nil { + if !errors.Is(err, io.EOF) { + return err + } else { + return nil + } + } + + if c.MaxLoop > 0 { + c.MaxLoop -= 1 + if c.MaxLoop == 0 { + return nil + } else if c.BytePerLoop == 1<<17 && leftN != 0 { + buf = buf[:leftN] + } + } + if c.BytePerSec != 0 && readC >= c.BytePerSec { + <-ticker.C + readC = 0 + } + } +} diff --git a/io/io_test.go b/io/io_test.go index 2dce990..9a35c66 100644 --- a/io/io_test.go +++ b/io/io_test.go @@ -1,10 +1,24 @@ package part import ( + "bytes" "io" "testing" ) +func Test_CopyIO(t *testing.T) { + var s = make([]byte, 1<<17+2) + s[1<<17-1] = '1' + s[1<<17] = '2' + s[1<<17+1] = '3' + + var w = &bytes.Buffer{} + + if e := Copy(bytes.NewReader(s), w, CopyConfig{1<<17 + 1, 1, 0, 0}); e != nil || w.Len() != 1<<17+1 || w.Bytes()[1<<17-1] != '1' || w.Bytes()[1<<17] != '2' { + t.Fatal(e) + } +} + func Test_rwc(t *testing.T) { rwc := RWC{R: func(p []byte) (n int, err error) { return 1, nil }} rwc.Close() -- 2.39.2