Add support for custom query prerequisite logic #16073

mayankgarg1990 · 2021-05-11T23:09:19Z

This PR adds a mechanism using which admins can check for certain prerequisites to be satisfied before a query is queued in the Presto queues. An example of this situation is in Facebook - where given our multi-region warehouse setup - we want to ensure that data is present in the region where the query is being executed before the query starts.

The high level design is:

Add to SPI an interface QueryPrerequisites which allows for custom logic to be implemented and injected using the Plugin interface. By doing so, admins can choose to add custom logic before query execution and can ensure if some prerequisites need to be met - that can be achieved.
Add a new state WAITING_FOR_PREREQUISITES as the first stage when a query starts and as the QueryPrerequisites is complete, it will move to QUEUED. This is done to ensure a clean line between when a query is ready to run and queued vs the admin needs to do some preparation before the query can be executed.

To install a custom QueryPrerequisites implementation, a factory instance QueryPrerequisitesFactory needs to be implemented which is exposed as a Plugin to the coordinator. Then the coordinator will load the correct factory by reading the configuration file etc/query-prerequisites.properties and the property query-prerequisites.factory. If nothing is specified, a default implementation will be loaded which finished immediately - hence it is not necessary to install a custom QueryPrerequisites instance.

Test plan - Added tests to TestLocalQueryDispatcher which contains the logic for invoking and dealing with the prerequisites stage. In addition to this, on a local cluster in Facebook, developed our custom QueryPrerequisites logic and built and injected using the Plugin mechanism and ensure that the system worked

== RELEASE NOTES ==

SPI Changes
* Add support for custom query prerequisites to be checked and satisfied through ``QueryPrerequisites`` interface. See :pr:`16073`.

presto-main/src/main/java/com/facebook/presto/execution/QueryState.java

swapsmagic · 2021-05-19T17:47:16Z

presto-main/src/main/java/com/facebook/presto/execution/QueryStateTimer.java

+        beginQueued(now);
+        queuedTime.compareAndSet(null, nanosSince(beginQueuedNanos, now));


it seems this two statement makes queuedTime almost 0. Is that intentional?

This is how this class works in general. Here is another function :

private void beginDispatching(long now) { beginWaitingForResources(now); resourceWaitingTime.compareAndSet(null, nanosSince(beginResourceWaitingNanos, now)); beginDispatchingNanos.compareAndSet(null, now); }

The logic is as following:

beginQueued(now) -> beginQueued has compareAndSet - this is there for situations where in the query goes from WAITING_FOR_PREREQUISITES to FAILED - in that case this chaining of functions ensures that all the states are set. So if the query had already crossed queuing state - compareAndSet will make this function a no-op

queuedTime.compareAndSet(null, nanosSince(beginQueuedNanos, now)); - the start of WaitingForResources indicates the end of queuing and hence we are setting the queuing time now

tdcmeehan · 2021-05-20T00:44:50Z

Add new query state WAITING_FOR_PREREQUISITES LGTM, % what @swapsmagic pointed out.

presto-spi/src/main/java/com/facebook/presto/spi/dispatcher/QueryPrerequisitesFactory.java

presto-main/src/main/java/com/facebook/presto/dispatcher/DefaultQueryPrerequisites.java

tdcmeehan · 2021-05-19T23:55:27Z

presto-main/src/main/java/com/facebook/presto/dispatcher/LocalDispatchQuery.java

+                }
+            });
+
+            prerequisitesFuture.whenCompleteAsync((result, throwable) -> {


Doesn't appear necesary to use the async version?

My understanding is that if I don't use the async function - then depending on when the future finishes, that will decide the execution -

If the future is finished before the callback is set - this will execute in the thread where the callback is registered (which is fine for us)

If the future is finished after the callback is set - this will execute in the thread that completes the future - which is not desirable in my opinion - we should not load the threads that are computing the prerequisites with actual query registration stuff.

So given (2) - I went with the async option

@Test public void testFutures() throws Exception { CompletableFuture<?> a = CompletableFuture.completedFuture(null); System.out.println(1); a.whenComplete((x, y) -> {System.out.println(Thread.currentThread().getName());}); System.out.println(2); Thread foo = new Thread(() -> a.whenComplete((x, y) -> {System.out.println(Thread.currentThread().getName());}), "test-thread"); foo.start(); foo.join(); CompletableFuture<?> b = new CompletableFuture<>(); System.out.println(3); b.whenComplete((x, y) -> {System.out.println(Thread.currentThread().getName());}); System.out.println(4); Thread bar = new Thread(() -> b.complete(null), "test-thread2"); bar.start(); bar.join();

Output:

1 main 2 test-thread 3 4 test-thread2

presto-main/src/main/java/com/facebook/presto/dispatcher/DispatchManager.java

presto-main/src/main/java/com/facebook/presto/dispatcher/QueryPrerequisitesManager.java

presto-main/src/main/java/com/facebook/presto/dispatcher/DefaultQueryPrerequisites.java

ajaygeorge · 2021-05-21T07:17:13Z

presto-spi/src/main/java/com/facebook/presto/spi/prerequisites/QueryPrerequisites.java

+    default void queryFinished(QueryId queryId)
+    {
+    }


doesn't feel like the right place to add logic which can be executed after the query is finished running since this is for just for query prerequisites.

This is to inform the QueryPrerequisites that a query has finished and they can cleanup state if needed. This is useful since the query may fail before the prerequisites are satisfied (example - a user may kill their query).

Can you clarify your concerns here?

The method signature read like actions to be done post the query finished executing. Since this class is for processing prerequisites before the query is run, mixing pre and post actions in the same class was my concern.

I somehow don't have the same concern and hence unable to think of a solution. Do you have suggestions on how what I mentioned can be achieved?

Having a separate class called PostQueryActions and add methods for queryFinished there. Since this is a high pri PR don't get blocked on this. Keep it in the backlog and see if this is possible.

presto-main/src/main/resources/webapp/src/components/QueryDetail.jsx

prithvip · 2021-05-21T17:48:23Z

presto-client/src/main/java/com/facebook/presto/client/StatementStats.java

@@ -97,6 +103,12 @@ public String getState()
        return state;
    }

+    @JsonProperty
+    public boolean isWaitingForPrerequisites()


Why do we need this if we already have state?

I am maintaining symmetry in this class. We already have isQueued despite the state being there and hence added this. One of the reasons why I feel this might have been done is that state is a String and not an enum - and hence this could provide people an easy way without figuring out the strings if the query has started or not.

prithvip · 2021-05-21T17:53:06Z

presto-main/src/main/java/com/facebook/presto/dispatcher/DispatchInfo.java

    private final Duration queuedTime;

-    public static DispatchInfo queued(Duration elapsedTime, Duration queuedTime)
+    public static DispatchInfo pendingDispatch(Duration elapsedTime, Duration waitingForPrerequisitesTime, Duration queuedTime)


Perhaps we should consider combining these timings into a wrapper class? Might make it easier to add more timings like this in the future without having to refactor DispatchInfo, BasicQueryStats, LocalDispatchQuery, etc.

Adding these states is not common and hence adding a wrapper class here might be more overhead than good.

prithvip · 2021-05-21T17:58:39Z

presto-main/src/main/java/com/facebook/presto/execution/QueryStateMachine.java

+    public boolean transitionToQueued()
+    {
+        queryStateTimer.beginQueued();
+        return queryState.setIf(QUEUED, currentState -> currentState.ordinal() < QUEUED.ordinal());


I guess this idiom pre-dates your change, but it feels brittle to compare enums by ordinal this way. We should consider writing a comparator, or add a method in QueryStateMachine, that we can write unit tests against.

We do have unittests for these states. Is there a specific case that you feel should be tested and is not ?

prithvip · 2021-05-21T18:02:39Z

presto-main/src/main/java/com/facebook/presto/server/protocol/QueuedStatementResource.java

    {
-        QueryState state = queryError.map(error -> FAILED).orElse(QUEUED);
+        QueryState state = queryError.map(error -> FAILED).orElse(WAITING_FOR_PREREQUISITES);
        return new QueryResults(


We should also modify the tests in TestQueryResource to account for this change

presto-main/src/main/java/com/facebook/presto/dispatcher/LocalDispatchQuery.java

prithvip · 2021-05-21T18:09:10Z

presto-main/src/main/java/com/facebook/presto/dispatcher/QueryPrerequisitesManager.java

+        }
+    }
+
+    public void loadQueryPrerequisites()


This seems to be a lot of boilerplate that is copy-pasted from InternalResourceGroupManager. Perhaps we can consider a refactor here to avoid code duplication.

Almost every Manager is written this way - search for the managers that are used in PluginManager. I did consider doing some refactor - but then every manager has something different which makes consolidating them nuanced

prithvip · 2021-05-21T18:15:44Z

presto-spi/src/main/java/com/facebook/presto/spi/prerequisites/QueryPrerequisitesContext.java

+    private final Map<String, String> systemProperties;
+    private final Map<ConnectorId, Map<String, String>> connectorProperties;
+
+    public QueryPrerequisitesContext(


Aren't all of these already in Session? Maybe add a convenience constructor that takes Session object so that it's easier to use this class.

Session is not a part of SPI and hence I cannot use it. Even I would have preferred it instead of making a new one

prithvip · 2021-05-21T18:52:12Z

presto-main/src/main/java/com/facebook/presto/dispatcher/LocalDispatchQuery.java

+                }
+            });
+
+            prerequisitesFuture.whenCompleteAsync((result, throwable) -> {


I think we should add a timeout here to be defensive against a bad implementation that takes too long to complete. We could cancel the thread from here, rather than relying on implementations to properly enforce their own timeouts and cancel themselves.

We already have the session properties query_max_run_time which controls how long can a query stay in the system. And if that is reached - the query will fail and that will in turn cancel the future. What extra value do you feel can adding an additional timeout here add ? The implementations can always take care of this themselves by ensuring that they add a timeout on their end if they want a specific timeout

This new query state will support adding custom logic to wait for admin defined prerequistes to finish before a query can start. This commit just introduces the new state and the actual custom prerequisite functionality will be added in a subsequent commit.

Implement mechanism to allow for custom logic to be plugged in to make queries to wait till the time all of the needed prerequisites for a query have been satisfied. Admins can decide to plugin their custom prerequisite logic by implementing the `QueryPrerequisites` interface and injecting that by implementing the factory interface `QueryPrerequisitesFactory` and exposing it through a `Plugin`. If this is not needed, no action is required since a default no-op implementation will be loaded.

tdcmeehan · 2021-05-24T16:01:56Z

presto-main/src/main/java/com/facebook/presto/dispatcher/LocalDispatchQuery.java

+                queueQuery();
+            }, queryExecutor);
+        }
+        catch (Throwable t) {


What will throw here?

waitForPrerequisites is a plugin that any customer can implement so that plugin can throw

mayankgarg1990 force-pushed the prereq branch 10 times, most recently from fe7dd65 to 1e97206 Compare May 18, 2021 00:38

mayankgarg1990 mentioned this pull request May 18, 2021

[RFC] Support custom waiting logic for query in coordinator #16045

Closed

mayankgarg1990 changed the title ~~[WIP] Implement WAITING_FOR_PREREQUISITES logic in coordinator~~ [WIP] Add support for custom pre-query logic May 18, 2021

mayankgarg1990 force-pushed the prereq branch from 1e97206 to 7b6e175 Compare May 18, 2021 04:58

mayankgarg1990 changed the title ~~[WIP] Add support for custom pre-query logic~~ [WIP] Add support for custom query prerequisite logic May 18, 2021

mayankgarg1990 requested review from tdcmeehan, basarhamdionat and jeremyteboul May 18, 2021 05:01

mayankgarg1990 force-pushed the prereq branch from 7b6e175 to 5f2eadf Compare May 18, 2021 05:03

mayankgarg1990 changed the title ~~[WIP] Add support for custom query prerequisite logic~~ Add support for custom query prerequisite logic May 18, 2021

mayankgarg1990 force-pushed the prereq branch 3 times, most recently from 514fb6c to ed5cc3d Compare May 19, 2021 12:29

swapsmagic requested changes May 19, 2021

View reviewed changes

tdcmeehan requested a review from prithvip May 19, 2021 20:33

tdcmeehan reviewed May 20, 2021

View reviewed changes

mayankgarg1990 force-pushed the prereq branch 2 times, most recently from 7fbf9a3 to 53ca0b6 Compare May 20, 2021 06:55

ajaygeorge reviewed May 21, 2021

View reviewed changes

presto-main/src/main/java/com/facebook/presto/dispatcher/DefaultQueryPrerequisites.java Show resolved Hide resolved

ajaygeorge reviewed May 21, 2021

View reviewed changes

presto-main/src/main/resources/webapp/src/components/QueryDetail.jsx Show resolved Hide resolved

mayankgarg1990 force-pushed the prereq branch from 53ca0b6 to 472b81a Compare May 21, 2021 09:58

prithvip reviewed May 21, 2021

View reviewed changes

mayankgarg1990 force-pushed the prereq branch from 472b81a to 7e139dd Compare May 24, 2021 01:31

mayankgarg1990 added 2 commits May 24, 2021 08:22

mayankgarg1990 force-pushed the prereq branch from 7e139dd to d8f3afb Compare May 24, 2021 15:24

tdcmeehan approved these changes May 24, 2021

View reviewed changes

basarhamdionat approved these changes May 24, 2021

View reviewed changes

tdcmeehan merged commit bf94d65 into prestodb:master May 25, 2021

jainxrohit mentioned this pull request Jun 5, 2021

Add release notes for 0.255 #16216

Merged

4 tasks

ajaygeorge mentioned this pull request Jun 9, 2021

Add release notes for 0.255 #16232

Closed

4 tasks

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Add support for custom query prerequisite logic #16073

Add support for custom query prerequisite logic #16073

mayankgarg1990 commented May 11, 2021 •

edited

swapsmagic May 19, 2021

mayankgarg1990 May 20, 2021

tdcmeehan commented May 20, 2021

tdcmeehan May 19, 2021

mayankgarg1990 May 20, 2021

ajaygeorge May 21, 2021

mayankgarg1990 May 21, 2021

ajaygeorge May 21, 2021 •

edited

mayankgarg1990 May 22, 2021

ajaygeorge May 22, 2021

prithvip May 21, 2021

mayankgarg1990 May 22, 2021

prithvip May 21, 2021

mayankgarg1990 May 22, 2021

prithvip May 21, 2021

mayankgarg1990 May 22, 2021

prithvip May 21, 2021

prithvip May 21, 2021

mayankgarg1990 May 22, 2021

prithvip May 21, 2021

mayankgarg1990 May 22, 2021

prithvip May 21, 2021

mayankgarg1990 May 22, 2021

tdcmeehan May 24, 2021

mayankgarg1990 May 24, 2021

		beginQueued(now);
		queuedTime.compareAndSet(null, nanosSince(beginQueuedNanos, now));

Add support for custom query prerequisite logic #16073

Add support for custom query prerequisite logic #16073

Conversation

mayankgarg1990 commented May 11, 2021 • edited

Choose a reason for hiding this comment

Choose a reason for hiding this comment

tdcmeehan commented May 20, 2021

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

ajaygeorge May 21, 2021 • edited

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

mayankgarg1990 commented May 11, 2021 •

edited

ajaygeorge May 21, 2021 •

edited