summaryrefslogtreecommitdiff
path: root/lib
diff options
context:
space:
mode:
Diffstat (limited to 'lib')
-rw-r--r--lib/streamio/runescanner.go90
1 files changed, 90 insertions, 0 deletions
diff --git a/lib/streamio/runescanner.go b/lib/streamio/runescanner.go
new file mode 100644
index 0000000..421859d
--- /dev/null
+++ b/lib/streamio/runescanner.go
@@ -0,0 +1,90 @@
+// Copyright (C) 2022-2023 Luke Shumaker <lukeshu@lukeshu.com>
+//
+// SPDX-License-Identifier: GPL-2.0-or-later
+
+package streamio
+
+import (
+ "bufio"
+ "context"
+ "io"
+ "os"
+ "time"
+
+ "github.com/datawire/dlib/dlog"
+
+ "git.lukeshu.com/btrfs-progs-ng/lib/textui"
+)
+
+type RuneScanner interface {
+ io.RuneScanner
+ io.Closer
+}
+
+type runeScanner struct {
+ ctx context.Context //nolint:containedctx // For detecting shutdown from methods
+ progress textui.Portion[int64]
+ progressWriter *textui.Progress[textui.Portion[int64]]
+ unreadCnt uint64
+ reader *bufio.Reader
+ closer io.Closer
+}
+
+// NewRuneScanner returns an io.RuneScanner (and io.Closer) that
+// bufferes a file, similar to bufio.NewReader. There are two
+// advantages over bufio.NewReader:
+//
+// - It takes a Context, and causes reads to fail once the Context is
+// canceled; allowing large parse operations to be gracefully cut
+// short.
+//
+// - It logs the progress of reading the file via textui.Progress.
+func NewRuneScanner(ctx context.Context, fh *os.File) (RuneScanner, error) {
+ fi, err := fh.Stat()
+ if err != nil {
+ return nil, err
+ }
+ ret := &runeScanner{
+ ctx: ctx,
+ progress: textui.Portion[int64]{
+ D: fi.Size(),
+ },
+ progressWriter: textui.NewProgress[textui.Portion[int64]](ctx, dlog.LogLevelInfo, textui.Tunable(1*time.Second)),
+ reader: bufio.NewReader(fh),
+ closer: fh,
+ }
+ return ret, nil
+}
+
+// ReadRune implements io.RuneReader.
+func (rs *runeScanner) ReadRune() (r rune, size int, err error) {
+ if err := rs.ctx.Err(); err != nil {
+ return 0, 0, err
+ }
+ r, size, err = rs.reader.ReadRune()
+ if rs.unreadCnt > 0 {
+ rs.unreadCnt--
+ } else {
+ rs.progress.N += int64(size)
+ rs.progressWriter.Set(rs.progress)
+ }
+ return
+}
+
+// ReadRune implements io.RuneScanner.
+func (rs *runeScanner) UnreadRune() error {
+ if err := rs.ctx.Err(); err != nil {
+ return err
+ }
+ if err := rs.reader.UnreadRune(); err != nil {
+ return err
+ }
+ rs.unreadCnt++
+ return nil
+}
+
+// ReadRune implements io.Closer.
+func (rs *runeScanner) Close() error {
+ rs.progressWriter.Done()
+ return rs.closer.Close()
+}