From cf18158da47a8d86630241698f5a39f57397eb89 Mon Sep 17 00:00:00 2001 From: Alex Vanin 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 --- 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 e291b446..258aabfe 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 00000000..6ebf6910 --- /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()) + }) + }) +}