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 }