asyncwriter.go

package main

import (
	"bytes"
	"fmt"
	"io"
	"sync"
)

type AsyncWriter struct {
	maxSize int

	l     sync.Mutex
	w     io.Writer
	buf   bytes.Buffer
	state state
	err   error
}

type state int32

const (
	idle state = iota
	writing

	defaultMaxSize = 32 * 1024
)

type opt func(*AsyncWriter) error

func New(w io.Writer, opts ...opt) (*AsyncWriter, error) {
	aw := &AsyncWriter{w: w, maxSize: defaultMaxSize}
	for _, o := range opts {
		if err := o(aw); err != nil {
			return nil, err
		}
	}
	return aw, nil
}

func WithMaxSize(n int) opt {
	return func(w *AsyncWriter) error {
		if n <= 0 {
			return fmt.Errorf("illegal max size %d", n)
		}
		w.maxSize = n
		return nil
	}
}

func (w *AsyncWriter) Buffered() int {
	w.l.Lock()
	defer w.l.Unlock()
	return (&w.buf).Len()
}

func (w *AsyncWriter) Write(b []byte) (int, error) {
	return w.write(b, false)
}

func (w *AsyncWriter) AtomicWrite(b []byte) (int, error) {
	return w.write(b, true)
}

func (w *AsyncWriter) write(b []byte, allOrNothing bool) (int, error) {
	w.l.Lock()
	defer w.l.Unlock()

	if w.err != nil {
		return 0, w.err
	}

	wr := len(b)
	if mwr := w.maxSize - w.buf.Len(); wr > mwr {
		if allOrNothing {
			return 0, &ErrorNotEnoughSpace{mwr}
		}
		wr = mwr
	}
	wr, _ = w.buf.Write(b[:wr])

	if w.state == idle && w.buf.Len() > 0 {
		w.state = writing
		go w.asyncWrite()
	}

	return wr, nil
}

func (w *AsyncWriter) asyncWrite() {
	var b bytes.Buffer
	locked := false

	defer func() {
		if r := recover(); r != nil {
			err := &ErrorBackgroundWrite{}
			if e, ok := r.(error); ok {
				err.CausedBy = e
			} else {
				err.CausedBy = fmt.Errorf("%v", r)
			}

			if !locked {
				w.l.Lock()
			}
			err.NotWritten = b.Len() + w.buf.Len()
			w.err = err
			w.state = idle
			w.l.Unlock()
		}
	}()

	for {
		w.l.Lock()
		locked = true

		if w.buf.Len() == 0 {
			w.state = idle
			locked = false
			w.l.Unlock()
			return
		}

		w.buf.WriteTo(&b)

		locked = false
		w.l.Unlock()

		for b.Len() > 0 {
			_, err := (&b).WriteTo(w.w)
			if err != nil {
				panic(err)
			}
		}
	}
}

type ErrorNotEnoughSpace struct {
	Available int
}

func (e *ErrorNotEnoughSpace) Error() string {
	return fmt.Sprintf("not enough space: %d bytes available", e.Available)
}

type ErrorBackgroundWrite struct {
	CausedBy   error
	NotWritten int
}

func (e *ErrorBackgroundWrite) Error() string {
	return fmt.Sprintf("background write error: %v, %d bytes not written", e.CausedBy, e.NotWritten)
}