fix small Read() big rate for slow.Reader

main
bel 2025-05-17 20:53:15 -06:00
parent d73b63f43c
commit 8f18cbae3a
3 changed files with 33 additions and 9 deletions

View File

@ -11,9 +11,23 @@ import (
"os" "os"
"path" "path"
"show-rss/src/slow" "show-rss/src/slow"
"strconv"
"golang.org/x/time/rate"
"time" "time"
) )
var EnvCksumBPS = func() int {
s := os.Getenv("CKSUM_BPS")
if s == "" {
return 50_000_000
}
n, err := strconv.Atoi(s)
if err != nil || n < 1 {
panic(err)
}
return n
}()
func One(ctx context.Context, p string) error { func One(ctx context.Context, p string) error {
shortp := path.Join("...", path.Base(path.Dir(p)), path.Base(p)) shortp := path.Join("...", path.Base(path.Dir(p)), path.Base(p))
@ -86,6 +100,6 @@ func Cksum(ctx context.Context, p string) (string, error) {
defer f.Close() defer f.Close()
hasher := md5.New() hasher := md5.New()
_, err = io.Copy(hasher, slow.NewReader(ctx, 50_000_000, f)) _, err = io.Copy(hasher, slow.NewReader(ctx, rate.Limit(EnvCksumBPS), f))
return base64.StdEncoding.EncodeToString(hasher.Sum(nil)), err return base64.StdEncoding.EncodeToString(hasher.Sum(nil)), err
} }

View File

@ -17,20 +17,24 @@ func TestOne(t *testing.T) {
p := path.Join(d, "f.mkv") p := path.Join(d, "f.mkv")
os.WriteFile(p, b, os.ModePerm) os.WriteFile(p, b, os.ModePerm)
t.Logf("initial cksum...")
cksum, _ := asses.Cksum(context.Background(), p) cksum, _ := asses.Cksum(context.Background(), p)
t.Logf("one && one...")
if err := asses.One(ctx, p); err != nil { if err := asses.One(ctx, p); err != nil {
t.Fatal(err) t.Fatal(err)
} else if err := asses.One(ctx, p); err != nil { } else if err := asses.One(ctx, p); err != nil {
t.Fatal(err) t.Fatal(err)
} }
t.Logf("test -f...")
if _, err := os.Stat(p); err != nil { if _, err := os.Stat(p); err != nil {
t.Fatalf("lost original mkv: %v", err) t.Fatalf("lost original mkv: %v", err)
} else if _, err := os.Stat(path.Join(d, "f.srt")); err != nil { } else if _, err := os.Stat(path.Join(d, "f.srt")); err != nil {
t.Fatalf("no new srt: %v", err) t.Fatalf("no new srt: %v", err)
} }
t.Logf("final cksum...")
newCksum, _ := asses.Cksum(context.Background(), p) newCksum, _ := asses.Cksum(context.Background(), p)
if cksum == newCksum { if cksum == newCksum {
t.Fatalf("cksum unchanged") t.Fatalf("cksum unchanged")

View File

@ -9,30 +9,36 @@ import (
type Reader struct { type Reader struct {
ctx context.Context ctx context.Context
limiter *rate.Limiter limiter rate.Limiter
r io.Reader r io.Reader
} }
var _ io.Reader = Reader{} var _ io.Reader = &Reader{}
func NewReader(ctx context.Context, bps rate.Limit, r io.Reader) Reader { func NewReader(ctx context.Context, bps rate.Limit, r io.Reader) *Reader {
return Reader{ return &Reader{
ctx: ctx, ctx: ctx,
limiter: rate.NewLimiter(bps, int(bps)), limiter: *rate.NewLimiter(bps, int(bps)),
r: r, r: r,
} }
} }
func (r Reader) Read(b []byte) (int, error) { func (r *Reader) Read(b []byte) (int, error) {
n, err := r.r.Read(b) n, err := r.r.Read(b)
m := 0 m := 0
burst := r.limiter.Burst() burst := r.limiter.Burst()
for m < n { for m < n {
if err := r.limiter.WaitN(r.ctx, burst); err != nil { page := burst
if left := n - m; page > left {
page = left
}
if err := r.limiter.WaitN(r.ctx, page); err != nil {
return n, err return n, err
} }
m += burst
m += page
} }
return n, err return n, err