buffer copies

Author: tjungblu

recordio/async_writer.go

@@ -5,7 +5,6 @@ package recordio
 import (
 	"github.com/godzie44/go-uring/uring"
 	"os"
-	"sync/atomic"
 )
 
 // AsyncWriter takes an uring and executes all writes asynchronously. There are only two barriers: flush and close.
@@ -28,13 +27,19 @@ func (w *AsyncWriter) Write(p []byte) (int, error) {
 		}
 	}
 
-	err := w.ring.QueueSQE(uring.Write(w.file.Fd(), p, w.offset), 0, 0)
+	// TODO(thomas): we would need to make a defensive copy for p, which actually is not optimal
+	// the reason is the buffer pooling (or the header reuse). It so happens that the original backing array was written
+	// a couple times before the ring was submitted. That caused some funny offsets to be written and eventually fail reading.
+	pc := make([]byte, len(p))
+	copy(pc, p)
+
+	err := w.ring.QueueSQE(uring.Write(w.file.Fd(), pc, w.offset), 0, 0)
 	if err != nil {
 		return 0, err
 	}
 
-	atomic.AddInt32(&w.submittedSQEs, 1)
-	atomic.AddUint64(&w.offset, uint64(len(p)))
+	w.submittedSQEs++
+	w.offset += uint64(len(p))
 
 	return len(p), nil
 }
@@ -58,7 +63,7 @@ func (w *AsyncWriter) submitAwaitOne() error {
 		return err
 	}
 
-	atomic.AddInt32(&w.submittedSQEs, -1)
+	w.submittedSQEs--
 	w.ring.SeenCQE(cqe)
 
 	err = cqe.Error()
@@ -92,20 +97,20 @@ func (w *AsyncWriter) Close() error {
 	return w.file.Close()
 }
 
-func NewAsyncWriter(filePath string, numRingEntries uint32, opts ...uring.SetupOption) (WriteCloserFlusher, error) {
+func NewAsyncWriter(filePath string, numRingEntries uint32, opts ...uring.SetupOption) (WriteCloserFlusher, *os.File, error) {
 	ring, err := uring.New(numRingEntries, opts...)
 	if err != nil {
-		return nil, err
+		return nil, nil, err
 	}
 
 	writeFile, err := os.OpenFile(filePath, os.O_WRONLY|os.O_CREATE, 0666)
 	if err != nil {
-		return nil, err
+		return nil, nil, err
 	}
 
 	err = ring.RegisterFiles([]int{int(writeFile.Fd())})
 	if err != nil {
-		return nil, err
+		return nil, nil, err
 	}
 
 	writer := &AsyncWriter{
@@ -114,5 +119,5 @@ func NewAsyncWriter(filePath string, numRingEntries uint32, opts ...uring.SetupO
 		ring:     ring,
 	}
 
-	return writer, nil
+	return writer, writeFile, nil
 }

recordio/async_writer_test.go

@@ -3,14 +3,15 @@
 package recordio
 
 import (
+	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
-	"github.com/thomasjungblut/go-sstables/recordio/iouring"
 	"io/ioutil"
+	"os"
 	"testing"
 )
 
 func TestAsyncWriter_HappyPath(t *testing.T) {
-	ok, err := iouring.IsIOUringAvailable()
+	ok, err := IsIOUringAvailable()
 	require.NoError(t, err)
 	if !ok {
 		t.Skip("iouring not available here")
@@ -21,13 +22,68 @@ func TestAsyncWriter_HappyPath(t *testing.T) {
 	require.NoError(t, err)
 	defer closeCleanFile(t, temp)
 
-	writer, err := NewAsyncWriter(temp.Name(), 4)
+	writer, file, err := NewAsyncWriter(temp.Name(), 4)
 	require.NoError(t, err)
+	require.NotNil(t, file)
 
-	for i := 0; i < 10000; i++ {
-		_, err = writer.Write(randomRecordOfSize(1024))
+	var expected []byte
+	for i := 0; i < 100; i++ {
+		s := randomRecordOfSize(10)
+		_, err = writer.Write(s)
 		require.NoError(t, err)
+		expected = append(expected, s...)
 	}
 
 	require.NoError(t, writer.Close())
+	fileContentEquals(t, file, expected)
+}
+
+func TestAsyncWriter_GuardAgainstBufferReuse(t *testing.T) {
+	ok, err := IsIOUringAvailable()
+	require.NoError(t, err)
+	if !ok {
+		t.Skip("iouring not available here")
+		return
+	}
+
+	temp, err := ioutil.TempFile("", "TestAsyncWriter_GuardAgainstBufferReuse")
+	require.NoError(t, err)
+	defer closeCleanFile(t, temp)
+
+	writer, file, err := NewAsyncWriter(temp.Name(), 4)
+	require.NoError(t, err)
+	require.NotNil(t, file)
+
+	reusedSlice := []byte{13, 06, 91}
+	// we are writing the same slice, three times before a forced flush due to capacity
+	writeBuf(t, writer, reusedSlice)
+	writeBuf(t, writer, reusedSlice)
+	writeBuf(t, writer, reusedSlice)
+	// fourth time we change the slice in-place
+	reusedSlice[0] = 29
+	writeBuf(t, writer, reusedSlice)
+	writeBuf(t, writer, reusedSlice)
+	require.NoError(t, writer.Close())
+
+	fileContentEquals(t, file, []byte{
+		13, 06, 91,
+		13, 06, 91,
+		13, 06, 91,
+		29, 06, 91,
+		29, 06, 91,
+	})
+}
+
+func fileContentEquals(t *testing.T, file *os.File, expectedContent []byte) {
+	f, err := os.Open(file.Name())
+	require.NoError(t, err)
+	all, err := ioutil.ReadAll(f)
+	require.NoError(t, err)
+	assert.Equal(t, expectedContent, all)
+}
+
+func writeBuf(t *testing.T, writer WriteCloserFlusher, buf []byte) {
+	o, err := writer.Write(buf)
+	require.NoError(t, err)
+	assert.Equal(t, len(buf), o)
 }

recordio/file_writer.go

@@ -4,6 +4,7 @@ import (
 	"encoding/binary"
 	"errors"
 	"fmt"
+	"github.com/godzie44/go-uring/uring"
 	"os"
 
 	pool "github.com/libp2p/go-buffer-pool"
@@ -21,14 +22,14 @@ type FileWriter struct {
 	open   bool
 	closed bool
 
-	file              *os.File
-	bufWriter         WriteCloserFlusher
-	currentOffset     uint64
-	compressionType   int
-	compressor        compressor.CompressionI
-	recordHeaderCache []byte
-	bufferPool        *pool.BufferPool
-	directIOEnabled   bool
+	file               *os.File
+	bufWriter          WriteCloserFlusher
+	currentOffset      uint64
+	compressionType    int
+	compressor         compressor.CompressionI
+	recordHeaderCache  []byte
+	bufferPool         *pool.BufferPool
+	alignedBlockWrites bool
 }
 
 var DirectIOSyncWriteErr = errors.New("currently not supporting directIO with sync writing")
@@ -59,7 +60,7 @@ func (w *FileWriter) Open() error {
 
 	// we flush early to get a valid file with header written, this is important in crash scenarios
 	// when directIO is enabled however, we can't write misaligned blocks - thus this is not executed
-	if !w.directIOEnabled {
+	if !w.alignedBlockWrites {
 		err = w.bufWriter.Flush()
 		if err != nil {
 			return fmt.Errorf("flushing header in file at '%s' failed with %w", w.file.Name(), err)
@@ -160,9 +161,10 @@ func (w *FileWriter) Write(record []byte) (uint64, error) {
 	return prevOffset, nil
 }
 
-// WriteSync appends a record of bytes and forces a disk sync, returns the current offset this item was written to
+// WriteSync appends a record of bytes and forces a disk sync, returns the current offset this item was written to.
+// When directIO is enabled however, we can't write misaligned blocks and immediately returns DirectIOSyncWriteErr
 func (w *FileWriter) WriteSync(record []byte) (uint64, error) {
-	if w.directIOEnabled {
+	if w.alignedBlockWrites {
 		return 0, DirectIOSyncWriteErr
 	}
 
@@ -205,11 +207,14 @@ func (w *FileWriter) Size() uint64 {
 // options
 
 type FileWriterOptions struct {
-	path            string
-	file            *os.File
-	compressionType int
-	bufferSizeBytes int
-	useDirectIO     bool
+	path                  string
+	file                  *os.File
+	compressionType       int
+	bufferSizeBytes       int
+	enableDirectIO        bool
+	enableIOUring         bool
+	ioUringNumRingEntries uint32
+	ioUringOpts           []uring.SetupOption
 }
 
 type FileWriterOption func(*FileWriterOptions)
@@ -246,21 +251,35 @@ func BufferSizeBytes(p int) FileWriterOption {
 	}
 }
 
-// DirectIO is experimental: this flag enables DirectIO while writing, this currently might not work due to the misaligned allocations
+// DirectIO is experimental: this flag enables DirectIO while writing. This has some limitation when writing headers and
+// disables the ability to use WriteSync.
 func DirectIO() FileWriterOption {
 	return func(args *FileWriterOptions) {
-		args.useDirectIO = true
+		args.enableDirectIO = true
+	}
+}
+
+// IOUring is experimental: this flag enables async writes using io_uring. This has some limitation around platform, it
+// needs Linux and recent 5.x kernel to work. This currently also does not work together with DirectIO.
+func IOUring(numRingEntries uint32, opts ...uring.SetupOption) FileWriterOption {
+	return func(args *FileWriterOptions) {
+		args.enableIOUring = true
+		args.ioUringNumRingEntries = numRingEntries
+		args.ioUringOpts = opts
 	}
 }
 
 // NewFileWriter creates a new writer with the given options, either Path or File must be supplied, compression is optional.
 func NewFileWriter(writerOptions ...FileWriterOption) (WriterI, error) {
 	opts := &FileWriterOptions{
-		path:            "",
-		file:            nil,
-		compressionType: CompressionTypeNone,
-		bufferSizeBytes: DefaultBufferSize,
-		useDirectIO:     false,
+		path:                  "",
+		file:                  nil,
+		compressionType:       CompressionTypeNone,
+		bufferSizeBytes:       DefaultBufferSize,
+		enableDirectIO:        false,
+		enableIOUring:         false,
+		ioUringNumRingEntries: 4,
+		ioUringOpts:           nil,
 	}
 
 	for _, writeOption := range writerOptions {
@@ -271,13 +290,19 @@ func NewFileWriter(writerOptions ...FileWriterOption) (WriterI, error) {
 		return nil, errors.New("NewFileWriter: either os.File or string path must be supplied, never both")
 	}
 
+	if opts.enableIOUring && opts.enableDirectIO {
+		return nil, errors.New("NewFileWriter: either directIO or io_uring must be enabled, never both")
+	}
+
 	if opts.path == "" {
 		opts.path = opts.file.Name()
 	}
 
 	var factory ReaderWriterCloserFactory
-	if opts.useDirectIO {
+	if opts.enableDirectIO {
 		factory = DirectIOFactory{}
+	} else if opts.enableIOUring {
+		factory = NewIOUringFactory(opts.ioUringNumRingEntries, opts.ioUringOpts...)
 	} else {
 		factory = BufferedIOFactory{}
 	}
@@ -294,18 +319,18 @@ func NewFileWriter(writerOptions ...FileWriterOption) (WriterI, error) {
 	if err != nil {
 		return nil, fmt.Errorf("failed to create new Writer at '%s' failed with %w", opts.path, err)
 	}
-	return newCompressedFileWriterWithFile(file, writer, opts.compressionType, opts.useDirectIO)
+	return newCompressedFileWriterWithFile(file, writer, opts.compressionType, opts.enableDirectIO)
 }
 
 // creates a new writer with the given os.File, with the desired compression
-func newCompressedFileWriterWithFile(file *os.File, bufWriter WriteCloserFlusher, compType int, directIOEnabled bool) (WriterI, error) {
+func newCompressedFileWriterWithFile(file *os.File, bufWriter WriteCloserFlusher, compType int, alignedBlockWrites bool) (WriterI, error) {
 	return &FileWriter{
-		file:            file,
-		bufWriter:       bufWriter,
-		directIOEnabled: directIOEnabled,
-		open:            false,
-		closed:          false,
-		compressionType: compType,
-		currentOffset:   0,
+		file:               file,
+		bufWriter:          bufWriter,
+		alignedBlockWrites: alignedBlockWrites,
+		open:               false,
+		closed:             false,
+		compressionType:    compType,
+		currentOffset:      0,
 	}, nil
 }

recordio/file_writer_test.go

@@ -158,6 +158,11 @@ func TestWriterInitNoPath(t *testing.T) {
 	assert.Equal(t, errors.New("NewFileWriter: either os.File or string path must be supplied, never both"), err)
 }
 
+func TestWriterDirectIOAndIOUringDisabled(t *testing.T) {
+	_, err := NewFileWriter(Path("/tmp/abc"), DirectIO(), IOUring(4))
+	assert.Equal(t, errors.New("NewFileWriter: either directIO or io_uring must be enabled, never both"), err)
+}
+
 func TestWriterCrashCreatesValidHeader(t *testing.T) {
 	tmpFile, err := ioutil.TempFile("", "recordio_CrashCreatesValidHeader")
 	require.Nil(t, err)

recordio/io_uring.go

@@ -0,0 +1,44 @@
+package recordio
+
+import (
+	"github.com/godzie44/go-uring/uring"
+	"os"
+)
+
+type IOUringFactory struct {
+	numRingEntries uint32
+	opts           []uring.SetupOption
+}
+
+func (f *IOUringFactory) CreateNewReader(filePath string, bufSize int) (*os.File, ByteReaderResetCount, error) {
+	//TODO implement me
+	panic("implement me")
+}
+
+func (f *IOUringFactory) CreateNewWriter(filePath string, _ int) (*os.File, WriteCloserFlusher, error) {
+	writer, file, err := NewAsyncWriter(filePath, f.numRingEntries, f.opts...)
+	if err != nil {
+		return nil, nil, err
+	}
+
+	return file, writer, nil
+}
+
+func NewIOUringFactory(numRingEntries uint32, opts ...uring.SetupOption) *IOUringFactory {
+	return &IOUringFactory{
+		numRingEntries: numRingEntries,
+		opts:           opts,
+	}
+}
+
+// IsIOUringAvailable tests whether io_uring is supported by the kernel.
+// It will return (true, nil) if that's the case, if it's not available it will be (false, nil).
+// Any other error will be indicated by the error (either true/false).
+func IsIOUringAvailable() (available bool, err error) {
+	ring, err := uring.New(1)
+	defer func() {
+		err = ring.Close()
+	}()
+
+	return err == nil, err
+}

recordio/iouring/iouring_test.go renamed to recordio/io_uring_test.go

@@ -1,4 +1,4 @@
-package iouring
+package recordio
 
 import (
 	"github.com/stretchr/testify/require"