zeroshade commented on code in PR #4955:
URL: https://github.com/apache/arrow-adbc/pull/4955#discussion_r4495287813
##########
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:
No need — the wrapper now holds a single reference to the underlying reader
and releases it when its own count reaches zero. Done in 3daaa840d.
--
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]
