diff options
| -rw-r--r-- | bilibili/streaming.go | 8 | ||||
| -rw-r--r-- | common/copy.go | 78 | ||||
| -rw-r--r-- | common/minmax.go | 27 | ||||
| -rw-r--r-- | go.mod | 3 | ||||
| -rw-r--r-- | go.sum | 6 | ||||
| -rw-r--r-- | recording/runner.go | 25 |
6 files changed, 100 insertions, 47 deletions
diff --git a/bilibili/streaming.go b/bilibili/streaming.go index e8a6911..64dc26d 100644 --- a/bilibili/streaming.go +++ b/bilibili/streaming.go @@ -5,8 +5,8 @@ import ( "context" "errors" "fmt" - "io" "net/http" + "os" "strings" ) @@ -15,7 +15,9 @@ func (b Bilibili) CopyLiveStream( ctx context.Context, roomId common.RoomId, stream StreamingUrlInfo, - out io.Writer, + out *os.File, + buffer []byte, + readChunkSize int, ) (err error) { url := stream.URL if !strings.HasPrefix(url, "https://") && @@ -60,7 +62,7 @@ func (b Bilibili) CopyLiveStream( defer cancelGuardian() // blocking copy - n, err := common.Copy(ctx, out, resp.Body) + n, err := common.CopyToFileWithBuffer(ctx, out, resp.Body, buffer, readChunkSize, false) if err != nil && !errors.Is(err, context.Canceled) { // real error happens diff --git a/common/copy.go b/common/copy.go index 6f47a62..bee6515 100644 --- a/common/copy.go +++ b/common/copy.go @@ -7,36 +7,58 @@ Copied from https://ixday.github.io/post/golang-cancel-copy/ import ( "context" "io" + "os" ) -// here is some syntax sugar inspired by the Tomas Senart's video, -// it allows me to inline the Reader interface -type readerFunc func(p []byte) (n int, err error) - -func (rf readerFunc) Read(p []byte) (n int, err error) { return rf(p) } - -// Copy slightly modified function signature: -// - context has been added in order to propagate cancellation -// - (undo by Keuin) I do not return the number of bytes written, has it is not useful in my use case -func Copy(ctx context.Context, out io.Writer, in io.Reader) (written int64, err error) { - // Copy will call the Reader and Writer interface multiple time, in order - // to copy by chunk (avoiding loading the whole file in memory). - // I insert the ability to cancel before read time as it is the earliest - // possible in the call process. - written, err = io.Copy(out, readerFunc(func(p []byte) (int, error) { - - // golang non-blocking channel: https://gobyexample.com/non-blocking-channel-operations - select { - - // if context has been canceled - case <-ctx.Done(): - // stop process and propagate "context canceled" error - return 0, ctx.Err() - default: - // otherwise just run default io.Reader implementation - return in.Read(p) +// CopyToFileWithBuffer copies data from io.Reader to os.File with given buffer and read chunk size. +// The reader and file won't be closed. +// If syncFile is set, the file will be synced after every read. +func CopyToFileWithBuffer( + ctx context.Context, + out *os.File, + in io.Reader, + buffer []byte, + chunkSize int, + syncFile bool, +) (written int64, err error) { + bufSize := len(buffer) + off := 0 // offset to the end of data in buffer + nRead := 0 // how many bytes were read in the last read + defer func() { + if off+nRead > 0 { + // write unwritten data in buffer + nWrite, _ := out.Write(buffer[:off+nRead]) + written += int64(nWrite) + if syncFile { + _ = out.Sync() + } } - })) + }() - return + for { + if err = ctx.Err(); err != nil { + return + } + nRead, err = in.Read(buffer[off:Min[int](off+chunkSize, bufSize)]) + if err != nil { + return + } + off += nRead + if off == bufSize { + // buffer is full + var nWritten int + nWritten, err = out.Write(buffer) + if err != nil { + return + } + if syncFile { + err = out.Sync() + if err != nil { + return + } + } + written += int64(nWritten) + off = 0 + } + } } diff --git a/common/minmax.go b/common/minmax.go new file mode 100644 index 0000000..b670887 --- /dev/null +++ b/common/minmax.go @@ -0,0 +1,27 @@ +package common + +/* +Golang is a piece of shit. Its creators are paranoids. +*/ + +import ( + "golang.org/x/exp/constraints" +) + +type Number interface { + constraints.Integer | constraints.Float +} + +func Min[T Number](t1 T, t2 T) T { + if t1 < t2 { + return t1 + } + return t2 +} + +func Max[T Number](t1 T, t2 T) T { + if t1 > t2 { + return t1 + } + return t2 +} @@ -7,6 +7,7 @@ require ( github.com/andybalholm/brotli v1.0.4 github.com/lunixbochs/struc v0.0.0-20200707160740-784aaebc1d40 github.com/spf13/viper v1.13.0 + golang.org/x/exp v0.0.0-20220907003533-145caa8ea1d0 nhooyr.io/websocket v1.8.7 ) @@ -23,7 +24,7 @@ require ( github.com/spf13/jwalterweatherman v1.1.0 // indirect github.com/spf13/pflag v1.0.5 // indirect github.com/subosito/gotenv v1.4.1 // indirect - golang.org/x/sys v0.0.0-20220520151302-bc2c85ada10a // indirect + golang.org/x/sys v0.0.0-20220722155257-8c9f86f7a55f // indirect golang.org/x/text v0.3.7 // indirect gopkg.in/ini.v1 v1.67.0 // indirect gopkg.in/yaml.v2 v2.4.0 // indirect @@ -238,6 +238,8 @@ golang.org/x/exp v0.0.0-20191227195350-da58074b4299/go.mod h1:2RIsYlXP63K8oxa1u0 golang.org/x/exp v0.0.0-20200119233911-0405dc783f0a/go.mod h1:2RIsYlXP63K8oxa1u096TMicItID8zy7Y6sNkU49FU4= golang.org/x/exp v0.0.0-20200207192155-f17229e696bd/go.mod h1:J/WKrq2StrnmMY6+EHIKF9dgMWnmCNThgcyBT1FY9mM= golang.org/x/exp v0.0.0-20200224162631-6cc2880d07d6/go.mod h1:3jZMyOhIsHpP37uCMkUooju7aAi5cS1Q23tOzKc+0MU= +golang.org/x/exp v0.0.0-20220907003533-145caa8ea1d0 h1:17k44ji3KFYG94XS5QEFC8pyuOlMh3IoR+vkmTZmJJs= +golang.org/x/exp v0.0.0-20220907003533-145caa8ea1d0/go.mod h1:cyybsKvd6eL0RnXn6p/Grxp8F5bW7iYuBgsNCOHpMYE= golang.org/x/image v0.0.0-20190227222117-0694c2d4d067/go.mod h1:kZ7UVZpmo3dzQBMxlp+ypCbDeSB+sBbTgSJuh5dn5js= golang.org/x/image v0.0.0-20190802002840-cff245a6509b/go.mod h1:FeLwcggjj3mMvU+oOTbSwawSJRM1uh48EjtB4UJZlP0= golang.org/x/lint v0.0.0-20181026193005-c67002cb31c3/go.mod h1:UVdnD1Gm6xHRNCYTkRU2/jEulfH38KcIWyp/GAMgvoE= @@ -347,8 +349,8 @@ golang.org/x/sys v0.0.0-20210225134936-a50acf3fe073/go.mod h1:h1NjWce9XRLGQEsW7w golang.org/x/sys v0.0.0-20210423185535-09eb48e85fd7/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs= golang.org/x/sys v0.0.0-20210615035016-665e8c7367d1/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg= golang.org/x/sys v0.0.0-20220412211240-33da011f77ad/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg= -golang.org/x/sys v0.0.0-20220520151302-bc2c85ada10a h1:dGzPydgVsqGcTRVwiLJ1jVbufYwmzD3LfVPLKsKg+0k= -golang.org/x/sys v0.0.0-20220520151302-bc2c85ada10a/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg= +golang.org/x/sys v0.0.0-20220722155257-8c9f86f7a55f h1:v4INt8xihDGvnrfjMDVXGxw9wrfxYyCjk0KbXjhR55s= +golang.org/x/sys v0.0.0-20220722155257-8c9f86f7a55f/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg= golang.org/x/term v0.0.0-20201126162022-7de9c90e9dd1/go.mod h1:bj7SfCRtBDWHUb9snDiAeCFNEtKQo2Wmx5Cou7ajbmo= golang.org/x/text v0.0.0-20170915032832-14c0d48ead0c/go.mod h1:NqM8EUOU14njkJ3fqMW+pc6Ldnwhi/IjpwHt7yyuwOQ= golang.org/x/text v0.3.0/go.mod h1:NqM8EUOU14njkJ3fqMW+pc6Ldnwhi/IjpwHt7yyuwOQ= diff --git a/recording/runner.go b/recording/runner.go index 974fa0e..2f15abd 100644 --- a/recording/runner.go +++ b/recording/runner.go @@ -8,7 +8,6 @@ package recording import ( "bilibili-livestream-archiver/bilibili" "bilibili-livestream-archiver/common" - "bufio" "context" "encoding/json" "errors" @@ -26,6 +25,8 @@ type TaskResult struct { Error error } +const kReadChunkSize = 64 * 1024 + // RunTask start a monitor&download task and // put its execution result into a channel. func RunTask(ctx context.Context, wg *sync.WaitGroup, task *TaskConfig) { @@ -189,19 +190,17 @@ func record( } defer func() { _ = file.Close() }() - // buffered writer - fWriter := bufio.NewWriterSize(file, task.Download.DiskWriteBufferBytes) - defer func() { - err := fWriter.Flush() - if err != nil { - logger := log.Default() - logger.Printf("Failed to flush buffered file write data: %v", err) - } - }() - logger.Printf("Write buffer size: %v byte", fWriter.Size()) - + writeBufferSize := task.Download.DiskWriteBufferBytes + if writeBufferSize < kReadChunkSize { + writeBufferSize = kReadChunkSize + } + if mod := writeBufferSize % kReadChunkSize; mod != 0 { + writeBufferSize += kReadChunkSize - mod + } + writeBuffer := make([]byte, writeBufferSize) + logger.Printf("Write buffer size: %v byte", writeBufferSize) logger.Printf("Recording live stream to file \"%v\"...", filePath) - err = bi.CopyLiveStream(ctx, task.RoomId, streamSource, fWriter) + err = bi.CopyLiveStream(ctx, task.RoomId, streamSource, file, writeBuffer, kReadChunkSize) cancelled = err == nil || errors.Is(err, context.Canceled) if !cancelled { // real error happens |