storj.io/uplink@v1.13.0/private/eestream/scheduler/scheduler.go (about)

     1  // Copyright (C) 2023 Storj Labs, Inc.
     2  // See LICENSE for copying information.
     3  
     4  package scheduler
     5  
     6  import (
     7  	"context"
     8  	"sync"
     9  )
    10  
    11  // Scheduler is a type to regulate a number of resources held by handles
    12  // with the property that earlier acquired handles get preference for
    13  // new resources over later acquired handles.
    14  type Scheduler struct {
    15  	opts  Options
    16  	rsema chan struct{}
    17  	hsema chan struct{}
    18  
    19  	mu      sync.Mutex
    20  	prio    int
    21  	waiters []*handle
    22  }
    23  
    24  // Options controls the parameters of the Scheduler.
    25  type Options struct {
    26  	MaximumConcurrent        int // number of maximum concurrent resources
    27  	MaximumConcurrentHandles int // number of maximum concurrent handles
    28  }
    29  
    30  // New constructs a new Scheduler.
    31  func New(opts Options) *Scheduler {
    32  	var hsema chan struct{}
    33  	if opts.MaximumConcurrentHandles > 0 {
    34  		hsema = make(chan struct{}, opts.MaximumConcurrentHandles)
    35  	}
    36  
    37  	return &Scheduler{
    38  		opts:  opts,
    39  		rsema: make(chan struct{}, opts.MaximumConcurrent),
    40  		hsema: hsema,
    41  	}
    42  }
    43  
    44  func (s *Scheduler) resourceGet(ctx context.Context, h *handle) bool {
    45  	// ensure that we don't return new resources if the context is
    46  	// already canceled.
    47  	if ctx.Err() != nil {
    48  		return false
    49  	}
    50  
    51  	// fast path: if we have a semaphore slot, then immediately return it.
    52  	select {
    53  	default:
    54  	case s.rsema <- struct{}{}:
    55  		return true
    56  	}
    57  
    58  	// slow path: add ourselves to a list of waiters so the best priority
    59  	// waiter gets the resource when it becomes available.
    60  	s.mu.Lock()
    61  	s.waiters = append(s.waiters, h)
    62  	s.mu.Unlock()
    63  
    64  	for {
    65  		select {
    66  		// someone has acquired a resource token and signaled to us.
    67  		case <-h.sig:
    68  			return true
    69  
    70  		// if we acquired a resource, then we're responsible for informing
    71  		// the appropriate handler.
    72  		case s.rsema <- struct{}{}:
    73  			// find the most appropriate handler and forward them the token.
    74  			var w *handle
    75  			s.mu.Lock()
    76  
    77  			// this condition is pretty subtle, but imagine two people join
    78  			// to get a resource at the same time, and they both remove each
    79  			// other. then the list of waiters is empty. the next time through
    80  			// the loop, they might nondeterministically get this case again
    81  			// and there are no waiters, so who gets the token? so, if there
    82  			// are no waiters, then we must have been removed from the list
    83  			// and so we must be ready to return the token. wait and do that.
    84  			if len(s.waiters) == 0 {
    85  				<-s.rsema
    86  				s.mu.Unlock()
    87  
    88  				<-h.sig
    89  				return true
    90  			}
    91  
    92  			s.waiters, w = removeBestHandle(s.waiters)
    93  			s.mu.Unlock()
    94  
    95  			w.sig <- struct{}{}
    96  
    97  		// if the context is done, we're done waiting for a resource.
    98  		case <-ctx.Done():
    99  			// try to remove ourselves from the set of waiters. if we
   100  			// couldn't be found, then someone else is going to try to
   101  			// send us the token, so we need to read it and succeed.
   102  			var removed bool
   103  			s.mu.Lock()
   104  			s.waiters, removed = removeHandle(s.waiters, h)
   105  			s.mu.Unlock()
   106  
   107  			if removed {
   108  				return false
   109  			}
   110  
   111  			<-h.sig
   112  			return true
   113  		}
   114  	}
   115  }
   116  
   117  func (s *Scheduler) numWaiters() int {
   118  	s.mu.Lock()
   119  	defer s.mu.Unlock()
   120  
   121  	return len(s.waiters)
   122  }
   123  
   124  // Join acquires a new Handle that can be used to acquire Resources.
   125  func (s *Scheduler) Join(ctx context.Context) (Handle, bool) {
   126  	if ctx.Err() != nil {
   127  		return nil, false
   128  	} else if s.hsema != nil {
   129  		select {
   130  		case <-ctx.Done():
   131  			return nil, false
   132  		case s.hsema <- struct{}{}:
   133  		}
   134  	}
   135  
   136  	s.mu.Lock()
   137  	defer s.mu.Unlock()
   138  
   139  	s.prio++
   140  
   141  	return &handle{
   142  		prio:  s.prio,
   143  		sched: s,
   144  		sig:   make(chan struct{}, 1),
   145  	}, true
   146  }
   147  
   148  // Handle is the interface describing acquired handles from a scheduler.
   149  type Handle interface {
   150  	// Get attempts to acquire a Resource. It will return nil and false if
   151  	// the handle is already Done or if the context is canceled before the
   152  	// resource is acquired. It should not be called concurrently with Done.
   153  	Get(context.Context) (Resource, bool)
   154  
   155  	// Done signals that no more resources should be acquired with Get, and
   156  	// it waits for existing resources to be done. It should not be called
   157  	// concurrently with Get.
   158  	Done()
   159  }
   160  
   161  type handle struct {
   162  	prio  int
   163  	wg    sync.WaitGroup
   164  	sched *Scheduler
   165  	sig   chan struct{}
   166  
   167  	mu   sync.Mutex
   168  	done bool
   169  }
   170  
   171  func (h *handle) Done() {
   172  	h.mu.Lock()
   173  	done := h.done
   174  	h.done = true
   175  	h.mu.Unlock()
   176  
   177  	h.wg.Wait()
   178  
   179  	if !done && h.sched.hsema != nil {
   180  		<-h.sched.hsema
   181  	}
   182  }
   183  
   184  func (h *handle) Get(ctx context.Context) (Resource, bool) {
   185  	if h.done {
   186  		return nil, false
   187  	}
   188  	ok := h.sched.resourceGet(ctx, h)
   189  	if !ok {
   190  		return nil, false
   191  	}
   192  	h.wg.Add(1)
   193  	return (*resource)(h), true
   194  }
   195  
   196  // Resource is the interface describing acquired resources from a scheduler.
   197  type Resource interface {
   198  	// Done signals that the resource is no longer in use and must be called
   199  	// exactly once.
   200  	Done()
   201  }
   202  
   203  type resource handle
   204  
   205  func (r *resource) Done() {
   206  	<-(*handle)(r).sched.rsema
   207  	(*handle)(r).wg.Done()
   208  }
   209  
   210  func removeBestHandle(hs []*handle) ([]*handle, *handle) {
   211  	if len(hs) == 0 {
   212  		return hs, nil
   213  	}
   214  	bh, bi := hs[0], 0
   215  	for i, h := range hs {
   216  		if h.prio < bh.prio {
   217  			bh, bi = h, i
   218  		}
   219  	}
   220  	return append(hs[:bi], hs[bi+1:]...), bh
   221  }
   222  
   223  func removeHandle(hs []*handle, x *handle) ([]*handle, bool) {
   224  	for i, h := range hs {
   225  		if h == x {
   226  			return append(hs[:i], hs[i+1:]...), true
   227  		}
   228  	}
   229  	return hs, false
   230  }