context_watcher.go

Documentation: github.com/jackc/pgconn/internal/ctxwatch

     1  package ctxwatch
     2  
     3  import (
     4  	"context"
     5  	"sync"
     6  )
     7  
     8  // ContextWatcher watches a context and performs an action when the context is canceled. It can watch one context at a
     9  // time.
    10  type ContextWatcher struct {
    11  	onCancel             func()
    12  	onUnwatchAfterCancel func()
    13  	unwatchChan          chan struct{}
    14  
    15  	lock              sync.Mutex
    16  	watchInProgress   bool
    17  	onCancelWasCalled bool
    18  }
    19  
    20  // NewContextWatcher returns a ContextWatcher. onCancel will be called when a watched context is canceled.
    21  // OnUnwatchAfterCancel will be called when Unwatch is called and the watched context had already been canceled and
    22  // onCancel called.
    23  func NewContextWatcher(onCancel func(), onUnwatchAfterCancel func()) *ContextWatcher {
    24  	cw := &ContextWatcher{
    25  		onCancel:             onCancel,
    26  		onUnwatchAfterCancel: onUnwatchAfterCancel,
    27  		unwatchChan:          make(chan struct{}),
    28  	}
    29  
    30  	return cw
    31  }
    32  
    33  // Watch starts watching ctx. If ctx is canceled then the onCancel function passed to NewContextWatcher will be called.
    34  func (cw *ContextWatcher) Watch(ctx context.Context) {
    35  	cw.lock.Lock()
    36  	defer cw.lock.Unlock()
    37  
    38  	if cw.watchInProgress {
    39  		panic("Watch already in progress")
    40  	}
    41  
    42  	cw.onCancelWasCalled = false
    43  
    44  	if ctx.Done() != nil {
    45  		cw.watchInProgress = true
    46  		go func() {
    47  			select {
    48  			case <-ctx.Done():
    49  				cw.onCancel()
    50  				cw.onCancelWasCalled = true
    51  				<-cw.unwatchChan
    52  			case <-cw.unwatchChan:
    53  			}
    54  		}()
    55  	} else {
    56  		cw.watchInProgress = false
    57  	}
    58  }
    59  
    60  // Unwatch stops watching the previously watched context. If the onCancel function passed to NewContextWatcher was
    61  // called then onUnwatchAfterCancel will also be called.
    62  func (cw *ContextWatcher) Unwatch() {
    63  	cw.lock.Lock()
    64  	defer cw.lock.Unlock()
    65  
    66  	if cw.watchInProgress {
    67  		cw.unwatchChan <- struct{}{}
    68  		if cw.onCancelWasCalled {
    69  			cw.onUnwatchAfterCancel()
    70  		}
    71  		cw.watchInProgress = false
    72  	}
    73  }
    74
View as plain text