From cf18158da47a8d86630241698f5a39f57397eb89 Mon Sep 17 00:00:00 2001
From: Alex Vanin <a.vanin@yadro.com>
Date: Mon, 20 Mar 2023 11:13:56 +0300
Subject: [PATCH] [#60] Implement periodic white space XML writer

Periodic white space XML writer sends XML header
and white spaces to the io.Writer.

Signed-off-by: Alex Vanin <a.vanin@yadro.com>
---
 api/handler/multipart_upload.go      | 51 ++++++++++++++++++++++++++++
 api/handler/multipart_upload_test.go | 48 ++++++++++++++++++++++++++
 2 files changed, 99 insertions(+)
 create mode 100644 api/handler/multipart_upload_test.go

diff --git a/api/handler/multipart_upload.go b/api/handler/multipart_upload.go
index e291b44..258aabf 100644
--- a/api/handler/multipart_upload.go
+++ b/api/handler/multipart_upload.go
@@ -2,6 +2,7 @@ package handler
 
 import (
 	"encoding/xml"
+	"io"
 	"net/http"
 	"net/url"
 	"strconv"
@@ -681,3 +682,53 @@ func encodeListPartsToResponse(info *layer.ListPartsInfo, params *layer.ListPart
 		Parts:            info.Parts,
 	}
 }
+
+// periodicXMLWriter creates go routine to write xml header and whitespaces
+// over time to avoid connection drop from the client. To work properly,
+// pass `http.ResponseWriter` with implemented `http.Flusher` interface.
+// Returns stop function which returns boolean if writer has been used
+// during goroutine execution. To disable writer, pass 0 duration value.
+func periodicXMLWriter(w io.Writer, dur time.Duration) (stop func() bool) {
+	if dur == 0 { // 0 duration disables periodic writer
+		return func() bool { return false }
+	}
+
+	whitespaceChar := []byte(" ")
+	closer := make(chan struct{})
+	done := make(chan struct{})
+	headerWritten := false
+
+	go func() {
+		defer close(done)
+
+		tick := time.NewTicker(dur)
+		defer tick.Stop()
+
+		for {
+			select {
+			case <-tick.C:
+				if !headerWritten {
+					_, err := w.Write([]byte(xml.Header))
+					headerWritten = err == nil
+				}
+				_, err := w.Write(whitespaceChar)
+				if err != nil {
+					return // is there anything we can do better than ignore error?
+				}
+				if buffered, ok := w.(http.Flusher); ok {
+					buffered.Flush()
+				}
+			case <-closer:
+				return
+			}
+		}
+	}()
+
+	stop = func() bool {
+		close(closer)
+		<-done // wait for goroutine to stop
+		return headerWritten
+	}
+
+	return stop
+}
diff --git a/api/handler/multipart_upload_test.go b/api/handler/multipart_upload_test.go
new file mode 100644
index 0000000..6ebf691
--- /dev/null
+++ b/api/handler/multipart_upload_test.go
@@ -0,0 +1,48 @@
+package handler
+
+import (
+	"bytes"
+	"encoding/xml"
+	"testing"
+	"time"
+
+	"github.com/stretchr/testify/require"
+)
+
+func TestPeriodicWriter(t *testing.T) {
+	const dur = 100 * time.Millisecond
+	const whitespaces = 8
+	expected := []byte(xml.Header)
+	for i := 0; i < whitespaces; i++ {
+		expected = append(expected, []byte(" ")...)
+	}
+
+	t.Run("writes data", func(t *testing.T) {
+		buf := bytes.NewBuffer(nil)
+		stop := periodicXMLWriter(buf, dur)
+
+		// N number of whitespaces + half durations to guarantee at least N writes in buffer
+		time.Sleep(whitespaces*dur + dur/2)
+		require.True(t, stop())
+		require.Equal(t, expected, buf.Bytes())
+
+		t.Run("no additional data after stop", func(t *testing.T) {
+			time.Sleep(2 * dur)
+			require.Equal(t, expected, buf.Bytes())
+		})
+	})
+
+	t.Run("does not write data", func(t *testing.T) {
+		buf := bytes.NewBuffer(nil)
+		stop := periodicXMLWriter(buf, dur)
+		time.Sleep(dur / 2)
+		require.False(t, stop())
+		require.Empty(t, buf.Bytes())
+
+		t.Run("disabled", func(t *testing.T) {
+			stop = periodicXMLWriter(buf, 0)
+			require.False(t, stop())
+			require.Empty(t, buf.Bytes())
+		})
+	})
+}