github.com/keybase/client/go@v0.0.0-20241007131713-f10651d043c8/kbfs/libkbfs/coalescing_context.go (about)

     1  package libkbfs
     2  
     3  import (
     4  	"reflect"
     5  	"sync"
     6  	"time"
     7  
     8  	"golang.org/x/net/context"
     9  )
    10  
    11  // CoalescingContext allows many contexts to be treated as one.  It waits on
    12  // all its contexts' Context.Done() channels, and when all of them have
    13  // returned, this CoalescingContext is canceled. At any point, a context can be
    14  // added to the list, and will subsequently also be part of the wait condition.
    15  // TODO: add timeout channel in case there is a goroutine leak.
    16  type CoalescingContext struct {
    17  	context.Context
    18  	doneCh   chan struct{}
    19  	mutateCh chan context.Context
    20  	selects  []reflect.SelectCase
    21  	start    sync.Once
    22  }
    23  
    24  const (
    25  	mutateChanSelectIndex       int = 0
    26  	closeChanSelectIndex        int = 1
    27  	numExplicitlyHandledSelects int = 2
    28  )
    29  
    30  func (ctx *CoalescingContext) loop() {
    31  	for {
    32  		chosen, val, _ := reflect.Select(ctx.selects)
    33  		switch chosen {
    34  		case mutateChanSelectIndex:
    35  			// request to mutate the select list
    36  			newCase := val.Interface().(context.Context)
    37  			if newCase != nil {
    38  				ctx.appendContext(newCase)
    39  			}
    40  		case closeChanSelectIndex:
    41  			// Done
    42  			close(ctx.doneCh)
    43  			return
    44  		default:
    45  			// The chosen channel has been closed. Remove it from our select list.
    46  			ctx.selects = append(ctx.selects[:chosen], ctx.selects[chosen+1:]...)
    47  			// If we have no more selects available, the request is done.
    48  			if len(ctx.selects) == numExplicitlyHandledSelects {
    49  				close(ctx.doneCh)
    50  				return
    51  			}
    52  		}
    53  	}
    54  }
    55  
    56  func (ctx *CoalescingContext) appendContext(other context.Context) {
    57  	ctx.selects = append(ctx.selects, reflect.SelectCase{
    58  		Dir:  reflect.SelectRecv,
    59  		Chan: reflect.ValueOf(other.Done()),
    60  	})
    61  }
    62  
    63  // NewCoalescingContext creates a new CoalescingContext. The context _must_ be
    64  // canceled to avoid a goroutine leak.
    65  func NewCoalescingContext(parent context.Context) (*CoalescingContext, context.CancelFunc) {
    66  	ctx := &CoalescingContext{
    67  		// Make the parent's `Value()` method available to consumers of this
    68  		// context. For example, this maintains the parent's log debug tags.
    69  		// TODO: Make _all_ parents' values available.
    70  		Context:  parent,
    71  		doneCh:   make(chan struct{}),
    72  		mutateCh: make(chan context.Context),
    73  	}
    74  	closeCh := make(chan struct{})
    75  	ctx.selects = []reflect.SelectCase{
    76  		{
    77  			Dir:  reflect.SelectRecv,
    78  			Chan: reflect.ValueOf(ctx.mutateCh),
    79  		},
    80  		{
    81  			Dir:  reflect.SelectRecv,
    82  			Chan: reflect.ValueOf(closeCh),
    83  		},
    84  	}
    85  	ctx.appendContext(parent)
    86  	cancelFunc := func() {
    87  		select {
    88  		case <-closeCh:
    89  		default:
    90  			close(closeCh)
    91  		}
    92  	}
    93  	return ctx, cancelFunc
    94  }
    95  
    96  func (ctx *CoalescingContext) startLoop() {
    97  	ctx.start.Do(func() {
    98  		go ctx.loop()
    99  	})
   100  }
   101  
   102  // Deadline overrides the default parent's Deadline().
   103  func (ctx *CoalescingContext) Deadline() (time.Time, bool) {
   104  	return time.Time{}, false
   105  }
   106  
   107  // Done returns a channel that is closed when the CoalescingContext is
   108  // canceled.
   109  func (ctx *CoalescingContext) Done() <-chan struct{} {
   110  	ctx.startLoop()
   111  	return ctx.doneCh
   112  }
   113  
   114  // Err returns context.Canceled if the CoalescingContext has been canceled, and
   115  // nil otherwise.
   116  func (ctx *CoalescingContext) Err() error {
   117  	ctx.startLoop()
   118  	select {
   119  	case <-ctx.doneCh:
   120  		return context.Canceled
   121  	default:
   122  	}
   123  	return nil
   124  }
   125  
   126  // AddContext adds a context to the set of contexts that we're waiting on.
   127  func (ctx *CoalescingContext) AddContext(other context.Context) error {
   128  	ctx.startLoop()
   129  	select {
   130  	case ctx.mutateCh <- other:
   131  		return nil
   132  	case <-ctx.doneCh:
   133  		return context.Canceled
   134  	}
   135  }