From 12c856a10cb9c59ae97504ce0fcd9fdb044bdd14 Mon Sep 17 00:00:00 2001
From: Keuin <keuinx@gmail.com>
Date: Thu, 15 Sep 2022 01:36:12 +0800
Subject: Use io.CopyN to utilize zero copy technique.

---
 bilibili/streaming.go | 29 ++++++++++++++++++++++-------
 1 file changed, 22 insertions(+), 7 deletions(-)

(limited to 'bilibili/streaming.go')

diff --git a/bilibili/streaming.go b/bilibili/streaming.go
index 0bda347..21bca01 100644
--- a/bilibili/streaming.go
+++ b/bilibili/streaming.go
@@ -5,6 +5,7 @@ import (
 	"errors"
 	"fmt"
 	"github.com/keuin/slbr/common"
+	"io"
 	"net/http"
 	"os"
 	"strings"
@@ -16,8 +17,7 @@ func (b Bilibili) CopyLiveStream(
 	roomId common.RoomId,
 	stream StreamingUrlInfo,
 	out *os.File,
-	buffer []byte,
-	readChunkSize int,
+	bufSize int64,
 ) (err error) {
 	url := stream.URL
 	if !strings.HasPrefix(url, "https://") &&
@@ -53,15 +53,30 @@ func (b Bilibili) CopyLiveStream(
 	defer func() { _ = resp.Body.Close() }()
 
 	b.logger.Info("Copying live stream...")
-	// blocking copy
-	n, err := common.CopyToFileWithBuffer(ctx, out, resp.Body, buffer, false, uint(len(buffer)/readChunkSize))
 
-	if err != nil && !errors.Is(err, context.Canceled) {
-		b.logger.Error("Stream copying was interrupted unexpectedly: %v", err)
+	var n int64
+
+	// blocking copy
+copyLoop:
+	for err == nil {
+		select {
+		case <-ctx.Done():
+			// cancelled
+			err = ctx.Err()
+			break copyLoop
+		default:
+			var sz int64
+			sz, err = io.CopyN(out, resp.Body, bufSize)
+			n += sz
+		}
 	}
 
-	if err == nil {
+	if errors.Is(err, context.Canceled) {
+		b.logger.Info("Stop copying...")
+	} else if errors.Is(err, io.EOF) {
 		b.logger.Info("The live is ended. (room %v)", roomId)
+	} else {
+		b.logger.Error("Stream copying was interrupted unexpectedly: %v", err)
 	}
 
 	b.logger.Info("Total downloaded: %v", common.PrettyBytes(uint64(n)))
-- 
cgit v1.2.3