diff --git a/task/doc.go b/task/doc.go new file mode 100644 index 000000000..ccfbea653 --- /dev/null +++ b/task/doc.go @@ -0,0 +1,42 @@ +// Copyright 2022 Molecula Corp. All rights reserved. + +// Package task provides an interface for indicating when an operation has +// been blocked, so that a worker pool which wants to be doing N things at +// a time can start trying new things when some things are blocked. +// +// To understand this, you have to start with the original context: We have +// a worker pool, which can handle up to N tasks at once. Tasks come in +// in batches, asynchronously. At most one write task can be active on a given +// database at a time, but many read tasks can be active on the database, +// with or without a write task. Each read task completes only when its entire +// containing operation completes. Write tasks can *partially* complete +// immediately, but in some cases, must wait for read tasks to finish before +// they can do crucial bookkeeping work. +// +// Regardless of the workload, we always have tasks which can progress +// available, and if we do them, eventually everything will complete. However, +// for some workloads, it is possible to pick N tasks *all of which are +// blocked*. In this case, the worker pool becomes useless. Furthermore, +// even if we don't hit that state, we can hit a state where nearly all worker +// pool tasks are blocked. +// +// To address this, we need a way for a worker pool to recognize that a worker +// has become blocked, and *start another worker*. This can result in running +// more than N workers at once. However, it rarely results in running *many* +// more. The typical case would be that we have a worker pool of N, and M of +// them are blocked waiting for write access to a given database. If one of them +// becomes unblocked, we may end up with N+1 active workers, but the other M-1 +// waiting on that database are still blocked. +// +// It might seem like the simplest thing to do is use a buffered channel as a +// semaphore, this being a standard Go idiom for pools. It's a great idiom, but +// in our case, it runs into a problem. When each worker starts, it writes into +// a buffered channel. When it becomes blocked, it reads from the channel to +// free up a slot. When it becomes unblocked, then, it has to write to the +// channel to indicate that it's taking up a slot again. But writes to the +// channel are contested, and usually only become possible when something else +// either blocks or exits... Meaning that, precisely at the moment that we have +// gained a highly contested lock and are able to proceed, we block for an +// indeterminate period of time *while holding that lock*. This is the opposite +// of what we want. +package task diff --git a/task/pool.go b/task/pool.go new file mode 100644 index 000000000..2102b468b --- /dev/null +++ b/task/pool.go @@ -0,0 +1,151 @@ +// Copyright 2022 Molecula Corp. All rights reserved. + +package task + +import ( + "sync" + "sync/atomic" +) + +// Pool represents a worker-pool type thing, which will call a given +// function in parallel aiming for a given level of concurrency. +// To use a pool, you create it, passing in a worker function; it +// then spawns goroutines to run that function in a loop. If the Pool's +// Block method is called, this marks one instance of the worker goroutine +// as blocked; the Unblock method marks it as unblocked. When there are +// insufficient unblocked goroutines, more are spawned. When there are +// excess goroutines, they exit. +// +// The pool can be shut down by calling Close(), setting its target number +// of workers to 0. +type Pool struct { + mu sync.Mutex // locker used for cond + cond *sync.Cond // notify of exiting workers + step func() + targetN int32 // desired number + unblocked int32 // currently active and unblocked + live int32 // currently active including blocked + stats PoolStats +} + +type PoolStats interface { + PoolSize(int) // reports current pool size +} + +// NewPool creates a pool that attempts to keep targetN goroutines +// active, executing step() repeatedly. It updates poolSize with the +// current size of the pool when that changes. +func NewPool(targetN int, step func(), stats PoolStats) *Pool { + p := &Pool{targetN: int32(targetN), step: step, stats: stats} + p.cond = sync.NewCond(&p.mu) + p.mu.Lock() + defer p.mu.Unlock() + for i := 0; i < targetN; i++ { + p.addWorker() + } + return p +} + +// Block marks a worker as blocked, indicating that we may need a new worker +// spawned because the caller is about to be blocked for an indeterminate +// period of time. If a new worker is needed, it's spawned immediately before +// Block returns. +func (p *Pool) Block() { + p.mu.Lock() + defer p.mu.Unlock() + unblocked := atomic.AddInt32(&p.unblocked, -1) + target := atomic.LoadInt32(&p.targetN) + if unblocked < target { + p.addWorker() + } +} + +// Unblock marks a worker as unblocked, potentially allowing the pool to +// retire a worker thread at some point in the future. +func (p *Pool) Unblock() { + atomic.AddInt32(&p.unblocked, 1) +} + +// Shutdown tells a pool to terminate by setting its desired pool size +// to zero, but does not wait for the jobs in it to stop. It is safe to +// call this before calling Close. +func (p *Pool) Shutdown() { + atomic.StoreInt32(&p.targetN, 0) +} + +// Stats reports on the pool's current state -- total live workers it +// has, how many it thinks are unblocked, and what its target is. +// These numbers are sampled individually, and there's no locking, so they +// are not guaranteed to be consistent. This is useful for approximate +// monitoring. +func (p *Pool) Stats() (live, unblocked, target int) { + return int(atomic.LoadInt32(&p.live)), int(atomic.LoadInt32(&p.unblocked)), int(atomic.LoadInt32(&p.targetN)) +} + +// Close is a Shutdown followed by waiting for all jobs to exit. +func (p *Pool) Close() { + p.mu.Lock() + p.Shutdown() + live := atomic.LoadInt32(&p.live) + for live > 0 { + p.cond.Wait() + // This line occurs while we hold p.mu. addWorker can't be called + // except from inside something that would also hold the lock. + // So, if the value can't be stale and increasing, and it can't + // increase anyway once targetN is 0. + live = atomic.LoadInt32(&p.live) + } +} + +// addWorker increments the number of unblocked things, and starts a worker. +// The unblocked count is technically wrong until the worker gets running, but +// it's right "soon". The live count maintenance is done inside the worker. +func (p *Pool) addWorker() { + // update worker count. we don't notify the condition variable because + // increasing workers can't make us more-closed. + live := atomic.AddInt32(&p.live, 1) + if p.stats != nil { + p.stats.PoolSize(int(live)) + } + atomic.AddInt32(&p.unblocked, 1) + go p.work() +} + +// work runs the provided work function in a loop as long as there's not +// too many unblocked goroutines, otherwise it exits. +func (p *Pool) work() { + defer func() { + live := atomic.AddInt32(&p.live, -1) + if p.stats != nil { + p.stats.PoolSize(int(live)) + } + // notify any waiters that we're done + if live == 0 { + p.cond.Broadcast() + } + }() + for { + unblocked := atomic.LoadInt32(&p.unblocked) + target := atomic.LoadInt32(&p.targetN) + for unblocked > target { + // Might have too many! + swapped := atomic.CompareAndSwapInt32(&p.unblocked, unblocked, unblocked-1) + if swapped { + // we've successfully removed ourselves from the unblocked count. + // now return, letting the deferred add above remove us from the live + // count as well. + return + } + // If the swap failed, unblocked increased or decreased. We + // re-extract it, and try the loop again. If it's no longer higher + // than the target, this loop ends and we continue running. + // If it's higher than the target, we'll try again with this new + // value. + // We also reload target because someone could have told us to + // terminate. + unblocked = atomic.LoadInt32(&p.unblocked) + target = atomic.LoadInt32(&p.targetN) + } + p.step() + } +} diff --git a/task/pool_test.go b/task/pool_test.go new file mode 100644 index 000000000..dbdf5f210 --- /dev/null +++ b/task/pool_test.go @@ -0,0 +1,430 @@ +// Copyright 2021 Molecula Corp. All rights reserved. + +package task + +import ( + "fmt" + "golang.org/x/sync/errgroup" + "math/rand" + "sync" + "sync/atomic" + "testing" + "time" +) + +// db represents a thing which can be locked, and which can +// perform read and write operations, which are modeled as channels which +// a workload can wait on writes to, and which embeds a lockable RWMutex. +// The RWMutex actually makes this slightly stricter than the semantics +// of RBF, which usually allows writes and reads to coexist, but fairly +// accurately represents the specific issue that RBF can't *finish* a write +// while an older read is active. Not the same, but has similar impact. +type db struct { + read, write chan struct{} + sync.Mutex +} + +// server represents a set of dbs, which jobs can be run against. They're +// [26] because they're denoted by lowercase/uppercase letters. +type server struct { + dbs [26]db + mu sync.Mutex // mutex to govern access to readers + readers [26][]workload // a list of readers associated with each db + waiters [26]struct { + mu sync.Mutex + cond *sync.Cond + } + pool *Pool + jobs chan *job + tb testing.TB +} + +// a job represents a single operation on a server, and a receiver +// waiting to hear back when it's done. It also has a reference to the +// bitmasks of read/write locks so that the parent operation can clean +// them all up when it's done. This is roughly parallel to the Qcx/Tx +// locking behavior in featurebase. +type job struct { + descr workload // the workload that generated this job, used to identify them + id int + write bool + locked *uint32 // bitmask of read-locked jobs + ch chan<- struct{} +} + +// workload represents a series of jobs as letters; +// lowercase letters read from the read channel of a component, uppercase +// letters read from the write channel of the corresponding lowercase +// component. each operation locks components as it reaches them for +// the first time, then unlocks all of them at the end of the string. +type workload string + +// runJob grabs a single job from the server's job queue, does it, and +// notifies the waiter. To "do" a job is to acquire the appropriate +// lock (for read or write), mark the appropriate bit in a bitmap of +// active locks, and then read from either a read or write channel, which +// then corresponds to values being passed to Satisfy. +func (s *server) runJob() { + j, ok := <-s.jobs + if !ok { + return + } + if j.write { + s.pool.Block() + s.dbs[j.id].Lock() + s.pool.Unblock() + s.mu.Lock() + // obtain list of existing readers + waiting := make([]workload, len(s.readers[j.id])) + copy(waiting, s.readers[j.id]) + s.mu.Unlock() + <-s.dbs[j.id].write + // In RBF, the write lock can't be released until the last outstanding + // reader predating this write terminates, but that's asynchronous + // from the actual request processing. So, similarly, we launch a thing + // that will unlock this slot in the database, once it's done waiting + // for any readers. We do that without the pool marked as blocked. + go func() { + // but the write can't actually complete until any pending readers + // that were already in play complete + if len(waiting) > 0 { + // We might need to wait for things. We need to be sure, + // though, that the server's list of readers for this isn't + // changing while we're checking it. So, we grab the specific + // lock, then check the reader list, and if we think we need + // to wait, we wait on a condition variable which then + // releases that lock so something else can update the reader + // list and notify us. + func() { + s.waiters[j.id].mu.Lock() + defer s.waiters[j.id].mu.Unlock() + // we have to check this with the specific lock held, so if + // anything were to change the list, it'd have to wait + // until we're done or waiting on the cond. + stillWaiting := s.stillWaiting(j.id, waiting) + for stillWaiting { + s.waiters[j.id].cond.Wait() + stillWaiting = s.stillWaiting(j.id, waiting) + } + }() + } + s.dbs[j.id].Unlock() + }() + } else { + s.pool.Block() + // attach us to the list of known readers, which must exit before + // any writers starting after them can exit + s.mu.Lock() + s.readers[j.id] = append(s.readers[j.id], j.descr) + s.mu.Unlock() + s.pool.Unblock() + cur := atomic.LoadUint32(j.locked) + // mask this bit in + for (cur>>j.id)&1 == 0 { + added := cur | (1 << j.id) + atomic.CompareAndSwapUint32(j.locked, cur, added) + cur = atomic.LoadUint32(j.locked) + } + <-s.dbs[j.id].read + } + j.ch <- struct{}{} +} + +// stillWaiting determines whether we're still waiting on anything in +// a given list terminating. +func (s *server) stillWaiting(id int, waitingOn []workload) bool { + s.mu.Lock() + readers := s.readers[id] + s.mu.Unlock() + for _, waiter := range waitingOn { + for _, reader := range readers { + if waiter == reader { + return true + } + } + } + return false +} + +// runWorkload runs the tasks within a workload, passing them to the worker +// queue, and then waiting for them all to complete. When it's done waiting +// for them, it releases any locks they obtained. +func (s *server) runWorkload(w workload) { + var locked uint32 + defer func() { + // unlock everything marked as locked + read := atomic.LoadUint32(&locked) + for i := 0; i < 32; i++ { + if (read>>i)&1 != 0 { + s.waiters[i].mu.Lock() + s.mu.Lock() + // remove us from readers list + for j := range s.readers[i] { + if s.readers[i][j] == w { + copy(s.readers[i][j:], s.readers[i][j+1:]) + s.readers[i] = s.readers[i][:len(s.readers[i])-1] + break + } + } + s.mu.Unlock() + s.waiters[i].mu.Unlock() + // and wake up anything that was waiting for this. + s.waiters[i].cond.Broadcast() + } + } + }() + ch := make(chan struct{}) + eg := &errgroup.Group{} + j := job{ch: ch, locked: &locked, descr: w} + for _, c := range w { + switch { + case c >= 'a' && c <= 'z': + j.id = int(c - 'a') + j.write = false + case c >= 'A' && c <= 'Z': + j.id = int(c - 'A') + j.write = true + default: + s.tb.Logf("unhandled character '%c'", c) + continue + } + j := j + eg.Go(func() error { + s.jobs <- &j + <-ch + return nil + }) + } + _ = eg.Wait() +} + +// newServer creates a server associated with the given testing.TB, +// allowing us to log things. +func newServer(tb testing.TB) *server { + s := &server{tb: tb, jobs: make(chan *job)} + for i := range s.dbs { + s.dbs[i].read = make(chan struct{}) + s.dbs[i].write = make(chan struct{}) + s.waiters[i].cond = sync.NewCond(&s.waiters[i].mu) + } + return s +} + +// close shuts the server down by closing all of its channels, and may +// not really be necessary. +func (s *server) close() { + for i := range s.dbs { + db := &s.dbs[i] + db.Lock() + close(db.read) + close(db.write) + db.Unlock() + } + close(s.jobs) +} + +// Satisfy satisfies the given read or write operations asynchronously, +// but waits for all of them in this batch to complete before returning. +func (s *server) satisfy(w workload) { + var eg errgroup.Group + for _, c := range w { + var id int + var write bool + switch { + case c >= 'a' && c <= 'z': + id = int(c - 'a') + write = false + case c >= 'A' && c <= 'Z': + id = int(c - 'A') + write = true + default: + s.tb.Logf("unhandled character '%c'", c) + continue + } + eg.Go(func() error { + if write { + s.dbs[id].write <- struct{}{} + } else { + s.dbs[id].read <- struct{}{} + } + return nil + }) + } + _ = eg.Wait() +} + +// makeWorkload generates a sequence of letters, some of which may be +// capitalized, in order +func makeWorkload() workload { + var letters [26]byte + var n int + write := rand.Intn(8) == 0 + for i := 0; i < 26; i++ { + if rand.Intn(4) == 0 { + if write { + letters[n] = 'A' + byte(i) + } else { + letters[n] = 'a' + byte(i) + } + n++ + } + } + return workload(letters[:n]) +} + +// testRandomWorkload makes up an arbitrary workload and tries to run +// the server against it. +func testRandomWorkload(t *testing.T) { + s := newServer(t) + eg := &errgroup.Group{} + p := NewPool(2, s.runJob, nil) + s.pool = p + defer p.Close() + defer s.close() + var workloads []workload // the requests we make + var quick []workload // the requests that get satisfied soon + var slow []workload // the requests that don't get satisfied until later + for i := 0; i < 10; i++ { + w := makeWorkload() + if len(w) == 0 { + continue + } + workloads = append(workloads, w) + partial := rand.Intn(26) + // possibly truncate and postpone some + if partial < len(w) { + quick = append(quick, w[:partial]) + slow = append(slow, w[partial:]) + } else { + quick = append(quick, w) + } + } + for _, w := range workloads { + w := w + eg.Go(func() error { + s.runWorkload(w) + return nil + }) + } + for _, w := range quick { + w := w + eg.Go(func() error { + s.satisfy(w) + return nil + }) + } + l, u, target := p.Stats() + // Only one worker at a time can be invoking the mark-as-blocked logic, + // so you can run after it marks that, but before the new worker is spawned, + // but the next worker can't invoke the blocked logic until that completes. + // + // Live count always decreases after unblocked count on the exit path, and + // increases before unblocked count on the startup path. So even if the + // samples are interrupted, I think it should be impossible for live + // to be less than unblocked. + if u < target-1 || l < u { + t.Fatalf("inconsistent pool stats: %d live, %d unblocked, %d target", l, u, target) + } + for _, w := range slow { + w := w + eg.Go(func() error { + s.satisfy(w) + return nil + }) + } + _ = eg.Wait() +} + +// TestRandomWorkloads makes up some arbitrary workloads, then tries to +// satisfy them out of order. +// In theory, this should work for any sequence of operations as long as +// no operation has the same letter for both read and write ops, and +// ops always occur in order. +func TestRandomWorkloads(t *testing.T) { + for i := 0; i < 10; i++ { + t.Run(fmt.Sprintf("%d", i), func(t *testing.T) { + testRandomWorkload(t) + }) + } +} + +func TestServer(t *testing.T) { + s := newServer(t) + eg := &errgroup.Group{} + p := NewPool(3, s.runJob, nil) + s.pool = p + defer p.Close() + defer s.close() + request := func(w workload) { + eg.Go(func() error { + s.runWorkload(w) + return nil + }) + } + // requesting "abcd" means that the reader "abcd" will still be active on + // a until all the other letters show up. + request("abcd") + // satisfy won't complete until at least two of the jobs have happened, + // so there's a decent chance that we've marked ourselves as a reader on + // a. + s.satisfy("abc") + // so we request a write on A. we get the write lock, but we can't + // relinquish it until "d" shows up. + request("A") + // satisfy that request immediately, but to no avail. + s.satisfy("A") + // three more requests come in. if they get pool slots, they definitely + // block; that request on A can't have finished yet. so they could fully + // block our work pool. + request("A") + request("A") + request("A") + // spawn something to provide "efg" + go s.satisfy("efg") + // runWorkload means we actually block waiting for it. if all the worker + // pool is blocked waiting on A, we can't do that. + s.runWorkload("efg") + // now we provide the missing d, which should allow the first request to + // finally complete, and then the next three A, which should finish + // the rest. + s.satisfy("dAAA") + // If we spawned new jobs, this should complete. Otherwise it should hang + // because the requests can't be satisfied because the queue is full + // of blocked operations. + _ = eg.Wait() +} + +func TestPoolStartup(t *testing.T) { + var counter int32 + started := make(chan struct{}) + done := make(chan struct{}) + addAndWait := func() { + <-started + atomic.AddInt32(&counter, 1) + <-done + } + // we expect this to spawn three counters + p := NewPool(3, addAndWait, nil) + time.Sleep(50 * time.Millisecond) + v := atomic.LoadInt32(&counter) + if v != 0 { + t.Fatalf("expected no adds yet, got %d", v) + } + close(started) + time.Sleep(50 * time.Millisecond) + v = atomic.LoadInt32(&counter) + if v != 3 { + t.Fatalf("expected 3 adds, got %d", v) + } + // Tell the pool to stop processing jobs + p.Shutdown() + // Allow the jobs to complete. Since this happens after the + // shutdown has set desired pool size to zero, they should now all exit. + close(done) + p.Close() + time.Sleep(50 * time.Millisecond) + v = atomic.LoadInt32(&counter) + if v != 3 { + t.Fatalf("expected no more adds, got %d including previous 3", v) + } +}