lidavidm commented on code in PR #4955:
URL: https://github.com/apache/arrow-adbc/pull/4955#discussion_r4487751284
##########
go/adbc/drivermgr/wrapper.go:
##########
@@ -722,18 +757,145 @@ func (c *cnxn) GetOptionDouble(key string) (float64,
error) {
type stmt struct {
st *C.struct_AdbcStatement
+
+ mu sync.Mutex // serializes cancellation and release
+ closed bool
+ active *cancellationWatch
+}
+
+type cancellationWatch struct {
+ stop func() bool
+ done chan struct{}
+ once sync.Once
+}
+
+func (w *cancellationWatch) stopAndJoin() {
+ w.once.Do(func() {
+ // AfterFunc's stop does not wait for an already-started
callback.
+ if !w.stop() {
+ <-w.done
+ }
+ })
+}
+
+func (s *stmt) stopActiveWatch() {
+ s.mu.Lock()
+ old := s.active
+ s.active = nil
+ s.mu.Unlock()
+ if old != nil {
+ old.stopAndJoin()
+ }
+}
+
+func (s *stmt) stopWatching(w *cancellationWatch) {
+ if w == nil {
+ return
+ }
+ s.mu.Lock()
+ if s.active == w {
+ s.active = nil
+ }
+ s.mu.Unlock()
+ w.stopAndJoin()
+}
+
+func (s *stmt) watchCancellation(ctx context.Context) *cancellationWatch {
+ if ctx.Done() == nil {
+ return nil
+ }
+
+ w := &cancellationWatch{done: make(chan struct{})}
+ s.mu.Lock()
+ w.stop = context.AfterFunc(ctx, func() {
+ s.mu.Lock()
+ if !s.closed && s.active == w {
+ var err C.struct_AdbcError
+ C.AdbcStatementCancel(s.st, &err)
+ C.releaseErr(&err) // Cancellation is best-effort,
including NOT_IMPLEMENTED.
+ }
+ s.mu.Unlock()
+ close(w.done)
+ })
+ s.active = w
+ s.mu.Unlock()
+ return w
}
var _ adbc.StatementExecuteSchema = (*stmt)(nil)
func (s *stmt) Close() error {
+ s.stopActiveWatch()
+ s.mu.Lock()
+ defer s.mu.Unlock()
+ if s.closed {
+ return nil
+ }
+
var err C.struct_AdbcError
+ s.closed = true // The driver manager invalidates the statement even if
release fails.
if code := adbc.Status(C.AdbcStatementRelease(s.st, &err)); code !=
adbc.StatusOK {
return toAdbcError(code, &err)
}
return nil
}
+type cancellationReader struct {
+ array.RecordReader
+ ctx context.Context
+ stmt *stmt
+ watch *cancellationWatch
+ refs atomic.Int64
+ err error
+}
+
+func (r *cancellationReader) stopWatching() {
+ r.stmt.stopWatching(r.watch)
+}
+func (r *cancellationReader) Retain() {
+ r.RecordReader.Retain()
+ r.refs.Add(1)
+}
+
+func (r *cancellationReader) Release() {
+ if r.refs.Add(-1) == 0 {
+ r.stopWatching()
+ }
+ r.RecordReader.Release()
+}
Review Comment:
nit: is there a need to retain/release the underlying reader on every
retain/release? It seems that we could just release the underlying reader when
the wrapper's refcount hits 0, and save the extra atomic instruction. (Not that
it matters that much.)
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
