improved API

Author: Noam Preil

arrow/ipc/message.go

@@ -20,10 +20,8 @@ import (
 	"encoding/binary"
 	"fmt"
 	"io"
-	"sync/atomic"
 	"unsafe"
 
-	"github.com/apache/arrow-go/v18/arrow/internal/debug"
 	"github.com/apache/arrow-go/v18/arrow/internal/flatbuf"
 	"github.com/apache/arrow-go/v18/arrow/memory"
 )
@@ -128,9 +126,9 @@ type MessageReader interface {
 
 // MessageReader reads messages from an io.Reader.
 type messageReader struct {
+	memory.Refcount
 	r io.Reader
 
-	refCount atomic.Int64
 	msg      *Message
 
 	mem    memory.Allocator
@@ -145,30 +143,11 @@ func NewMessageReader(r io.Reader, opts ...Option) MessageReader {
 	}
 
 	mr := &messageReader{r: r, mem: cfg.alloc}
-	mr.refCount.Add(1)
+	mr.Retain()
+	mr.ReferenceDependency(unsafe.Pointer(&mr.msg))
 	return mr
 }
 
-// Retain increases the reference count by 1.
-// Retain may be called simultaneously from multiple goroutines.
-func (r *messageReader) Retain() {
-	r.refCount.Add(1)
-}
-
-// Release decreases the reference count by 1.
-// When the reference count goes to zero, the memory is freed.
-// Release may be called simultaneously from multiple goroutines.
-func (r *messageReader) Release() {
-	debug.Assert(r.refCount.Load() > 0, "too many releases")
-
-	if r.refCount.Add(-1) == 0 {
-		if r.msg != nil {
-			r.msg.Release()
-			r.msg = nil
-		}
-	}
-}
-
 // Message returns the current message that has been extracted from the
 // underlying stream.
 // It is valid until the next call to Message.

arrow/ipc/reader.go

@@ -21,6 +21,7 @@ import (
 	"fmt"
 	"io"
 	"sync/atomic"
+	"unsafe"
 
 	"github.com/apache/arrow-go/v18/arrow"
 	"github.com/apache/arrow-go/v18/arrow/array"
@@ -97,7 +98,8 @@ func NewReader(r io.Reader, opts ...Option) (rr *Reader, err error) {
 	}()
 	cfg := newConfig(opts...)
 	mr := &messageReader{r: r, mem: cfg.alloc}
-	mr.refCount.Add(1)
+	mr.Retain()
+	mr.ReferenceDependency(unsafe.Pointer(&mr.msg))
 	rr = &Reader{
 		r:        mr,
 		refCount: atomic.Int64{},

arrow/memory/refcount.go

@@ -11,7 +11,7 @@ import (
 
 type Refcount struct {
 	count        atomic.Int64
-	dependencies []**Refcount
+	dependencies []unsafe.Pointer
 	buffers      []**Buffer
 	derived      []unsafe.Pointer
 }
@@ -20,7 +20,7 @@ type Refcount struct {
 // When this object is completely unreferenced, all dependencies will
 // be unreferenced by it and, if this was the only object still
 // referencing them, they will be freed as well, recursively.
-func (r *Refcount) ReferenceDependency(d ...**Refcount) {
+func (r *Refcount) ReferenceDependency(d ...unsafe.Pointer) {
 	r.dependencies = d
 }
 
@@ -31,9 +31,8 @@ func (r *Refcount) ReferenceBuffer(b ...**Buffer) {
 	r.buffers = b
 }
 
-// Must only be called once per object, with a list of pointers that are
-// _derived from_ allocations owned by or referenced by this object.
-// When this object is unreferenced, all such pointers will be nilled.
+// Must only be called once per object, with a list of pointers that need to be
+// cleared when the object becomes unreferenced.
 // Note: this needs the _address of_ the pointers to nil, _not_ the pointers
 // themselves!
 func (r *Refcount) ReferenceDerived(p ...unsafe.Pointer) {
@@ -52,8 +51,16 @@ func (r *Refcount) Release() {
 			*buffer = nil
 		}
 		for _, dependency := range r.dependencies {
-			(*dependency).Release()
-			*dependency = nil
+			ptr := (*unsafe.Pointer)(dependency)
+			if *ptr != nil {
+				// Ptr should be a **T, where T has a Refcount
+				// embedded at the front.
+				// So, if *ptr != nil, we should be able to cast *ptr
+				// to a *Refcount.
+				rc := (*Refcount)(*ptr)
+				rc.Release()
+				*ptr = nil
+			}
 		}
 		r.buffers = nil
 		r.dependencies = nil