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

@@ -88,7 +88,7 @@ func fileHeaderAsByteSlice(compressionType uint32) []byte {
 }
 
 // for legacy reference still around, main paths unused - mostly for tests writing old versions
-//noinspection GoUnusedFunction
+// noinspection GoUnusedFunction
 func writeRecordHeaderV1(writer *FileWriter, payloadSizeUncompressed uint64, payloadSizeCompressed uint64) (int, error) {
 	// 4 byte magic number, 8 byte uncompressed size, 8 bytes for compressed size = 20 bytes
 	bytes := make([]byte, RecordHeaderSizeBytes)

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"

recordio/iouring/iouring.go

@@ -61,7 +61,7 @@ func TestReadWriteEndToEndDirectIO(t *testing.T) {
 		return
 	}
 
-	tmpFile, err := ioutil.TempFile("", "recordio_EndToEnd")
+	tmpFile, err := ioutil.TempFile("", "recordio_EndToEndDirectIO")
 	require.NoError(t, err)
 	defer func() { require.NoError(t, os.Remove(tmpFile.Name())) }()
 	writer, err := NewFileWriter(File(tmpFile), DirectIO())