common/utils/io2/io.go

package io2

import (
	"io"
	"sync"
)

type PromiseWriteCloser[T any] interface {
	io.Writer
	Abort(err error)    // 中断发送文件
	Finish() (T, error) // 发送文件完成，等待返回结果
}

func WriteAll(writer io.Writer, data []byte) error {
	pos := 0
	dataLen := len(data)

	for pos < dataLen {
		writeLen, err := writer.Write(data[pos:])
		if err != nil {
			return err
		}

		pos += writeLen
	}

	return nil
}

const (
	onceDisabled = 0
	onceEnabled  = 1
	onceDone     = 2
)

type readCloserHook struct {
	readCloser io.ReadCloser
	callback   func(closer io.ReadCloser)
	once       int
	isBefore   bool // callback调用时机，true则在closer的Close之前调用
	lock       *sync.Mutex
}

func (hook *readCloserHook) Read(buf []byte) (n int, err error) {
	return hook.readCloser.Read(buf)
}

func (hook *readCloserHook) Close() error {
	hook.lock.Lock()
	defer hook.lock.Unlock()

	if hook.once == onceDone {
		return hook.readCloser.Close()
	}

	if hook.isBefore {
		hook.callback(hook.readCloser)
	}

	err := hook.readCloser.Close()

	if !hook.isBefore {
		hook.callback(hook.readCloser)
	}

	if hook.once == onceEnabled {
		hook.once = onceDone
	}

	return err
}

func BeforeReadClosing(closer io.ReadCloser, callback func(closer io.ReadCloser)) io.ReadCloser {
	return &readCloserHook{
		readCloser: closer,
		callback:   callback,
		once:       onceDisabled,
		isBefore:   true,
		lock:       &sync.Mutex{},
	}
}

func AfterReadClosed(closer io.ReadCloser, callback func(closer io.ReadCloser)) io.ReadCloser {
	return &readCloserHook{
		readCloser: closer,
		callback:   callback,
		once:       onceDisabled,
		isBefore:   false,
		lock:       &sync.Mutex{},
	}
}

func AfterReadClosedOnce(closer io.ReadCloser, callback func(closer io.ReadCloser)) io.ReadCloser {
	return &readCloserHook{
		readCloser: closer,
		callback:   callback,
		once:       onceEnabled,
		isBefore:   false,
		lock:       &sync.Mutex{},
	}
}

type afterEOF struct {
	inner    io.ReadCloser
	callback func(str io.ReadCloser, err error)
	lock     *sync.Mutex
}

func (hook *afterEOF) Read(buf []byte) (n int, err error) {
	n, err = hook.inner.Read(buf)

	hook.lock.Lock()
	defer hook.lock.Unlock()

	if hook.callback != nil {
		if err == io.EOF {
			hook.callback(hook.inner, nil)
			hook.callback = nil
		} else if err != nil {
			hook.callback(hook.inner, err)
			hook.callback = nil
		}
	}
	return n, err
}

func (hook *afterEOF) Close() error {
	err := hook.inner.Close()

	hook.lock.Lock()
	defer hook.lock.Unlock()

	if hook.callback != nil {
		hook.callback(hook.inner, io.ErrClosedPipe)
		hook.callback = nil
	}
	return err
}

func AfterEOF(str io.ReadCloser, callback func(str io.ReadCloser, err error)) io.ReadCloser {
	return &afterEOF{
		inner:    str,
		callback: callback,
		lock:     &sync.Mutex{},
	}
}

type readerWithCloser struct {
	reader io.Reader
	closer func(reader io.Reader) error
}

func (hook *readerWithCloser) Read(buf []byte) (n int, err error) {
	return hook.reader.Read(buf)
}
func (c *readerWithCloser) Close() error {
	return c.closer(c.reader)
}

func WithCloser(reader io.Reader, closer func(reader io.Reader) error) io.ReadCloser {
	return &readerWithCloser{
		reader: reader,
		closer: closer,
	}
}

type LazyReadCloser struct {
	open   func() (io.ReadCloser, error)
	stream io.ReadCloser
}

func (r *LazyReadCloser) Read(buf []byte) (n int, err error) {
	if r.stream == nil {
		var err error
		r.stream, err = r.open()
		if err != nil {
			return 0, err
		}
	}

	return r.stream.Read(buf)
}

func (r *LazyReadCloser) Close() error {
	if r.stream == nil {
		return nil
	}

	return r.stream.Close()
}

func Lazy(open func() (io.ReadCloser, error)) *LazyReadCloser {
	return &LazyReadCloser{
		open: open,
	}
}

func ToReaders(strs []io.ReadCloser) ([]io.Reader, func()) {
	var readers []io.Reader
	for _, s := range strs {
		readers = append(readers, s)
	}

	return readers, func() {
		for _, s := range strs {
			s.Close()
		}
	}
}

func DropWithBuf(str io.Reader, buf []byte) {
	for {
		_, err := str.Read(buf)
		if err != nil {
			break
		}
	}
}

func ReadMost(str io.Reader, n int) ([]byte, error) {
	buf := make([]byte, n)
	n, err := io.ReadFull(str, buf)
	if err == nil {
		return buf, nil
	}

	if err == io.EOF || err == io.ErrUnexpectedEOF {
		return buf[:n], nil
	}

	return buf[:n], err
}