[clangd] The new threading implementation

[ilya-biryukov]

Summary

In the new threading model clangd creates one thread per file to manage
the AST and one thread to process each of the incoming requests.
The number of actively running threads is bounded by the semaphore to
avoid overloading the system.

[ilya-biryukov]

clangd/ClangdServer.h

101

These fields should probably be grouped into multiple groups:

• Inputs - capture the latest input. Can only be used on the main thread.
• Resource->Preamble - capture some built preamble (that's always available). Can be accessed on both the main thread and any of the worker threads.
• LastRequestIsUpdate and LastUpdateCF - capture the state of last update to allow cancelling it. Can be accessed only on the main thread.
• Queue and Resources->AST - capture the latest AST, Queue is used to schedule requests on the main thread, Resources->AST can only be accessed from actions on scheduled on the corresponding queue.

[sammccall]

Thanks, this helps me understand where previous patch is coming from!

I have some comments on the ThreadPool part, which basically amount to trying to represent the same abstract structure with fewer pieces.
But I still want to consider whether that's the right structure (specifically: whether the Queue abstraction makes it awkward to express task interactions like diagnostics debouncing).

So please don't do any work based on these comments yet! More thoughts to come...

clangd/ClangdServer.h

96–97

So overall here, I think that we can drop ThreadPool without much impact on the design.
Queue would stay, as a wrapper around std::thread that lets you add tasks to its runloop. It would be owned by FileData, and shutdown would be triggered by its destructor.

The advantage is one less layer here to understand, and an attractive nuisance to tweak over time.
The downsides I see:

• threading is no longer abstracted away, so changes to it are less isolated I think it's likely that we get away with the model staying simple. If it gets complicated, we'll want to add the abstraction back in. But this isn't java, we can abstract when we need it :-)
• RunSynchronously escapes into Scheduler. This is for the worse I think, but only slightly.

130

As discussed offline, Queue's semantics are very similar to a thread, except:

• cooperatively scheduled: Queues do not yield until idle
• limited in parallelism by the size of the threadpool

In the interests of keeping things simple and familiar, I think we should start out by using std::thread.
We can use a semaphore to limit the parallelism (I'm fine with doing this in the first iteration, but urge you to consider leaving it out until we see actual problems with using the CPU as our limit). This should give both properties above (if the Queue only releases the semaphore when idle).
If we still have performance problems we may need to switch to a multiplexing version, though I have a hard time imagining it (e.g. even on windows, thread spawn should be <1us, and memory usage is trivial compared to anything that touches an AST).

131

Similarly, the free requests look a lot like standalone threads, with a few enhancements that are implementable but also possible YAGNI.

• running code complete first is roughly[1] equivalent to elevating the thread's priority (no portability shim in llvm yet, but it's pretty easy). I don't think this needs to be in the first patch.
• limiting parallelism can be done with semaphores. In fact we can easily express something like "up to 18 queued tasks, up to 20 free tasks, up to 20 total", which nice for latency.

[1] I see both advantages and disadvantages, happy to discuss more!

[sammccall]

OK, here's a braindump, probably nothing surprising after our offline discussion.

TL;DR: the only design I *know* is extensible in ways we care about is the one that basically shards Scheduler into a bunch of per-TU schedulers.
I think that's probably the way to go, but if you can show how to cleanly extend one of the other options, we should compare those on the merits.

I think if we do adopt and stick with the sharding design, then a public API that reflects it is slightly nicer overall.
But it's not that important to me, tastes may vary, and above all I don't want to block the API patch on coming to a decision on the question above.
So I think we should agree on a name and then land the API patch with the interface it has now.

clangd/ClangdServer.cpp

192 ↗

(On Diff #131567)

So... scheduler and the queue abstraction.
This works pretty well for the functionality that's in this patch, which is at least as good as what's in the live version.
However I think we want to e.g. debounce diagnostic requests, and likely other such things, e.g. Indexing is a nice candidate for debouncing *and* preemption.

(Current diagnostics behavior: we compute and send diagnostics for the syntax error introduced by the first keystroke. This is wasted load, and increases latency for diagnostics on the *last* keystroke, since we have to wait for the previous one to finish)

So we should consider some other options...

193 ↗

(On Diff #131567)

Option 1: Put all timing-related logic in scheduler.

I think this probably means Scheduler itself gets a thread whose job it is to wake up at certain times and schedule or cancel tasks on the queues.
update would set an alarm to add a task on the queue, which would be invalidated by any subsequent update. The easiest way to implement this is probably to track a version number that increments on update calls and is captured by the alarm.

This gets complicated by reads after updates that haven't finished yet - you need force the update to happen now, so you need to schedule the thing and prevent the alarm from scheduling it.
It seems like in general trying to do nontrivial things may end up with a bunch of interacting actors and shared data structures that are hard to reason about and verify.

194 ↗

(On Diff #131567)

Option 2: Share timing-related additions between scheduler and queue, but keep the tasks opaque to the queue.
This means extending the queue abstraction e.g. instead of just tasks, understand (delay, priority, task) tuples.

This isn't all that different from option 1, but moving some of the pieces out of scheduler might reduce the ball of hair to a manageable size, and add some conceptual clarity.

195 ↗

(On Diff #131567)

Option 3: extend the responsibilities of Queue so this is naturally its concern.
It needs to understand the timing, dependencies between tasks, and it probably makes sense to know when they become redundant etc too.

The easiest way to express this is that Queue owns a thread and implements the run-loop for actions on a single TU, and Scheduler's operations basically delegate to Queue.

FWIW, this can be transformed into a version without dedicated threads. First we write the thread version in this form:

runLoop() {
  while (!shutdown) {
    nextThing = decideWhatToDo(); // either action, or a sleep duration
    if (nextThing.action)
      nextThing.action();
    else
      interrupt.wait(nextThing.sleepDuration); // any update(), withAST etc interrupts
  }
}

I don't think it'd be hard to write a threadpool-based executor that runs "workloads" defined by their decideWhatToDo() function and their interactions with the interrupt condition variable. (I don't know exactly what a nice API would look like, but details...)

That said, I definitely think the thread/runloop based one is an easier place to start, and suspect it would be good enough.

196 ↗

(On Diff #131567)

My conclusion (feel free to disagree!)
We should not add timing functionality in the first patch, but we should strongly prefer a design that we know it can be added to.

• Option 1 looks pretty messy/complex to me. If you want to go this route, I'd really like to see a prototype with timing stuff.
• Option 2 looks maybe workable, I can't say. If you can describe the right extension to Queue, this might work. I worry about baking in an abstraction that's too general to have a simple implementation, or too specific and needs extension in the future, but let's discuss.
• Option 3 seems to be obviously implementable with the timing extensions. It also seems pretty flexible to modify in the future, without high abstraction barriers between policy and implementation.

My preference/advice would be to take option 3 now, I think it's at least as good as the others and doesn't need additional research. Of course I'm not opposed if you want to dig into options 1/2 and show that they can work.

(Separately, I'd prefer to initially use std::thread rather than a threadpool abstraction for the reasons laid out in the previous comment. Especially for option 3, where the API for the threadpool equivalent isn't obvious or standard. But *mostly* these questions are orthogonal)

clangd/threading/Cancellation.h

33 ↗

(On Diff #131567)

This is a nice little class!

I think callsites will look more natural if this is spelled operator bool() - you can call the variable Cancelled and write if (Cancelled) and Cancelled.set()

Maybe there's a case for spelling set as operator= too, and asserting that the arg is true, but i'm less sure of that one!

[ilya-biryukov]

• An initial version of thread-per-file approach.

This is by no means a final version, we should definitely move things between files, do some renames, etc. before landing the final version.
Some things are not used anymore (e.g. ThreadPool), but are still in the patch, we'll need to remove them too.

The new version does not drop updates that are immidieately followed by other updates, this seems like an oversight and should be fixed too.
Please take a look at the overall design and let me know what you think. Feel free to add suggestions on how we can improve things, too!

[sammccall]

Not a complete review, but this looks pretty good!
You probably want to read the last comment first - the gc thread is the threadpooliest bit of code left, and you may (or may not) want to eliminate it.

clangd/TUScheduler.cpp

31

I tend to find thread bodies more readable as a member function, up to you

75

in the spirit of "just spawn a thread, and write direct code"...

can we just spawn a shared_ptr<std::thread> to do the work here, and replace GCThread with a vector<weak_ptr<std::thread>>.
Then ~TUScheduler could just loop through:

for (const auto &WeakThread : Cleanups)
  if (auto SharedThread = WeakThread.lock()) // thread may have died already
    SharedThread->join();

might be simpler? You need to lock the Cleanups vector, but no fiddling with CVs etc.

110

this seems like where shared_ptr might help us out if we're willing to deal with it.
If we captured a shared_ptr<File> here, then we'd want to call getPossiblyStalePreamble inside rather than outside the lambda.

So this would look something like:

• FileASTThread does thread shutdown as part of destructor, and requires the TUScheduler to outlive it
• TUScheduler keeps (conceptually): active files StringMap<shared_ptr<FileASTThread>>, and shutdown handles vector<weak_ptr<thread>>
• removal takes the shared_ptr out of the map, transfers ownership of the shared_ptr to a new thread which does nothing (except destroy the shared_ptr), and stashes a weak_ptr to the thread as a shutdown handle, which get looped over in the destructor.

Is this simpler? I'm not sure. It's maybe more thematically consistent :-) and it solves your "preamble got built" problem.

clangd/TUScheduler.h

26

Probably belongs in Threading.h instead.

Maybe just Semaphore? If we have *two* semaphores in clangd, we've jumped the shark!
(Isn't a non-counting semaphore just a mutex?)

39

Maybe FileWorker? this has-a thread, rather than is-a

48

AFAICS, we require callers to setDone() before destroying, and the caller always does this immediately before destroying.

Why not just make the destructor do this?

52

do we want to move the callback to be a clangd-global thing, rather than a per-update thing, before this patch or later?

(And will that change be mirrored here, or will FileASTThread::update still take a callback?)

54

consider mirroring the runWithAST name

59

(no action required, just random thoughts)

one day we should really get the threading annotations (REQUIRES_EXCLUSIVE_LOCK etc macros) set up in LLVM - they're *implemented* in llvm for crying out loud...

61

Having trouble seeing which parts of this are really used.
My understanding is that the ThreadPool+Queue abstraction is going away.
This means this is a deque + lock + ... a "done" state of some kind? Is more of it still live?

I'd probably expect the 'done' state to be redundant with the one here... If it's just a deque + a lock, I'd lean towards making it really thin, maybe even a plain struct, and putting it in this file.

But not sure of your thoughts here.

[ilya-biryukov]

clangd/TUScheduler.cpp

75

The problem is that stored vector will have unbounded growth, and I think it shouldn't.
What we could do instead, though, is to have a counter that increments with each spawned request and wait for it to go down to 0 when TUScheduler dies. That should be simpler than what we have now and we have vectors that can grow indefinitely.

[ilya-biryukov]

clangd/TUScheduler.h

52

I think we're fine both ways.
Those changes are pretty much independent of the other parts of the code.
I'd personally first finish this patch and remove the callback afterwards. But I could also later update the patch to have the new callback semantics if you beat me to it.

59

Yeah, they're certainly great. The sole reason we don't have them enabled is because LLVM does not use threading that much?

61

I'll probably remove RequestsQueue or turn it into something really simple like you suggest.
It certainly has more state than we need with this patch. This is one of the reasons this patch is far from being ready.

[ilya-biryukov]

Cleaned up the patch and added the missing bits.
This is in a much better shape now and should be ready for review.
I'll have a go through the existing review comments to make sure all concerns were addressed.

[ilya-biryukov]

clangd/ClangdServer.h

96–97

I tried fiddling with the idea and ended up abandoning the RequestQueue in favor of the std::queue with explicit locks.
There's only one place where we queue requests now (processing the updates/reads of the ASTs), everything else is managed by creating separate std::threads.

130

Done exactly that. There's a new abstraction in the patch that manages creating threads and waiting for them to finish for us. It seems pretty simple, please take a look and let me know what you think.

[ilya-biryukov]

clangd/ClangdServer.h

131

We're not setting the priorities in the first version, but certainly agree we should add that later.
And we're using a simple semaphore to limit the number of actively running threads.

clangd/TUScheduler.cpp

31

Moved it as a member function and made a (questionable) decision that clients of the class are responsible for running that function on a separate thread themselves.
It allows to reuse the abstraction that runs the threads and waits for them to finish (also used when spawning threads for completion, and similar).
It shouldn't be too bad, given that we have just a single client. Even though the interface is not particularly nice.

Happy to chat about it.

110

I ran into related problem with completion threads: we'd have to store somewhere in order to wait for there completion.
Ended up creating a thing that counts all running threads and allows to wait for all of them to finish.
I feel it looks very simple now, please let me know what you think

clangd/TUScheduler.h

26

Renamed to Semaphore and moved to Threading.h.

39

Renamed it to ASTWorker. It doesn't really know anything really caret about files directly, just manages the AST.

48

I still think that setDone() is useful to check the invariant that new requests are not being scheduled after FileData is removed.

61

queue + lock + done is how we do this in the new patch. RequestQueue is removed

[ilya-biryukov]

• Removed redundant includes

[sammccall]

Really coming together!

clangd/ASTWorker.cpp

2

This file could really use some high-level comments about the scheduling strategy, throughout

21

?!

103

This strategy has some upsides:

• we eventually get diagnostics for the latest version
• we don't begin building an AST that we know can never be read
• we don't "wait" when there's something useful to do
• it's simple to implement (apart from cancellation)

And downsides:

• diagnostics are built on the first keystroke
• indexing blocks interactive actions, because we don't (and can't reasonably) respect cancellation before calling OnUpdated
• it requires CancellationFlag even though our actions aren't cancelable

What's the goal here - is this a strategy to keep? Or a placeholder? Or trying to maximize compatibility with the previous code?

clangd/ASTWorker.h

2

ASTWorker is an implementation detail of TUScheduler. Can we move it to TUScheduler.cpp, instead of exposing it, and remove this header?

20

This seems like an "implementation header" - nobody should depend on ASTWorker AIUI.

So InputsAndAST shouldn't really go here, as it's part of the TUScheduler public interface.

25

(InputsAndPreamble is entirely unused here, I think)

30

updated -> updates

39

As discussed offline, this lifecycle is a bit complicated :-)
ASTWorker basically manages a thread, and it'd be nice if we could align the object lifetime and thread lifetime more closely.

The difficulty seems to be that we want to TUScheduler to be able to discard ASTWorkers instantly, even though we can't interrupt them.
After offline brainstorming, we came up with some sort of "ASTWorkerHandle" which:

• owns the thread that run() is called on
• detaches that thread on destruction
• has a weak_ptr or shared_ptr to the worker itself, which it exposes

We still have run() and setDone(), but they're private details.

44

stop()? (or stopSoon/requestStop to be more explicit)

46

can we reorder/group these by purpose/sequence?
e.g (no need for the comments, just illustrative).

//lifecycle
run();
setDone();

//write
update()

//read
runWithAST()
getPossiblyStalePreamble()

//misc
getUsedBytes()

55

I think this might actually be easier to read with just using Request = UniqueFunction<void()> and then spelling out pair<Request, Context>. up to you though.

60

nit: "combine ... into one class"?

I'd hope we won't end up with ASTWorker, CppFile, FileInputs, *and* another thing?

70

Why?
e.g. This allows us to cancel an update that was never read, if a subsequent update comes in

clangd/TUScheduler.cpp

38

destructor will do this

clangd/TUScheduler.h

77

This would be superseded by ASTWorkerHandle, right?

90

inline this? or come find a more descriptive name :-)

clangd/Threading.h

62

nit: "allows to run tasks" --> "Runs tasks"

A little more motivation would be useful here: this describes what is done, but not really why.

e.g. Objects that need to spawn threads can own an AsyncTaskRunner to ensure they all complete on destruction.

65

comment

67

any need to expose this separately from the destructor?

[ilya-biryukov]

• Changed interface of ASTWorker so that it runs the processing loop itself.

[ilya-biryukov]

• Removed ASTWorker files, moved all the code to TUScheduler.cpp
• Renamed setDone to stop
• Added a comment to TUScheduler.cpp
• Addressed other review comments

[ilya-biryukov]

clangd/ASTWorker.cpp

2

I've added the docs to TUScheduler.cpp

103

Trying to maximize the compatibility with existing code. There are certainly other strategies to schedule updates.
I'm perfectly happy to explore those, but would try to keep it out of this patch, it's touching quite a lot of things already. And we should probably remove the callback from update in order to make implementing the timeouts simpler

clangd/ASTWorker.h

46

I used a different grouping:

//lifecycle
run();
setDone();

//async write
update()
// async read
runWithAST()

// sync reads
getPossiblyStalePreamble()
getUsedBytes()

Does that make sense? Or do you feel read/write vs sync/async is a better distinction?

55

I'd rather keep it as is. I specifically came up with this typedef because I got annoyed of typing pair<Request, Context>.
Happy to change it if you feel that's totally unreadable, of course.

60

Removed the FIXME

clangd/TUScheduler.cpp

38

It's much safer to call it explicitly to not depend no the order of fields in the class.

clangd/TUScheduler.h

77

It's still there, stores inputs and ASTWorkerHandle.

Inputs in the ASTWorker correspond to the latest processed update, we don't expose them in the API.
These inputs correspond to the latest inputs of the file (they are used by runWithPreamble to provide proper inputs).

clangd/Threading.h

67

It's a useful helper to have.
We can use it to rewrite tests when removing future returned from addDocument/removeDocument.
Also see my other comment about calling it in destructor of TUScheduler.

[sammccall]

Very nice! Thanks for adding the docs to TUScheduler implementation, makes a big difference.

Rest is almost all readability/comment bits. Substantive stuff is:

• getUsedBytes() looks racy
• I'm not sure we're choosing the right preamble

My understanding is the threading-related bits of CppFile (locking, and deferRebuild etc) are obsolete after this patch.
It's fine not to delete them here (I guess there could be some fallout), but maybe you want to add a comment in this patch marking them as obsolete?

Testing: I think this is mostly covered by the existing TUScheduler tests. I'd suggest adding a unit test to AsyncTaskRunner though: e.g. start a bunch of threads while holding a mutex that prevents them from making progress (to verify actually async), then have them increment a counter and verify that the counter is N after waitForAll() returns.

clangd/ASTWorker.cpp

103

This sounds fine - can you add a little bit of this rationale and things we might want to change? Maybe at the end of the header comment?

As it stands, it would be very easy for us to land this, move onto other things for a month, and lose the reasoning why the strategy is this way.

clangd/TUScheduler.cpp

10

nit: store -> create? storage is a bit more complicated, but doesn't need to clutter the opening sentence.

12

nit: blank line after "...FIFO order"? separating "this is the hierarchy" from "this is how we manage lifetimes".

In fact, I'd consider moving the whole lifetime explanation down and merging it with the ASTWorker class comment. It's more detail than strategy.

13

is shared by?

14

I can't really follow these two sentences. I think we may want to give a little motivation (it's not clear why the destructor needs to not block) before introducing ASTWorkerHandle. e.g.

The TUScheduler should discard an ASTWorker when remove() is called, but its thread may
be busy and we don't want to block. So the workers are accessed via an ASTWorkerHandle.
// Destroying the handle signals the worker to exit its run loop, gives up shared ownership of the worker, and detaches the thread.

26

any need for this constructor?

29

This could use a comment and/or a more descriptive variable name to distinguish it from Worker.FileInputs.

34

nit: just "Notify all workers that they need to stop", rather than echoing the code

38

Hmm, not sure it's much safer than putting a comment on a member, but up to you.

If we're not using the destructor-blocking behavior, remove the destructor or replace with an assert?

44–45

Maybe clearer than STL it->second stuff:

FileData &FD = Files[File]
if (!FD.Worker)
  FD.Worker = Create(...);
FD.Inputs = move(Inputs)

45–61

this looks like a second copy of Inputs, we should only need one

60

ah, this is my confusion - i mixed up "current state of File" with "current state of the file", which are almost opposites!
Do you think it's confusing to call the File member AST? All the callsites actually seem to make sense with that name.

62

Not totally obvious in this class what the mutex is for - consider a comment.
e.g. // Guards members used by both TUScheduler and the worker thread.

63

processing.

99

why do we want to make this copy rather than *only* getting the preamble within the task?
Is it to do with preferring preambles from the past, but accepting those from the future?
If I'm understanding right, the general case is these versions:

• V1: this preamble has finished building before runWithPreamble is called
• V2: this update() came before runWithPreamble, but wasn't built yet. It finished building before the read task ran.
• V3: this update() came before runWithPreamble, but wasn't built by the time the read task ran.
• V4: this update() came after runWithPreamble was called, but finished before the read task ran (semaphore race?)
• V5: this came after runWithPreamble and finished after the read task

Ideally I guess we'd prefer V2, falling back to V1 -> V4 -> null in that order.
Just calling getPossiblyStalePreamble() in the task would return V4 -> V2 -> V1 -> null.
The current code will return V1 -> V4 -> V2.

But we have to get really unlucky for V4 to exist, whereas V2 is quite likely to exist I think. This means that only calling getPossiblyStalePreamble() once, in the task, seems more likely to get this right to me.

114

no-op.
(If it's just an initlist, inline?)

137

This deserves a comment for why we don't check isCancelled again... (concretely I guess, why we need it to actually get diagnostics).

140

It'd separate the work from the mechanism a bit more to pull out:

startTask(task, isUpdate, args...):
  if runSync
    task(args...)
  else
    grab context, bind args, push request, notify

it looks like all cases fit?

164

why under lock? I thought this is the only thread that can access?

(if this locking is a no-op pending cleanup, fixme)

195

I don't think this works - other accesses of File aren't guarded by the mutex?
(And File is documented as only being used by the worker thread)

One option is to cache this in a member after each rebuild(), and put that member behind the mutex. It won't catch ASTs that grow over time, though :-(
A fundamental question is what this function should do if called while the AST is being rebuilt on the worker thread.

214

So when Done && !Requests.empty(), we complete the queue even though the file is removed.
This seems like the right behavior, but deserves a line-comment.

229

nit: inlining these trivial methods might aid readability

clangd/TUScheduler.h

91

nit: this name is pervasive, but it stopped me understanding what these are - PCHOps?

93

Can we make this an optional<AsyncTaskRunner>, empty if we're running synchronously?
That way we'll get an assert if we try to spawn threads in the wrong mode.
I guess RunSync can go away then.

(And Worker::Create() takes a pointer I guess)

clangd/Threading.h

27

hmm, I'm not sure there's enough value in this FIXME to make it ever worth doing. I think it will live forever and we should just drop it :-)

67

Yep, the testing use case makes sense. Let's replace the dtor with an assert though, if we don't think owning this object is a safe way to guard.

[ilya-biryukov]

• Renamed File to AST.
• Introduced startTask().
• Moved small methods of ASTWorkerHandle to have inline definitions.
• Removed constructor of FileData.
• Replaced find() with [].
• Removed RunSync, store optional<TasksRunner> instead.
• Added a test for AsyncTasksRunner.
• Documented approach for cancelling updates.
• Addressed other review comments.

clangd/ASTWorker.cpp

103

Added a section about cancellation to the comments in the .cpp file.

clangd/TUScheduler.cpp

14

Added your comment.

26

It allowed us to use llvm::make_unique.
Remove it and replaced with unique_ptr<FileData>(new FileData{}).

29

Added a comment. Worker.FileInputs are not exposed in the public API (for that specific reason), so hopefully shouldn't cause any confusion

38

I had trouble figuring out what went wrong more than once when hitting similar errors.
The worst part is that it usually pops up when editing the code (and missing the comment), not in the initial patch.
I'd leave as is.

I also think that an extra waitForAll() in destructor doesn't hurt and makes the interface of the class saner, i.e. makes it easier to use such a class as a local variable.
There is a new test that checks the destructor waits, too. Hope that's alright.

44–45

Thanks. Looks much better.

60

SG, renamed.

99

Good point.
The original intention was to pass both preambles to the actions and make them choose which one matches for their needs.
Given that we now cheat and use preambles anyway without checking if they match, not storing the copy seems the way to go.

164

That's the only way to access the AST with the current CppFile interface, I'll do a cleanup with the follow-up patch.
Added a comment

195

It work because CppFile is thread-safe at the moment. We don't even need the lock.
I'll do the associated cleanup with the follow-up patch.

One option is to cache this in a member after each rebuild(), and put that member behind the mutex. It won't catch ASTs that grow over time, though :-(

That was the plan for the cleanup of CppFile.
We could update this member after each read, too. This would give us a good estimate after AST grows.

A fundamental question is what this function should do if called while the AST is being rebuilt on the worker thread.

We only report estimated memory usage and it's only accurate when there are no active requests.
So the answer is: we report memory used by the previous AST while doing the rebuild and memory used by the new AST after rebuild is finished.
This seems good enough, we don't need the perfect estimates AFAIK.

clangd/Threading.h

27

Oh, I think there is if we ever want to use it as a more general cancellation mechanism.
Splitting the responsibilities of "checking for cancellation" and "cancelling the request" seems like a very useful thing to do.

Removed the FIXME, we can discuss that in more detail if CancellationFlag gets more widely used.

67

I don't think owning the object is a safe way to guard, but waiting in destructor still seems like a nicer interface than assert (see the other comment).
Added a test for that.

[sammccall]

This is really great. Just one test nit.

unittests/clangd/ThreadingTests.cpp

34 ↗

(On Diff #132967)

The current test passes if runAsync is synchronous.
I'd suggest scheduling while holding a mutex that prevents the tasks from getting to the point of incrementing the counter, (and maybe checking that the counter is 0 after scheduling for clarity).

[ilya-biryukov]

• Addressed the last review comment

clangd/ASTWorker.h

+//===--- ASTWorker.h ----------------------------------------------*-C++-*-===//
+//
+//                     The LLVM Compiler Infrastructure
+//
+// This file is distributed under the University of Illinois Open Source
+// License. See LICENSE.TXT for details.
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef LLVM_CLANG_TOOLS_EXTRA_CLANGD_ASTWORKER_H
+#define LLVM_CLANG_TOOLS_EXTRA_CLANGD_ASTWORKER_H
+
+#include "ClangdUnit.h"
+#include "Threading.h"
+#include <memory>
+#include <queue>
+
+namespace clang {
+namespace clangd {
+struct InputsAndAST {
+  const ParseInputs &Inputs;
+  ParsedAST &AST;
+};
+
+struct InputsAndPreamble {
+  const ParseInputs &Inputs;
+  const PreambleData *Preamble;
+};
+
+class ASTWorkerHandle;
+
+/// Owns one instance of the AST, schedules updated and reads of it.
+/// Also responsible for building and providing access to the preamble.
+/// Each ASTWorker processes the async requests sent to it on a separate
+/// dedicated thread.
+/// The clients can access this ASTWorker only via the smart-pointer-like
+/// ASTWorkerHandler.
+class ASTWorker {
+  friend class ASTWorkerHandle;
+  ASTWorker(Semaphore &Barrier, std::shared_ptr<CppFile> File, bool RunSync);
+
+public:
+  /// Create a new ASTWorker and return a handle to it.
+  /// The processing thread is spawned using \p Tasks. However, when \p RunSync
+  /// is true, all requests will be processed on the calling thread
+  /// synchronously instead. \p Barrier is acquired when processing each
+  /// request, it is be used to limit the number of actively running threads.
+  static ASTWorkerHandle Create(AsyncTaskRunner &Tasks, Semaphore &Barrier,
+                                std::shared_ptr<CppFile> File, bool RunSync);
+  ~ASTWorker();
+
+  std::shared_ptr<const PreambleData> getPossiblyStalePreamble() const;
+  std::size_t getUsedBytes() const;
+
+  void update(ParseInputs Inputs,
+              UniqueFunction<void(llvm::Optional<std::vector<DiagWithFixIts>>)>
+                  OnUpdated);
+  void runWithAST(UniqueFunction<void(llvm::Expected<InputsAndAST>)> Action);
+
+private:
+  /// Signal that run() should finish processing pending requests and exit.
+  void setDone();
+  // Must be called exactly once on processing thread. Will return after
+  // setDone() is called on a separate thread and all pending requests are
+  // processed.
+  void run();
+
+  using RequestWithCtx = std::pair<UniqueFunction<void()>, Context>;
+
+  const bool RunSync;
+  Semaphore &Barrier;
+  // File and FileInputs are only accessed on the processing thread from run().
+  // FIXME(ibiryukov): group CppFile and FileInputs into a separate class.
+  const std::shared_ptr<CppFile> File;
+  // Inputs, corresponding to the current state of File.
+  ParseInputs FileInputs;
+  mutable std::mutex Mutex;
+  // Set to true to signal run() to finish processin
+  bool Done;                           /* GUARDED_BY(Mutex) */
+  std::queue<RequestWithCtx> Requests; /* GUARDED_BY(Mutex) */
+  // Only set when last request is an update.
+  llvm::Optional<CancellationFlag> LastUpdateCF; /* GUARDED_BY(Mutex) */
+  std::condition_variable RequestsCV;
+};
+
+/// A smart-pointer-like class that points to an active ASTWorker.
+/// In destructor, signals to the underlying ASTWorker that no new requests will
+/// be sent and the processing loop may exit (after running all pending
+/// requests).
+class ASTWorkerHandle {
+  friend class ASTWorker;
+  ASTWorkerHandle(std::shared_ptr<ASTWorker> Worker);
+
+public:
+  ASTWorkerHandle(const ASTWorkerHandle &) = delete;
+  ASTWorkerHandle &operator=(const ASTWorkerHandle &) = delete;
+  ASTWorkerHandle(ASTWorkerHandle &&) = default;
+  ASTWorkerHandle &operator=(ASTWorkerHandle &&) = default;
+
+  ~ASTWorkerHandle();
+
+  ASTWorker &operator*();
+  ASTWorker *operator->();
+
+  /// Returns an owning reference to the underlying ASTWorker that can outlive
+  /// the ASTWorkerHandle. However, no new requests to an active ASTWorker can
+  /// be schedule via the returned reference, i.e. only reads of the preamble
+  /// are possible.
+  std::shared_ptr<const ASTWorker> lock();
+
+private:
+  std::shared_ptr<ASTWorker> Worker;
+};
+
+} // namespace clangd
+} // namespace clang
+
+#endif

clangd/ASTWorker.cpp

+//===--- ASTWorker.cpp --------------------------------------------*-C++-*-===//
+//
+//                     The LLVM Compiler Infrastructure
+//
+// This file is distributed under the University of Illinois Open Source
+// License. See LICENSE.TXT for details.
+//
+//===----------------------------------------------------------------------===//
+
+#include "ASTWorker.h"
+#include "llvm/Support/Errc.h"
+
+namespace clang {
+namespace clangd {
+
+ASTWorkerHandle ASTWorker::Create(AsyncTaskRunner &Tasks, Semaphore &Barrier,
+                                  std::shared_ptr<CppFile> File, bool RunSync) {
+  std::shared_ptr<ASTWorker> Worker(
+      new ASTWorker(Barrier, std::move(File), RunSync));
+  if (!RunSync)
+    Tasks.runAsync([Worker]() { Worker->run(); });
+
+  return ASTWorkerHandle(std::move(Worker));
+}
+
+ASTWorker::ASTWorker(Semaphore &Barrier, std::shared_ptr<CppFile> File,
+                     bool RunSync)
+    : RunSync(RunSync), Barrier(Barrier), File(std::move(File)), Done(false) {
+  if (RunSync)
+    return;
+}
+
+ASTWorker::~ASTWorker() {
+#ifndef NDEBUG
+  std::lock_guard<std::mutex> Lock(Mutex);
+  assert(Done && "handle was not destroyed");
+  assert(Requests.empty() && "unprocessed requests when destroying ASTWorker");
+#endif
+}
+
+std::shared_ptr<const PreambleData>
+ASTWorker::getPossiblyStalePreamble() const {
+  return File->getPossiblyStalePreamble();
+}
+
+std::size_t ASTWorker::getUsedBytes() const {
+  std::lock_guard<std::mutex> Lock(Mutex);
+  return File->getUsedBytes();
+}
+
+void ASTWorker::update(
+    ParseInputs Inputs,
+    UniqueFunction<void(llvm::Optional<std::vector<DiagWithFixIts>>)>
+        OnUpdated) {
+  auto Task = [=](CancellationFlag CF, decltype(OnUpdated) OnUpdated) mutable {
+    if (CF.isCancelled()) {
+      OnUpdated(llvm::None);
+      return;
+    }
+    FileInputs = Inputs;
+    auto Diags = File->rebuild(std::move(Inputs));
+    OnUpdated(std::move(Diags));
+  };
+
+  if (RunSync) {
+    Task(CancellationFlag(), std::move(OnUpdated));
+    return;
+  }
+
+  {
+    std::lock_guard<std::mutex> Lock(Mutex);
+    assert(!Done && "update() after setDone()");
+    if (!Requests.empty() && LastUpdateCF) {
+      // There were no reads for the last unprocessed update, let's cancel it to
+      // not waste time on it.
+      LastUpdateCF->cancel();
+    }
+    LastUpdateCF = CancellationFlag();
+    Requests.emplace(BindWithForward(Task, *LastUpdateCF, std::move(OnUpdated)),
+                     Context::current().clone());
+  }
+  RequestsCV.notify_one();
+}
+
+void ASTWorker::runWithAST(
+    UniqueFunction<void(llvm::Expected<InputsAndAST>)> Action) {
+  auto Task = [=](decltype(Action) Action) {
+    auto AST = File->getAST().get();
+    AST->runUnderLock([&](ParsedAST *AST) {
+      if (!AST) {
+        Action(llvm::make_error<llvm::StringError>(
+            "invalid AST", llvm::errc::invalid_argument));
+        return;
+      }
+      Action(InputsAndAST{FileInputs, *AST});
+    });
+  };
+
+  if (RunSync) {
+    Task(std::move(Action));
+    return;
+  }
+
+  {
+    std::lock_guard<std::mutex> Lock(Mutex);
+    assert(!Done && "runWithAST() after setDone()");
+    LastUpdateCF = llvm::None;
+    Requests.emplace(BindWithForward(Task, std::move(Action)),
+                     Context::current().clone());
+  }
+  RequestsCV.notify_one();
+}
+
+void ASTWorker::run() {
+  while (true) {
+    RequestWithCtx Req;
+    {
+      std::unique_lock<std::mutex> Lock(Mutex);
+      RequestsCV.wait(Lock, [&]() { return Done || !Requests.empty(); });
+      if (Requests.empty()) {
+        assert(Done);
+        return;
+      }
+
+      Req = std::move(Requests.front());
+      Requests.pop();
+    } // unlock Mutex
+
+    std::lock_guard<Semaphore> BarrierLock(Barrier);
+    WithContext Guard(std::move(Req.second));
+    Req.first();
+  }
+}
+
+void ASTWorker::setDone() {
+  {
+    std::lock_guard<std::mutex> Lock(Mutex);
+    assert(!Done && "setDone() called twice");
+    Done = true;
+  }
+  RequestsCV.notify_one();
+}
+
+ASTWorkerHandle::ASTWorkerHandle(std::shared_ptr<ASTWorker> Worker)
+    : Worker(std::move(Worker)) {
+  assert(this->Worker);
+}
+
+ASTWorkerHandle::~ASTWorkerHandle() {
+  if (Worker)
+    Worker->setDone();
+}
+
+ASTWorker &ASTWorkerHandle::operator*() {
+  assert(Worker && "Handle was moved from");
+  return *Worker;
+}
+
+ASTWorker *ASTWorkerHandle::operator->() {
+  assert(Worker && "Handle was moved from");
+  return Worker.get();
+}
+
+std::shared_ptr<const ASTWorker> ASTWorkerHandle::lock() { return Worker; }
+
+} // namespace clangd
+} // namespace clang

clangd/CMakeLists.txt

 set(LLVM_LINK_COMPONENTS
   Support
   )
 
 add_clang_library(clangDaemon
+  ASTWorker.cpp
   ClangdLSPServer.cpp
   ClangdServer.cpp
   ClangdUnit.cpp
-  ClangdUnitStore.cpp
   CodeComplete.cpp
   CodeCompletionStrings.cpp
   CompileArgsCache.cpp
   Compiler.cpp
   Context.cpp
   DraftStore.cpp
   FuzzyMatch.cpp
   GlobalCompilationDatabase.cpp

clangd/ClangdServer.h

 //===--- ClangdServer.h - Main clangd server code ----------------*- C++-*-===//
 //
 //                     The LLVM Compiler Infrastructure
 //
 // This file is distributed under the University of Illinois Open Source
 // License. See LICENSE.TXT for details.
 //
 //===----------------------------------------------------------------------===//
 
 #ifndef LLVM_CLANG_TOOLS_EXTRA_CLANGD_CLANGDSERVER_H
 #define LLVM_CLANG_TOOLS_EXTRA_CLANGD_CLANGDSERVER_H
 
 #include "ClangdUnit.h"
-#include "ClangdUnitStore.h"
 #include "CodeComplete.h"
 #include "CompileArgsCache.h"
 #include "DraftStore.h"
 #include "Function.h"
 #include "GlobalCompilationDatabase.h"
 #include "Protocol.h"
 #include "TUScheduler.h"
 #include "index/FileIndex.h"

 };
 
 class RealFileSystemProvider : public FileSystemProvider {
 public:
   /// \return getRealFileSystem() tagged with default tag, i.e. VFSTag()
   Tagged<IntrusiveRefCntPtr<vfs::FileSystem>>
   getTaggedFileSystem(PathRef File) override;
 };
 
 /// Provides API to manage ASTs for a collection of C++ files and request
 /// various language features.
 /// Currently supports async diagnostics, code completion, formatting and goto
 /// definition.
 class ClangdServer {
 public:
   /// Creates a new ClangdServer instance.
   /// To process parsing requests asynchronously, ClangdServer will spawn \p
   /// AsyncThreadsCount worker threads. However, if \p AsyncThreadsCount is 0,
   /// all requests will be processed on the calling thread.
   ///
   /// ClangdServer uses \p FSProvider to get an instance of vfs::FileSystem for
   /// each parsing request. Results of code completion and diagnostics also
public:
   ///
   /// After each parsing request finishes, ClangdServer reports diagnostics to
   /// \p DiagConsumer. Note that a callback to \p DiagConsumer happens on a
   /// worker thread. Therefore, instances of \p DiagConsumer must properly
   /// synchronize access to shared state.
   ///
   /// \p StorePreamblesInMemory defines whether the Preambles generated by
   /// clangd are stored in-memory or on disk.
   ///
   /// If \p BuildDynamicSymbolIndex is true, ClangdServer builds a dynamic
   /// in-memory index for symbols in all opened files and uses the index to
   /// augment code completion results.
   ///
   /// If \p StaticIdx is set, ClangdServer uses the index for global code
   /// completion.
   ClangdServer(GlobalCompilationDatabase &CDB,
                DiagnosticsConsumer &DiagConsumer,
                FileSystemProvider &FSProvider, unsigned AsyncThreadsCount,

clangd/ClangdUnitStore.h

-//===--- ClangdUnitStore.h - A container of CppFiles -------------*-C++-*-===//
-//
-//                     The LLVM Compiler Infrastructure
-//
-// This file is distributed under the University of Illinois Open Source
-// License. See LICENSE.TXT for details.
-//
-//===---------------------------------------------------------------------===//
-
-#ifndef LLVM_CLANG_TOOLS_EXTRA_CLANGD_CLANGDUNITSTORE_H
-#define LLVM_CLANG_TOOLS_EXTRA_CLANGD_CLANGDUNITSTORE_H
-
-#include "ClangdUnit.h"
-#include "GlobalCompilationDatabase.h"
-#include "Logger.h"
-#include "Path.h"
-#include "clang/Tooling/CompilationDatabase.h"
-#include <mutex>
-
-namespace clang {
-namespace clangd {
-
-class Logger;
-
-/// Thread-safe mapping from FileNames to CppFile.
-class CppFileCollection {
-public:
-  /// \p ASTCallback is called when a file is parsed synchronously. This should
-  /// not be expensive since it blocks diagnostics.
-  explicit CppFileCollection(bool StorePreamblesInMemory,
-                             std::shared_ptr<PCHContainerOperations> PCHs,
-                             ASTParsedCallback ASTCallback)
-      : ASTCallback(std::move(ASTCallback)), PCHs(std::move(PCHs)),
-        StorePreamblesInMemory(StorePreamblesInMemory) {}
-
-  std::shared_ptr<CppFile> getOrCreateFile(PathRef File) {
-    std::lock_guard<std::mutex> Lock(Mutex);
-    auto It = OpenedFiles.find(File);
-    if (It == OpenedFiles.end()) {
-      It = OpenedFiles
-               .try_emplace(File, CppFile::Create(File, StorePreamblesInMemory,
-                                                  PCHs, ASTCallback))
-               .first;
-    }
-    return It->second;
-  }
-
-  std::shared_ptr<CppFile> getFile(PathRef File) const {
-    std::lock_guard<std::mutex> Lock(Mutex);
-    auto It = OpenedFiles.find(File);
-    if (It == OpenedFiles.end())
-      return nullptr;
-    return It->second;
-  }
-
-  /// Removes a CppFile, stored for \p File, if it's inside collection and
-  /// returns it.
-  std::shared_ptr<CppFile> removeIfPresent(PathRef File);
-
-  /// Gets used memory for each of the stored files.
-  std::vector<std::pair<Path, std::size_t>> getUsedBytesPerFile() const;
-
-private:
-  mutable std::mutex Mutex;
-  llvm::StringMap<std::shared_ptr<CppFile>> OpenedFiles;
-  ASTParsedCallback ASTCallback;
-  std::shared_ptr<PCHContainerOperations> PCHs;
-  bool StorePreamblesInMemory;
-};
-} // namespace clangd
-} // namespace clang
-
-#endif

clangd/ClangdUnitStore.cpp

-//===--- ClangdUnitStore.cpp - A ClangdUnits container -----------*-C++-*-===//
-//
-//                     The LLVM Compiler Infrastructure
-//
-// This file is distributed under the University of Illinois Open Source
-// License. See LICENSE.TXT for details.
-//
-//===----------------------------------------------------------------------===//
-
-#include "ClangdUnitStore.h"
-#include "llvm/Support/Path.h"
-#include <algorithm>
-
-using namespace clang::clangd;
-using namespace clang;
-
-std::shared_ptr<CppFile> CppFileCollection::removeIfPresent(PathRef File) {
-  std::lock_guard<std::mutex> Lock(Mutex);
-
-  auto It = OpenedFiles.find(File);
-  if (It == OpenedFiles.end())
-    return nullptr;
-
-  std::shared_ptr<CppFile> Result = It->second;
-  OpenedFiles.erase(It);
-  return Result;
-}
-std::vector<std::pair<Path, std::size_t>>
-CppFileCollection::getUsedBytesPerFile() const {
-  std::lock_guard<std::mutex> Lock(Mutex);
-  std::vector<std::pair<Path, std::size_t>> Result;
-  Result.reserve(OpenedFiles.size());
-  for (auto &&PathAndFile : OpenedFiles)
-    Result.push_back(
-        {PathAndFile.first().str(), PathAndFile.second->getUsedBytes()});
-  return Result;
-}

clangd/TUScheduler.h

 //===--- TUScheduler.h -------------------------------------------*-C++-*-===//
 //
 //                     The LLVM Compiler Infrastructure
 //
 // This file is distributed under the University of Illinois Open Source
 // License. See LICENSE.TXT for details.
 //
 //===----------------------------------------------------------------------===//
 
 #ifndef LLVM_CLANG_TOOLS_EXTRA_CLANGD_TUSCHEDULER_H
 #define LLVM_CLANG_TOOLS_EXTRA_CLANGD_TUSCHEDULER_H
 
-#include "ClangdUnit.h"
-#include "ClangdUnitStore.h"
+#include "ASTWorker.h"
 #include "Function.h"
 #include "Threading.h"
+#include "llvm/ADT/StringMap.h"
 
 namespace clang {
 namespace clangd {
 /// Returns a number of a default async threads to use for TUScheduler.
 /// Returned value is always >= 1 (i.e. will not cause requests to be processed
 /// synchronously).
 unsigned getDefaultAsyncThreadsCount();
 
-struct InputsAndAST {
-  const ParseInputs &Inputs;
-  ParsedAST *
-};
-
-struct InputsAndPreamble {
-  const ParseInputs &Inputs;
-  const PreambleData *Preamble;
-};
-
 /// Handles running tasks for ClangdServer and managing the resources (e.g.,
 /// preambles and ASTs) for opened files.
 /// TUScheduler is not thread-safe, only one thread should be providing updates
 /// and scheduling tasks.
 /// Callbacks are run on a threadpool and it's appropriate to do slow work in
 /// them.
 class TUScheduler {
 public:
   TUScheduler(unsigned AsyncThreadsCount, bool StorePreamblesInMemory,
               ASTParsedCallback ASTCallback);
+  ~TUScheduler();
 
   /// Returns estimated memory usage for each of the currently open files.
   /// The order of results is unspecified.
   std::vector<std::pair<Path, std::size_t>> getUsedBytesPerFile() const;
 
   /// Schedule an update for \p File. Adds \p File to a list of tracked files if
   /// \p File was not part of it before.
   /// FIXME(ibiryukov): remove the callback from this function.
   void update(PathRef File, ParseInputs Inputs,
               UniqueFunction<void(llvm::Optional<std::vector<DiagWithFixIts>>)>
                   OnUpdated);
 
   /// Remove \p File from the list of tracked files and schedule removal of its
   /// resources. \p Action will be called when resources are freed.
   /// If an error occurs during processing, it is forwarded to the \p Action
   /// callback.
   /// FIXME(ibiryukov): the callback passed to this function is not used, we
   /// should remove it.
   void remove(PathRef File, UniqueFunction<void(llvm::Error)> Action);
 
   /// Schedule an async read of the AST. \p Action will be called when AST is
   /// ready. The AST passed to \p Action refers to the version of \p File
   /// tracked at the time of the call, even if new updates are received before
   /// \p Action is executed.
   /// If an error occurs during processing, it is forwarded to the \p Action
   /// callback.
   void runWithAST(PathRef File,
                   UniqueFunction<void(llvm::Expected<InputsAndAST>)> Action);
 
   /// Schedule an async read of the Preamble. Preamble passed to \p Action may
   /// be built for any version of the file, callers must not rely on it being
   /// consistent with the current version of the file.
   /// If an error occurs during processing, it is forwarded to the \p Action
   /// callback.
   void runWithPreamble(
       PathRef File,
       UniqueFunction<void(llvm::Expected<InputsAndPreamble>)> Action);
 
 private:
-  const ParseInputs &getInputs(PathRef File);
+  /// This class stores per-file data in the Files map.
+  struct FileData {
+    FileData(ParseInputs Inputs, ASTWorkerHandle Worker);
+
+    ParseInputs Inputs;
+    ASTWorkerHandle Worker;
+  };
+
+  struct AuxData {
+    bool StorePreamblesInMemory;
+    std::shared_ptr<PCHContainerOperations> PCHs;
+    ASTParsedCallback ASTCallback;
+  };
 
-  llvm::StringMap<ParseInputs> CachedInputs;
-  CppFileCollection Files;
-  ThreadPool Threads;
+  const bool RunSync;
+  const AuxData Data;
+  Semaphore Barrier;
+  llvm::StringMap<std::unique_ptr<FileData>> Files;
+  AsyncTaskRunner Tasks;
 };
 } // namespace clangd
 } // namespace clang
 
 #endif

clangd/TUScheduler.cpp

+//===--- TUScheduler.cpp -----------------------------------------*-C++-*-===//
+//
+//                     The LLVM Compiler Infrastructure
+//
+// This file is distributed under the University of Illinois Open Source
+// License. See LICENSE.TXT for details.
+//
+//===----------------------------------------------------------------------===//
 #include "TUScheduler.h"
 #include "clang/Frontend/PCHContainerOperations.h"
 #include "llvm/Support/Errc.h"
 
 namespace clang {
 namespace clangd {
 unsigned getDefaultAsyncThreadsCount() {
   unsigned HardwareConcurrency = std::thread::hardware_concurrency();
   // C++ standard says that hardware_concurrency()
   // may return 0, fallback to 1 worker thread in
   // that case.
   if (HardwareConcurrency == 0)
     return 1;
   return HardwareConcurrency;
 }
 
 TUScheduler::TUScheduler(unsigned AsyncThreadsCount,
                          bool StorePreamblesInMemory,
                          ASTParsedCallback ASTCallback)
+    : RunSync(AsyncThreadsCount == 0),
+      Data{StorePreamblesInMemory, std::make_shared<PCHContainerOperations>(),
+           std::move(ASTCallback)},
+      Barrier(AsyncThreadsCount) {}
+
+TUScheduler::~TUScheduler() {
+  // Clear all FileData objects to notify all workers that they need to stop.
+  Files.clear();
 
-    : Files(StorePreamblesInMemory, std::make_shared<PCHContainerOperations>(),
-            std::move(ASTCallback)),
-      Threads(AsyncThreadsCount) {}
+  // Wait for all in-flight tasks to finish.
+  Tasks.waitForAll();
+}
 
 void TUScheduler::update(
     PathRef File, ParseInputs Inputs,
     UniqueFunction<void(llvm::Optional<std::vector<DiagWithFixIts>>)>
         OnUpdated) {
-  CachedInputs[File] = Inputs;
-
-  auto Resources = Files.getOrCreateFile(File);
-  auto DeferredRebuild = Resources->deferRebuild(std::move(Inputs));
+  auto It = Files.find(File);
+  if (It == Files.end()) {
+    // Create a new worker to process the AST-related tasks.
+    ASTWorkerHandle Worker =
+        ASTWorker::Create(Tasks, Barrier,
+                          CppFile::Create(File, Data.StorePreamblesInMemory,
+                                          Data.PCHs, Data.ASTCallback),
+                          RunSync);
+    It = Files
+             .try_emplace(
+                 File, llvm::make_unique<FileData>(Inputs, std::move(Worker)))
+             .first;
+  } else {
+    It->second->Inputs = Inputs;
+  }
 
-  Threads.addToFront(
-      [](decltype(OnUpdated) OnUpdated,
-         decltype(DeferredRebuild) DeferredRebuild) {
-        auto Diags = DeferredRebuild();
-        OnUpdated(Diags);
-      },
-      std::move(OnUpdated), std::move(DeferredRebuild));
+  It->second->Worker->update(Inputs, std::move(OnUpdated));
 }
 
 void TUScheduler::remove(PathRef File,
                          UniqueFunction<void(llvm::Error)> Action) {
-  CachedInputs.erase(File);
-
-  auto Resources = Files.removeIfPresent(File);
-  if (!Resources) {
+  auto It = Files.find(File);
+  if (It == Files.end()) {
     Action(llvm::make_error<llvm::StringError>(
         "trying to remove non-added document", llvm::errc::invalid_argument));
     return;
   }
-
-  auto DeferredCancel = Resources->deferCancelRebuild();
-  Threads.addToFront(
-      [](decltype(Action) Action, decltype(DeferredCancel) DeferredCancel) {
-        DeferredCancel();
-        Action(llvm::Error::success());
-      },
-      std::move(Action), std::move(DeferredCancel));
+  Files.erase(It);
 }
 
 void TUScheduler::runWithAST(
     PathRef File, UniqueFunction<void(llvm::Expected<InputsAndAST>)> Action) {
-  auto Resources = Files.getFile(File);
-  if (!Resources) {
+  auto It = Files.find(File);
+  if (It == Files.end()) {
     Action(llvm::make_error<llvm::StringError>(
         "trying to get AST for non-added document",
         llvm::errc::invalid_argument));
     return;
   }
 
-  const ParseInputs &Inputs = getInputs(File);
-  // We currently block the calling thread until AST is available and run the
-  // action on the calling thread to avoid inconsistent states coming from
-  // subsequent updates.
-  // FIXME(ibiryukov): this should be moved to the worker threads.
-  Resources->getAST().get()->runUnderLock([&](ParsedAST *AST) {
-    if (AST)
-      Action(InputsAndAST{Inputs, *AST});
-    else
-      Action(llvm::make_error<llvm::StringError>(
-          "Could not build AST for the latest file update",
-          llvm::errc::invalid_argument));
-  });
+  It->second->Worker->runWithAST(std::move(Action));
 }
 
 void TUScheduler::runWithPreamble(
     PathRef File,
     UniqueFunction<void(llvm::Expected<InputsAndPreamble>)> Action) {
-  std::shared_ptr<CppFile> Resources = Files.getFile(File);
-  if (!Resources) {
+  auto It = Files.find(File);
+  if (It == Files.end()) {
     Action(llvm::make_error<llvm::StringError>(
         "trying to get preamble for non-added document",
         llvm::errc::invalid_argument));
     return;
   }
 
-  const ParseInputs &Inputs = getInputs(File);
+  ParseInputs InputsCopy = It->second->Inputs;
   std::shared_ptr<const PreambleData> Preamble =
-      Resources->getPossiblyStalePreamble();
-  Threads.addToFront(
-      [Resources, Preamble, Inputs](decltype(Action) Action) mutable {
-        if (!Preamble)
-          Preamble = Resources->getPossiblyStalePreamble();
+      It->second->Worker->getPossiblyStalePreamble();
 
-        Action(InputsAndPreamble{Inputs, Preamble.get()});
-      },
-      std::move(Action));
+  if (RunSync) {
+    Action(InputsAndPreamble{InputsCopy, Preamble.get()});
+    return;
   }
 
-const ParseInputs &TUScheduler::getInputs(PathRef File) {
-  auto It = CachedInputs.find(File);
-  assert(It != CachedInputs.end());
-  return It->second;
+  std::shared_ptr<const ASTWorker> Worker = It->second->Worker.lock();
+  auto Task = [InputsCopy, Preamble, Worker,
+               this](Context Ctx, decltype(Action) Action) mutable {
+    std::lock_guard<Semaphore> BarrierLock(Barrier);
+    WithContext Guard(std::move(Ctx));
+    if (!Preamble)
+      Preamble = Worker->getPossiblyStalePreamble();
+    Action(InputsAndPreamble{InputsCopy, Preamble.get()});
+  };
+
+  Tasks.runAsync(
+      BindWithForward(Task, Context::current().clone(), std::move(Action)));
 }
 
 std::vector<std::pair<Path, std::size_t>>
 TUScheduler::getUsedBytesPerFile() const {
-  return Files.getUsedBytesPerFile();
+  std::vector<std::pair<Path, std::size_t>> Result;
+  Result.reserve(Files.size());
+  for (auto &&PathAndFile : Files)
+    Result.push_back(
+        {PathAndFile.first(), PathAndFile.second->Worker->getUsedBytes()});
+  return Result;
 }
+
+TUScheduler::FileData::FileData(ParseInputs Inputs, ASTWorkerHandle Worker)
+    : Inputs(std::move(Inputs)), Worker(std::move(Worker)) {}
+
 } // namespace clangd
 } // namespace clang

clangd/Threading.h

 //===--- ThreadPool.h --------------------------------------------*- C++-*-===//
 //
 //                     The LLVM Compiler Infrastructure
 //
 // This file is distributed under the University of Illinois Open Source
 // License. See LICENSE.TXT for details.
 //
 //===----------------------------------------------------------------------===//
 
 #ifndef LLVM_CLANG_TOOLS_EXTRA_CLANGD_THREADING_H
 #define LLVM_CLANG_TOOLS_EXTRA_CLANGD_THREADING_H
 
 #include "Context.h"
 #include "Function.h"
+#include <atomic>
+#include <cassert>
 #include <condition_variable>
-#include <deque>
+#include <memory>
 #include <mutex>
-#include <thread>
 #include <vector>
 
 namespace clang {
 namespace clangd {
-/// A simple fixed-size thread pool implementation.
-class ThreadPool {
+
+/// A shared boolean flag indicating if the computation was cancelled.
+/// Once cancelled, cannot be returned to the previous state.
+/// FIXME: We should split this class it into consumers and producers of the
+/// cancellation flags.
+class CancellationFlag {
 public:
-  /// If \p AsyncThreadsCount is 0, requests added using addToFront and addToEnd
-  /// will be processed synchronously on the calling thread.
-  // Otherwise, \p AsyncThreadsCount threads will be created to schedule the
-  // requests.
-  ThreadPool(unsigned AsyncThreadsCount);
-  /// Destructor blocks until all requests are processed and worker threads are
-  /// terminated.
-  ~ThreadPool();
-
-  /// Add a new request to run function \p F with args \p As to the start of the
-  /// queue. The request will be run on a separate thread.
-  template <class Func, class... Args>
-  void addToFront(Func &&F, Args &&... As) {
-    if (RunSynchronously) {
-      std::forward<Func>(F)(std::forward<Args>(As)...);
-      return;
-    }
+  CancellationFlag();
 
-    {
-      std::lock_guard<std::mutex> Lock(Mutex);
-      RequestQueue.emplace_front(
-          BindWithForward(std::forward<Func>(F), std::forward<Args>(As)...),
-          Context::current().clone());
-    }
-    RequestCV.notify_one();
+  void cancel() {
+    assert(WasCancelled && "the object was moved");
+    WasCancelled->store(true);
   }
 
-  /// Add a new request to run function \p F with args \p As to the end of the
-  /// queue. The request will be run on a separate thread.
-  template <class Func, class... Args> void addToEnd(Func &&F, Args &&... As) {
-    if (RunSynchronously) {
-      std::forward<Func>(F)(std::forward<Args>(As)...);
-      return;
-    }
+  bool isCancelled() const {
+    assert(WasCancelled && "the object was moved");
+    return WasCancelled->load();
 
-    {
-      std::lock_guard<std::mutex> Lock(Mutex);
-      RequestQueue.emplace_back(
-          BindWithForward(std::forward<Func>(F), std::forward<Args>(As)...),
-          Context::current().clone());
-    }
-    RequestCV.notify_one();
-  }
+private:
+  std::shared_ptr<std::atomic<bool>> WasCancelled;
+};
+
+/// Limits the number of threads that can acquire the lock at the same time.
+class Semaphore {
+public:
+  Semaphore(std::size_t MaxLocks);
+
+  void lock();
+  void unlock();
+
+private:
+  std::mutex Mutex;
+  std::condition_variable SlotsChanged;
+  std::size_t FreeSlots;
+};
+
+/// Allows to run tasks on separate (detached) threads and wait for all tasks to
+/// finish.
+class AsyncTaskRunner {
+public:
+  ~AsyncTaskRunner();
+
+  void waitForAll();
+  void runAsync(UniqueFunction<void()> Action);
 
 private:
-  bool RunSynchronously;
-  mutable std::mutex Mutex;
-  /// We run some tasks on separate threads(parsing, CppFile cleanup).
-  /// These threads looks into RequestQueue to find requests to handle and
-  /// terminate when Done is set to true.
-  std::vector<std::thread> Workers;
-  /// Setting Done to true will make the worker threads terminate.
-  bool Done = false;
-  /// A queue of requests.
-  std::deque<std::pair<UniqueFunction<void()>, Context>> RequestQueue;
-  /// Condition variable to wake up worker threads.
-  std::condition_variable RequestCV;
+  std::mutex Mutex;
+  std::condition_variable TasksReachedZero;
+  std::size_t InFlightTasks = 0;
 };
 } // namespace clangd
 } // namespace clang
 #endif

clangd/Threading.cpp

 #include "Threading.h"
+#include "llvm/ADT/ScopeExit.h"
 #include "llvm/Support/FormatVariadic.h"
 #include "llvm/Support/Threading.h"
+#include <thread>
 
 namespace clang {
 namespace clangd {
-ThreadPool::ThreadPool(unsigned AsyncThreadsCount)
-    : RunSynchronously(AsyncThreadsCount == 0) {
-  if (RunSynchronously) {
-    // Don't start the worker thread if we're running synchronously
-    return;
-  }
 
-  Workers.reserve(AsyncThreadsCount);
-  for (unsigned I = 0; I < AsyncThreadsCount; ++I) {
-    Workers.push_back(std::thread([this, I]() {
-      llvm::set_thread_name(llvm::formatv("scheduler/{0}", I));
-      while (true) {
-        UniqueFunction<void()> Request;
-        Context Ctx;
+CancellationFlag::CancellationFlag()
+    : WasCancelled(std::make_shared<std::atomic<bool>>(false)) {}
 
-        // Pick request from the queue
-        {
-          std::unique_lock<std::mutex> Lock(Mutex);
-          // Wait for more requests.
-          RequestCV.wait(Lock,
-                         [this] { return !RequestQueue.empty() || Done; });
-          if (RequestQueue.empty()) {
-            assert(Done);
-            return;
-          }
+Semaphore::Semaphore(std::size_t MaxLocks) : FreeSlots(MaxLocks) {}
+
+void Semaphore::lock() {
+  SlotsChanged.wait(Lock, [&]() { return FreeSlots > 0; });
+  --FreeSlots;
 
-          // We process requests starting from the front of the queue. Users of
-          // ThreadPool have a way to prioritise their requests by putting
-          // them to the either side of the queue (using either addToEnd or
-          // addToFront).
-          std::tie(Request, Ctx) = std::move(RequestQueue.front());
-          RequestQueue.pop_front();
-        } // unlock Mutex
+void Semaphore::unlock() {
+  std::unique_lock<std::mutex> Lock(Mutex);
+  ++FreeSlots;
+  Lock.unlock();
 
-        WithContext WithCtx(std::move(Ctx));
-        Request();
-      }
-    }));
-  }
+  SlotsChanged.notify_one();
 }
 
-ThreadPool::~ThreadPool() {
-  if (RunSynchronously)
-    return; // no worker thread is running in that case
+AsyncTaskRunner::~AsyncTaskRunner() { waitForAll(); }
 
+void AsyncTaskRunner::waitForAll() {
+  std::unique_lock<std::mutex> Lock(Mutex);
+  TasksReachedZero.wait(Lock, [&]() { return InFlightTasks == 0; });
+}
+
+void AsyncTaskRunner::runAsync(UniqueFunction<void()> Action) {
   {
-    std::lock_guard<std::mutex> Lock(Mutex);
-    // Wake up the worker thread
-    Done = true;
-  } // unlock Mutex
-  RequestCV.notify_all();
+    std::unique_lock<std::mutex> Lock(Mutex);
+    ++InFlightTasks;
+  }
+
+  auto CleanupTask = llvm::make_scope_exit([this]() {
+    std::unique_lock<std::mutex> Lock(Mutex);
+    int NewTasksCnt = --InFlightTasks;
+    Lock.unlock();
 
-  for (auto &Worker : Workers)
-    Worker.join();
+    if (NewTasksCnt == 0)
+      TasksReachedZero.notify_one();
+  });
+
+  std::thread(
+      [](decltype(Action) Action, decltype(CleanupTask)) {
+        Action();
+        // Make sure function stored by Action is destroyed before CleanupTask
+        // is run.
+        Action = nullptr;
+      },
+      std::move(Action), std::move(CleanupTask))
+      .detach();
 }
 } // namespace clangd
 } // namespace clang