|
|
@@ -8,106 +8,159 @@ import (
|
|
|
|
|
|
// BufferToBytesWriter is a Writer that writes alloc.Buffer into underlying writer.
|
|
|
type BufferToBytesWriter struct {
|
|
|
- writer io.Writer
|
|
|
+ io.Writer
|
|
|
+}
|
|
|
+
|
|
|
+func NewBufferToBytesWriter(writer io.Writer) *BufferToBytesWriter {
|
|
|
+ return &BufferToBytesWriter{
|
|
|
+ Writer: writer,
|
|
|
+ }
|
|
|
}
|
|
|
|
|
|
-// Write implements Writer.Write(). Write() takes ownership of the given buffer.
|
|
|
-func (w *BufferToBytesWriter) Write(mb MultiBuffer) error {
|
|
|
+// WriteMultiBuffer implements Writer. This method takes ownership of the given buffer.
|
|
|
+func (w *BufferToBytesWriter) WriteMultiBuffer(mb MultiBuffer) error {
|
|
|
defer mb.Release()
|
|
|
|
|
|
bs := mb.ToNetBuffers()
|
|
|
- _, err := bs.WriteTo(w.writer)
|
|
|
+ _, err := bs.WriteTo(w)
|
|
|
return err
|
|
|
}
|
|
|
|
|
|
-type writerAdapter struct {
|
|
|
- writer MultiBufferWriter
|
|
|
+func (w *BufferToBytesWriter) ReadFrom(reader io.Reader) (int64, error) {
|
|
|
+ if readerFrom, ok := w.Writer.(io.ReaderFrom); ok {
|
|
|
+ return readerFrom.ReadFrom(reader)
|
|
|
+ }
|
|
|
+
|
|
|
+ var sc SizeCounter
|
|
|
+ err := Copy(NewReader(reader), w, CountSize(&sc))
|
|
|
+ return sc.Size, err
|
|
|
}
|
|
|
|
|
|
-// Write implements buf.MultiBufferWriter.
|
|
|
-func (w *writerAdapter) Write(mb MultiBuffer) error {
|
|
|
- return w.writer.WriteMultiBuffer(mb)
|
|
|
+type BufferedWriter struct {
|
|
|
+ writer Writer
|
|
|
+ legacyWriter io.Writer
|
|
|
+ buffer *Buffer
|
|
|
+ buffered bool
|
|
|
}
|
|
|
|
|
|
-type mergingWriter struct {
|
|
|
- writer io.Writer
|
|
|
- buffer []byte
|
|
|
+func NewBufferedWriter(writer Writer) *BufferedWriter {
|
|
|
+ w := &BufferedWriter{
|
|
|
+ writer: writer,
|
|
|
+ buffer: New(),
|
|
|
+ buffered: true,
|
|
|
+ }
|
|
|
+ if lw, ok := writer.(io.Writer); ok {
|
|
|
+ w.legacyWriter = lw
|
|
|
+ }
|
|
|
+ return w
|
|
|
}
|
|
|
|
|
|
-func (w *mergingWriter) Write(mb MultiBuffer) error {
|
|
|
- defer mb.Release()
|
|
|
+func (w *BufferedWriter) Write(b []byte) (int, error) {
|
|
|
+ if !w.buffered && w.legacyWriter != nil {
|
|
|
+ return w.legacyWriter.Write(b)
|
|
|
+ }
|
|
|
|
|
|
- for !mb.IsEmpty() {
|
|
|
- nBytes, _ := mb.Read(w.buffer)
|
|
|
- if _, err := w.writer.Write(w.buffer[:nBytes]); err != nil {
|
|
|
- return err
|
|
|
+ totalBytes := 0
|
|
|
+ for len(b) > 0 {
|
|
|
+ nBytes, err := w.buffer.Write(b)
|
|
|
+ totalBytes += nBytes
|
|
|
+ if err != nil {
|
|
|
+ return totalBytes, err
|
|
|
+ }
|
|
|
+ if w.buffer.IsFull() {
|
|
|
+ if err := w.Flush(); err != nil {
|
|
|
+ return totalBytes, err
|
|
|
+ }
|
|
|
}
|
|
|
+ b = b[nBytes:]
|
|
|
}
|
|
|
- return nil
|
|
|
+ return totalBytes, nil
|
|
|
}
|
|
|
|
|
|
-type seqWriter struct {
|
|
|
- writer io.Writer
|
|
|
-}
|
|
|
+func (w *BufferedWriter) WriteMultiBuffer(b MultiBuffer) error {
|
|
|
+ if !w.buffered {
|
|
|
+ return w.writer.WriteMultiBuffer(b)
|
|
|
+ }
|
|
|
|
|
|
-func (w *seqWriter) Write(mb MultiBuffer) error {
|
|
|
- defer mb.Release()
|
|
|
+ defer b.Release()
|
|
|
|
|
|
- for _, b := range mb {
|
|
|
- if b.IsEmpty() {
|
|
|
- continue
|
|
|
- }
|
|
|
- if _, err := w.writer.Write(b.Bytes()); err != nil {
|
|
|
+ for !b.IsEmpty() {
|
|
|
+ if err := w.buffer.AppendSupplier(ReadFrom(&b)); err != nil {
|
|
|
return err
|
|
|
}
|
|
|
+ if w.buffer.IsFull() {
|
|
|
+ if err := w.Flush(); err != nil {
|
|
|
+ return err
|
|
|
+ }
|
|
|
+ }
|
|
|
}
|
|
|
|
|
|
return nil
|
|
|
}
|
|
|
|
|
|
-var (
|
|
|
- _ MultiBufferWriter = (*bytesToBufferWriter)(nil)
|
|
|
-)
|
|
|
+func (w *BufferedWriter) Flush() error {
|
|
|
+ if !w.buffer.IsEmpty() {
|
|
|
+ if err := w.writer.WriteMultiBuffer(NewMultiBufferValue(w.buffer)); err != nil {
|
|
|
+ return err
|
|
|
+ }
|
|
|
|
|
|
-type bytesToBufferWriter struct {
|
|
|
- writer Writer
|
|
|
+ if w.buffered {
|
|
|
+ w.buffer = New()
|
|
|
+ } else {
|
|
|
+ w.buffer = nil
|
|
|
+ }
|
|
|
+ }
|
|
|
+ return nil
|
|
|
}
|
|
|
|
|
|
-// Write implements io.Writer.
|
|
|
-func (w *bytesToBufferWriter) Write(payload []byte) (int, error) {
|
|
|
- mb := NewMultiBufferCap(len(payload)/Size + 1)
|
|
|
- mb.Write(payload)
|
|
|
- if err := w.writer.Write(mb); err != nil {
|
|
|
- return 0, err
|
|
|
+func (w *BufferedWriter) SetBuffered(f bool) error {
|
|
|
+ w.buffered = f
|
|
|
+ if !f {
|
|
|
+ return w.Flush()
|
|
|
}
|
|
|
- return len(payload), nil
|
|
|
+ return nil
|
|
|
}
|
|
|
|
|
|
-func (w *bytesToBufferWriter) WriteMultiBuffer(mb MultiBuffer) error {
|
|
|
- return w.writer.Write(mb)
|
|
|
+func (w *BufferedWriter) ReadFrom(reader io.Reader) (int64, error) {
|
|
|
+ var sc SizeCounter
|
|
|
+ if !w.buffer.IsEmpty() {
|
|
|
+ sc.Size += int64(w.buffer.Len())
|
|
|
+ if err := w.Flush(); err != nil {
|
|
|
+ return sc.Size, err
|
|
|
+ }
|
|
|
+ }
|
|
|
+
|
|
|
+ if readerFrom, ok := w.writer.(io.ReaderFrom); ok {
|
|
|
+ return readerFrom.ReadFrom(reader)
|
|
|
+ }
|
|
|
+
|
|
|
+ w.buffered = false
|
|
|
+ err := Copy(NewReader(reader), w, CountSize(&sc))
|
|
|
+ return sc.Size, err
|
|
|
}
|
|
|
|
|
|
-func (w *bytesToBufferWriter) ReadFrom(reader io.Reader) (int64, error) {
|
|
|
- mbReader := NewReader(reader)
|
|
|
- totalBytes := int64(0)
|
|
|
- for {
|
|
|
- mb, err := mbReader.Read()
|
|
|
- if errors.Cause(err) == io.EOF {
|
|
|
- break
|
|
|
- } else if err != nil {
|
|
|
- return totalBytes, err
|
|
|
+type seqWriter struct {
|
|
|
+ writer io.Writer
|
|
|
+}
|
|
|
+
|
|
|
+func (w *seqWriter) WriteMultiBuffer(mb MultiBuffer) error {
|
|
|
+ defer mb.Release()
|
|
|
+
|
|
|
+ for _, b := range mb {
|
|
|
+ if b.IsEmpty() {
|
|
|
+ continue
|
|
|
}
|
|
|
- totalBytes += int64(mb.Len())
|
|
|
- if err := w.writer.Write(mb); err != nil {
|
|
|
- return totalBytes, err
|
|
|
+ if _, err := w.writer.Write(b.Bytes()); err != nil {
|
|
|
+ return err
|
|
|
}
|
|
|
}
|
|
|
- return totalBytes, nil
|
|
|
+
|
|
|
+ return nil
|
|
|
}
|
|
|
|
|
|
type noOpWriter struct{}
|
|
|
|
|
|
-func (noOpWriter) Write(b MultiBuffer) error {
|
|
|
+func (noOpWriter) WriteMultiBuffer(b MultiBuffer) error {
|
|
|
b.Release()
|
|
|
return nil
|
|
|
}
|