From ec4baa396c994d0bd89aa6e7b1e5496353b92c93 Mon Sep 17 00:00:00 2001
From: Vincent Paturet <vincent.paturet@entur.org>
Date: Thu, 30 May 2024 14:23:43 +0200
Subject: [PATCH] Add OTP request timeout GraphQL instrumentation

---
 .../OTPRequestTimeoutInstrumentation.java     | 36 +++++++++++++++++++
 .../apis/transmodel/TransmodelGraph.java      |  5 ++-
 2 files changed, 40 insertions(+), 1 deletion(-)
 create mode 100644 src/main/java/org/opentripplanner/apis/transmodel/OTPRequestTimeoutInstrumentation.java

diff --git a/src/main/java/org/opentripplanner/apis/transmodel/OTPRequestTimeoutInstrumentation.java b/src/main/java/org/opentripplanner/apis/transmodel/OTPRequestTimeoutInstrumentation.java
new file mode 100644
index 00000000000..acb27eb4e87
--- /dev/null
+++ b/src/main/java/org/opentripplanner/apis/transmodel/OTPRequestTimeoutInstrumentation.java
@@ -0,0 +1,36 @@
+package org.opentripplanner.apis.transmodel;
+
+import static graphql.execution.instrumentation.SimpleInstrumentationContext.noOp;
+
+import graphql.execution.instrumentation.Instrumentation;
+import graphql.execution.instrumentation.InstrumentationContext;
+import graphql.execution.instrumentation.InstrumentationState;
+import graphql.execution.instrumentation.parameters.InstrumentationFieldFetchParameters;
+import java.util.concurrent.atomic.AtomicLong;
+import org.opentripplanner.framework.application.OTPRequestTimeoutException;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+/**
+ * A GraphQL instrumentation that periodically checks the OTP request interruption status while the
+ * query is being processed.
+ */
+public class OTPRequestTimeoutInstrumentation implements Instrumentation {
+
+  private static final Logger LOG = LoggerFactory.getLogger(OTPRequestTimeoutInstrumentation.class);
+
+  private final AtomicLong fieldFetchCounter = new AtomicLong();
+
+  @Override
+  public InstrumentationContext<Object> beginFieldFetch(
+    InstrumentationFieldFetchParameters parameters,
+    InstrumentationState state
+  ) {
+    long fetched = fieldFetchCounter.incrementAndGet();
+    if (fetched % 100000 == 0) {
+      LOG.debug("Fetched {} fields", fetched);
+      OTPRequestTimeoutException.checkForTimeout();
+    }
+    return noOp();
+  }
+}
diff --git a/src/main/java/org/opentripplanner/apis/transmodel/TransmodelGraph.java b/src/main/java/org/opentripplanner/apis/transmodel/TransmodelGraph.java
index 051e369dde0..2cd405212c2 100644
--- a/src/main/java/org/opentripplanner/apis/transmodel/TransmodelGraph.java
+++ b/src/main/java/org/opentripplanner/apis/transmodel/TransmodelGraph.java
@@ -79,7 +79,10 @@ Response executeGraphQL(
   }
 
   private static Instrumentation createInstrumentation(int maxResolves, Iterable<Tag> tracingTags) {
-    Instrumentation instrumentation = new MaxQueryComplexityInstrumentation(maxResolves);
+    Instrumentation instrumentation = new ChainedInstrumentation(
+      new MaxQueryComplexityInstrumentation(maxResolves),
+      new OTPRequestTimeoutInstrumentation()
+    );
 
     if (OTPFeature.ActuatorAPI.isOn()) {
       instrumentation =