diff options
author | GitLab Bot <gitlab-bot@gitlab.com> | 2020-12-17 14:59:07 +0300 |
---|---|---|
committer | GitLab Bot <gitlab-bot@gitlab.com> | 2020-12-17 14:59:07 +0300 |
commit | 8b573c94895dc0ac0e1d9d59cf3e8745e8b539ca (patch) | |
tree | 544930fb309b30317ae9797a9683768705d664c4 /workhorse/internal/helper/writeafterreader.go | |
parent | 4b1de649d0168371549608993deac953eb692019 (diff) |
Add latest changes from gitlab-org/gitlab@13-7-stable-eev13.7.0-rc42
Diffstat (limited to 'workhorse/internal/helper/writeafterreader.go')
-rw-r--r-- | workhorse/internal/helper/writeafterreader.go | 144 |
1 files changed, 144 insertions, 0 deletions
diff --git a/workhorse/internal/helper/writeafterreader.go b/workhorse/internal/helper/writeafterreader.go new file mode 100644 index 00000000000..d583ae4a9b8 --- /dev/null +++ b/workhorse/internal/helper/writeafterreader.go @@ -0,0 +1,144 @@ +package helper + +import ( + "fmt" + "io" + "io/ioutil" + "os" + "sync" +) + +type WriteFlusher interface { + io.Writer + Flush() error +} + +// Couple r and w so that until r has been drained (before r.Read() has +// returned some error), all writes to w are sent to a tempfile first. +// The caller must call Flush() on the returned WriteFlusher to ensure +// all data is propagated to w. +func NewWriteAfterReader(r io.Reader, w io.Writer) (io.Reader, WriteFlusher) { + br := &busyReader{Reader: r} + return br, &coupledWriter{Writer: w, busyReader: br} +} + +type busyReader struct { + io.Reader + + error + errorMutex sync.RWMutex +} + +func (r *busyReader) Read(p []byte) (int, error) { + if err := r.getError(); err != nil { + return 0, err + } + + n, err := r.Reader.Read(p) + if err != nil { + if err != io.EOF { + err = fmt.Errorf("busyReader: %v", err) + } + r.setError(err) + } + return n, err +} + +func (r *busyReader) IsBusy() bool { + return r.getError() == nil +} + +func (r *busyReader) getError() error { + r.errorMutex.RLock() + defer r.errorMutex.RUnlock() + return r.error +} + +func (r *busyReader) setError(err error) { + if err == nil { + panic("busyReader: attempt to reset error to nil") + } + r.errorMutex.Lock() + defer r.errorMutex.Unlock() + r.error = err +} + +type coupledWriter struct { + io.Writer + *busyReader + + tempfile *os.File + tempfileMutex sync.Mutex + + writeError error +} + +func (w *coupledWriter) Write(data []byte) (int, error) { + if w.writeError != nil { + return 0, w.writeError + } + + if w.busyReader.IsBusy() { + n, err := w.tempfileWrite(data) + if err != nil { + w.writeError = fmt.Errorf("coupledWriter: %v", err) + } + return n, w.writeError + } + + if err := w.Flush(); err != nil { + w.writeError = fmt.Errorf("coupledWriter: %v", err) + return 0, w.writeError + } + + return w.Writer.Write(data) +} + +func (w *coupledWriter) Flush() error { + w.tempfileMutex.Lock() + defer w.tempfileMutex.Unlock() + + tempfile := w.tempfile + if tempfile == nil { + return nil + } + + w.tempfile = nil + defer tempfile.Close() + + if _, err := tempfile.Seek(0, 0); err != nil { + return err + } + if _, err := io.Copy(w.Writer, tempfile); err != nil { + return err + } + return nil +} + +func (w *coupledWriter) tempfileWrite(data []byte) (int, error) { + w.tempfileMutex.Lock() + defer w.tempfileMutex.Unlock() + + if w.tempfile == nil { + tempfile, err := w.newTempfile() + if err != nil { + return 0, err + } + w.tempfile = tempfile + } + + return w.tempfile.Write(data) +} + +func (*coupledWriter) newTempfile() (tempfile *os.File, err error) { + tempfile, err = ioutil.TempFile("", "gitlab-workhorse-coupledWriter") + if err != nil { + return nil, err + } + if err := os.Remove(tempfile.Name()); err != nil { + tempfile.Close() + return nil, err + } + + return tempfile, nil +} |