v2fly/common/io/writer.go

package io

import (
	"io"

	"github.com/v2ray/v2ray-core/common"
	"github.com/v2ray/v2ray-core/common/alloc"
)

// Writer extends io.Writer with alloc.Buffer.
type Writer interface {
	common.Releasable
	// Write writes an alloc.Buffer into underlying writer.
	Write(*alloc.Buffer) error
}

// AdaptiveWriter is a Writer that writes alloc.Buffer into underlying writer.
type AdaptiveWriter struct {
	writer io.Writer
}

// NewAdaptiveWriter creates a new AdaptiveWriter.
func NewAdaptiveWriter(writer io.Writer) *AdaptiveWriter {
	return &AdaptiveWriter{
		writer: writer,
	}
}

// Write implements Writer.Write(). Write() takes ownership of the given buffer.
func (this *AdaptiveWriter) Write(buffer *alloc.Buffer) error {
	nBytes, err := this.writer.Write(buffer.Value)
	if nBytes < buffer.Len() {
		_, err = this.writer.Write(buffer.Value[nBytes:])
	}
	buffer.Release()
	return err
}

func (this *AdaptiveWriter) Release() {
	this.writer = nil
}