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

     1  package slotctx
     2  
     3  import (
     4  	"context"
     5  	"sync"
     6  )
     7  
     8  // Slot is a slot in which only one context can thrive.
     9  type Slot struct {
    10  	mu     sync.Mutex
    11  	cancel context.CancelFunc
    12  }
    13  
    14  func New() *Slot {
    15  	return &Slot{}
    16  }
    17  
    18  // Use derives a context bound to the slot.
    19  // Cancels any context previously bound to the slot.
    20  func (s *Slot) Use(ctx context.Context) context.Context {
    21  	ctx, cancel := context.WithCancel(ctx)
    22  	s.mu.Lock()
    23  	s.cancel, cancel = cancel, s.cancel
    24  	s.mu.Unlock()
    25  	if cancel != nil {
    26  		cancel()
    27  	}
    28  	return ctx
    29  }
    30  
    31  // Stop cancels the running task if there is one.
    32  func (s *Slot) Stop() {
    33  	s.mu.Lock()
    34  	defer s.mu.Unlock()
    35  	if s.cancel != nil {
    36  		s.cancel()
    37  		s.cancel = nil
    38  	}
    39  }
    40  
    41  // PrioritySlot is a slot in which only one context can thrive.
    42  type PrioritySlot struct {
    43  	mu       sync.Mutex
    44  	cancel   context.CancelFunc
    45  	priority int
    46  	shutdown bool
    47  }
    48  
    49  func NewPriority() *PrioritySlot {
    50  	return &PrioritySlot{}
    51  }
    52  
    53  // Use derives a new context.
    54  // Whichever of the argument and the incumbent are lower priority is canceled.
    55  // In a tie the incumbent is canceled.
    56  func (s *PrioritySlot) Use(ctx context.Context, priority int) context.Context {
    57  	ctx, cancel := context.WithCancel(ctx)
    58  	s.mu.Lock()
    59  	defer s.mu.Unlock()
    60  	if s.shutdown {
    61  		// Not accepting new processes.
    62  		cancel()
    63  		return ctx
    64  	}
    65  	if s.cancel == nil {
    66  		// First use
    67  		s.cancel = cancel
    68  		s.priority = priority
    69  		return ctx
    70  	}
    71  	if s.priority <= priority {
    72  		// Argument wins
    73  		s.cancel()
    74  		s.cancel = cancel
    75  		s.priority = priority
    76  		return ctx
    77  	}
    78  	// Incumbent wins
    79  	cancel()
    80  	return ctx
    81  }
    82  
    83  // Stop cancels the running task if there is one.
    84  func (s *PrioritySlot) Stop() {
    85  	s.mu.Lock()
    86  	defer s.mu.Unlock()
    87  	if s.cancel != nil {
    88  		s.cancel()
    89  		s.cancel = nil
    90  		s.priority = 0
    91  	}
    92  }
    93  
    94  // Shutdown disables the slot forever.
    95  func (s *PrioritySlot) Shutdown() {
    96  	s.mu.Lock()
    97  	defer s.mu.Unlock()
    98  	if s.cancel != nil {
    99  		s.cancel()
   100  		s.cancel = nil
   101  		s.priority = 0
   102  	}
   103  	s.shutdown = true
   104  }