webapi: detect client disconnect on fetch (#1971)

Add a StreamOps.onDownstreamAbort operator that runs an action when the
downstream cancels before the upstream has completed, and use it on the
/api/v2/fetch response entity to record an atlas.webapi.clientDisconnect
counter when a client goes away mid-stream.

This is observe-only: no request handling behavior changes. A mid-stream
disconnect is currently invisible because the IPC access log entry is
completed when the response headers are produced, not when the body has
been delivered.

watchTermination is not usable here since it completes successfully on a
downstream cancellation and so cannot distinguish an abandoned stream
from one that finished normally.

The endpoint tag uses the matched request path to line up with the
endpoint tag on ipc.server.call, which is derived from the
Netflix-Endpoint header set by the endpointPath directives.

Author: brharrington

atlas-pekko/src/main/scala/com/netflix/atlas/pekko/StreamOps.scala

@@ -45,6 +45,7 @@ import org.apache.pekko.stream.stage.GraphStageWithMaterializedValue
 
 import java.util.concurrent.BlockingQueue
 import scala.concurrent.duration.FiniteDuration
+import scala.util.control.NonFatal
 
 /**
   * Utility functions for commonly used operations on Pekko streams. Most of these are for
@@ -317,6 +318,72 @@ object StreamOps extends StrictLogging {
     Flow[T].via(new MonitorFlow[T](registry, id))
   }
 
+  /**
+    * Performs the specified action if the downstream cancels before the upstream has
+    * completed. When used on the source for an HTTP response entity, this indicates the
+    * client went away before the full response could be delivered.
+    *
+    * Note this is distinct from `watchTermination`, which completes successfully for a
+    * downstream cancellation and so cannot be used to tell an abandoned stream apart from
+    * one that finished normally.
+    *
+    * @param action
+    *     Action to perform, passed the cause reported for the cancellation. For a normal
+    *     cancellation this will be `SubscriptionWithCancelException.NoMoreElementsNeeded`.
+    */
+  def onDownstreamAbort[T](action: Throwable => Unit): Flow[T, T, NotUsed] = {
+    Flow[T].via(new OnDownstreamAbort[T](action))
+  }
+
+  private final class OnDownstreamAbort[T](action: Throwable => Unit)
+      extends GraphStage[FlowShape[T, T]] {
+
+    private val in = Inlet[T]("OnDownstreamAbort.in")
+    private val out = Outlet[T]("OnDownstreamAbort.out")
+
+    override val shape: FlowShape[T, T] = FlowShape(in, out)
+
+    override def createLogic(inheritedAttributes: Attributes): GraphStageLogic = {
+      new GraphStageLogic(shape) with InHandler with OutHandler {
+
+        private var upstreamDone = false
+
+        override def onPush(): Unit = {
+          push(out, grab(in))
+        }
+
+        override def onPull(): Unit = {
+          pull(in)
+        }
+
+        override def onUpstreamFinish(): Unit = {
+          upstreamDone = true
+          completeStage()
+        }
+
+        override def onUpstreamFailure(t: Throwable): Unit = {
+          upstreamDone = true
+          failStage(t)
+        }
+
+        override def onDownstreamFinish(cause: Throwable): Unit = {
+          // Do not let a failure in the action mask the cancellation itself. The cancellation
+          // is propagated in a `finally` so it still happens if the action throws something
+          // that is not caught below.
+          try {
+            if (!upstreamDone) action(cause)
+          } catch {
+            case NonFatal(e) => logger.warn("onDownstreamAbort action failed", e)
+          } finally {
+            cancelStage(cause)
+          }
+        }
+
+        setHandlers(in, out, this)
+      }
+    }
+  }
+
   private final class MonitorFlow[T](registry: Registry, id: String)
       extends GraphStage[FlowShape[T, T]] {
 

atlas-pekko/src/test/scala/com/netflix/atlas/pekko/StreamOpsSuite.scala

@@ -23,6 +23,8 @@ import org.apache.pekko.stream.Materializer
 import org.apache.pekko.stream.scaladsl.Keep
 import org.apache.pekko.stream.scaladsl.Sink
 import org.apache.pekko.stream.scaladsl.Source
+import org.apache.pekko.stream.testkit.scaladsl.TestSink
+import org.apache.pekko.stream.testkit.scaladsl.TestSource
 import com.netflix.spectator.api.DefaultRegistry
 import com.netflix.spectator.api.ManualClock
 import com.netflix.spectator.api.Registry
@@ -533,4 +535,73 @@ class StreamOpsSuite extends FunSuite {
     val result = Await.result(future, Duration.Inf)
     assertEquals(result.size, chunkSize * numChunks)
   }
+
+  //
+  // onDownstreamAbort
+  //
+  // These tests drive both ends of the stream with probes so the ordering is explicit.
+  // `upstream.expectCancellation()` only returns once the cancellation has propagated all
+  // the way to the source, which is strictly after the stage under test has handled it.
+  // That gives a happens-before edge without relying on timing.
+  //
+
+  test("onDownstreamAbort: action runs when downstream cancels early") {
+    val causes = new ArrayBlockingQueue[Throwable](1)
+    val (upstream, downstream) = TestSource[Int]()
+      .via(StreamOps.onDownstreamAbort[Int](causes.add))
+      .toMat(TestSink[Int]())(Keep.both)
+      .run()
+
+    downstream.request(1)
+    upstream.sendNext(1)
+    downstream.expectNext(1)
+    downstream.cancel()
+
+    upstream.expectCancellation()
+    assertEquals(causes.size(), 1)
+  }
+
+  test("onDownstreamAbort: action does not run when upstream completes first") {
+    val causes = new ArrayBlockingQueue[Throwable](1)
+    val (upstream, downstream) = TestSource[Int]()
+      .via(StreamOps.onDownstreamAbort[Int](causes.add))
+      .toMat(TestSink[Int]())(Keep.both)
+      .run()
+
+    downstream.request(1)
+    upstream.sendNext(1)
+    downstream.expectNext(1)
+    upstream.sendComplete()
+    downstream.expectComplete()
+
+    assertEquals(causes.size(), 0)
+  }
+
+  test("onDownstreamAbort: action does not run when upstream fails") {
+    val causes = new ArrayBlockingQueue[Throwable](1)
+    val (upstream, downstream) = TestSource[Int]()
+      .via(StreamOps.onDownstreamAbort[Int](causes.add))
+      .toMat(TestSink[Int]())(Keep.both)
+      .run()
+
+    downstream.request(1)
+    upstream.sendError(new IllegalStateException("boom"))
+    downstream.expectError()
+
+    assertEquals(causes.size(), 0)
+  }
+
+  test("onDownstreamAbort: failure in the action does not mask the cancellation") {
+    val (upstream, downstream) = TestSource[Int]()
+      .via(StreamOps.onDownstreamAbort[Int](_ => throw new IllegalStateException("boom")))
+      .toMat(TestSink[Int]())(Keep.both)
+      .run()
+
+    downstream.request(1)
+    upstream.sendNext(1)
+    downstream.expectNext(1)
+    downstream.cancel()
+
+    upstream.expectCancellation()
+  }
 }

atlas-webapi/src/main/scala/com/netflix/atlas/webapi/FetchRequestSource.scala

@@ -27,6 +27,10 @@ import com.netflix.atlas.eval.graph.GraphConfig
 import com.netflix.atlas.eval.model.TimeSeriesMessage
 import com.netflix.atlas.json3.Json
 import com.netflix.atlas.pekko.DiagnosticMessage
+import com.netflix.atlas.pekko.StreamOps
+import com.netflix.spectator.api.Id
+import com.netflix.spectator.api.Registry
+import com.netflix.spectator.api.Spectator
 import org.apache.pekko.NotUsed
 import org.apache.pekko.actor.ActorRefFactory
 import org.apache.pekko.http.scaladsl.model.HttpEntity
@@ -64,6 +68,22 @@ object FetchRequestSource {
 
   private val chunkSize = ApiSettings.fetchChunkSize
 
+  /**
+    * Counter for clients that go away before the full response has been delivered. The
+    * `endpoint` tag uses the matched request path so it lines up with the `endpoint` tag on
+    * `ipc.server.call`, which comes from the `Netflix-Endpoint` header set by `endpointPath`.
+    */
+  private val clientDisconnectId =
+    Id.create("atlas.webapi.clientDisconnect").withTag("endpoint", "/api/v2/fetch")
+
+  // Atlas only allows `-._A-Za-z0-9^~` in tag values. A Scala object reports a simple name with
+  // a trailing `$` and an anonymous class reports an empty one, so clean up the class name
+  // before using it as a tag value.
+  private def causeTag(cause: Throwable): String = {
+    val name = cause.getClass.getSimpleName.stripSuffix("$")
+    if (name.isEmpty) "unknown" else name
+  }
+
   /**
     * Create an SSE source that can be used as the entity for the HttpResponse.
     */
@@ -139,7 +159,23 @@ object FetchRequestSource {
     * Returns an HttpResponse with an entity that is generated by the fetch source.
     */
   def createResponse(system: ActorRefFactory, graphCfg: GraphConfig): HttpResponse = {
+    createResponse(system, graphCfg, Spectator.globalRegistry())
+  }
+
+  private[webapi] def createResponse(
+    system: ActorRefFactory,
+    graphCfg: GraphConfig,
+    registry: Registry
+  ): HttpResponse = {
+    // Detect a client that goes away part way through the response. The access log entry is
+    // completed when the response headers are produced, so a mid-stream disconnect is not
+    // otherwise visible.
     val source = apply(system, graphCfg)
+      .via(StreamOps.onDownstreamAbort { cause =>
+        registry
+          .counter(clientDisconnectId.withTag("cause", causeTag(cause)))
+          .increment()
+      })
     HttpResponse(
       status = StatusCodes.OK,
       entity = HttpEntity.Chunked(MediaTypes.`text/event-stream`, source)

atlas-webapi/src/test/scala/com/netflix/atlas/webapi/FetchRequestSourceSuite.scala

@@ -0,0 +1,172 @@
+/*
+ * Copyright 2014-2026 Netflix, Inc.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package com.netflix.atlas.webapi
+
+import com.netflix.atlas.core.model.TimeSeries
+import com.netflix.atlas.eval.graph.Grapher
+import com.netflix.spectator.api.AbstractRegistry
+import com.netflix.spectator.api.Clock
+import com.netflix.spectator.api.Counter
+import com.netflix.spectator.api.DefaultRegistry
+import com.netflix.spectator.api.DistributionSummary
+import com.netflix.spectator.api.Gauge
+import com.netflix.spectator.api.Id
+import com.netflix.spectator.api.Measurement
+import com.netflix.spectator.api.Registry
+import com.netflix.spectator.api.Timer
+import com.typesafe.config.ConfigFactory
+import munit.FunSuite
+import org.apache.pekko.actor.Actor
+import org.apache.pekko.actor.ActorSystem
+import org.apache.pekko.actor.Props
+import org.apache.pekko.http.scaladsl.model.HttpEntity
+import org.apache.pekko.http.scaladsl.model.HttpRequest
+import org.apache.pekko.stream.scaladsl.Sink
+import org.apache.pekko.stream.scaladsl.Source
+import org.apache.pekko.stream.testkit.scaladsl.TestSink
+
+import java.util.concurrent.CountDownLatch
+import java.util.concurrent.TimeUnit
+import scala.concurrent.Await
+import scala.concurrent.duration.*
+import scala.jdk.CollectionConverters.*
+
+class FetchRequestSourceSuite extends FunSuite {
+
+  private val config = ConfigFactory.load()
+  private val grapher = Grapher(config)
+
+  /** Db actor that responds with empty data for every request. */
+  private class EmptyDb extends Actor {
+
+    def receive: Receive = {
+      case GraphApi.DataRequest(ctx, exprs, _) =>
+        val data = exprs.map(e => e -> List.empty[TimeSeries]).toMap
+        sender() ! GraphApi.DataResponse(ctx.step, data)
+    }
+  }
+
+  private def disconnects(registry: Registry): Double = {
+    registry
+      .counters()
+      .iterator()
+      .asScala
+      .filter(_.id().name() == "atlas.webapi.clientDisconnect")
+      .map(_.actualCount())
+      .sum
+  }
+
+  /**
+    * Counter that releases a latch once it has been updated. The disconnect is recorded on a
+    * stream thread and there is nothing upstream of the recording stage for the test to
+    * synchronize on, so the counter itself provides the happens-before edge.
+    */
+  private class SignalCounter(delegate: Counter, latch: CountDownLatch) extends Counter {
+
+    override def id(): Id = delegate.id()
+
+    override def measure(): java.lang.Iterable[Measurement] = delegate.measure()
+
+    override def hasExpired: Boolean = delegate.hasExpired
+
+    override def actualCount(): Double = delegate.actualCount()
+
+    override def add(amount: Double): Unit = {
+      delegate.add(amount)
+      latch.countDown()
+    }
+  }
+
+  /** Registry that releases `latch` when a counter named `name` is updated. */
+  private class SignalRegistry(name: String, latch: CountDownLatch)
+      extends AbstractRegistry(Clock.SYSTEM) {
+
+    private val delegate = new DefaultRegistry()
+
+    override protected def newCounter(id: Id): Counter = {
+      val c = delegate.counter(id)
+      if (id.name() == name) new SignalCounter(c, latch) else c
+    }
+
+    override protected def newDistributionSummary(id: Id): DistributionSummary =
+      delegate.distributionSummary(id)
+
+    override protected def newTimer(id: Id): Timer = delegate.timer(id)
+
+    override protected def newGauge(id: Id): Gauge = delegate.gauge(id)
+
+    override protected def newMaxGauge(id: Id): Gauge = delegate.maxGauge(id)
+  }
+
+  private def withFixture(registry: Registry)(
+    f: (ActorSystem, Registry, Source[HttpEntity.ChunkStreamPart, ?]) => Unit
+  ): Unit = {
+    val system = ActorSystem(s"FetchRequestSourceSuite-${System.nanoTime()}")
+    try {
+      system.actorOf(Props(new EmptyDb), "db")
+      val uri = "/api/v2/fetch?q=name,sps,:eq,:sum&s=e-6h&e=now&step=1h"
+      val graphCfg = grapher.toGraphConfig(HttpRequest(uri = uri))
+      val response = FetchRequestSource.createResponse(system, graphCfg, registry)
+      response.entity match {
+        case HttpEntity.Chunked(_, source) => f(system, registry, source)
+        case other                         => fail(s"expected a chunked entity, got $other")
+      }
+    } finally {
+      Await.ready(system.terminate(), 30.seconds)
+    }
+  }
+
+  test("no disconnect recorded when the full response is consumed") {
+    withFixture(new DefaultRegistry) { (system, registry, source) =>
+      implicit val sys: ActorSystem = system
+      // Awaiting the sink gives a happens-before edge: the stream has fully terminated,
+      // so the counter is stable by the time it is read.
+      Await.result(source.runWith(Sink.ignore), 30.seconds)
+      assertEquals(disconnects(registry), 0.0)
+    }
+  }
+
+  test("disconnect recorded when the client goes away") {
+    val latch = new CountDownLatch(1)
+    withFixture(new SignalRegistry("atlas.webapi.clientDisconnect", latch)) {
+      (system, registry, source) =>
+        implicit val sys: ActorSystem = system
+        // Cancel without requesting anything. With no demand the terminating `close` message
+        // cannot be pushed, so the upstream cannot have completed first and the cancellation is
+        // unambiguously an abort. Requesting an element first would be racy: for an empty
+        // response the close message is the only element, and the source completes as soon as
+        // it has been pushed.
+        val probe = source.runWith(TestSink[HttpEntity.ChunkStreamPart]())
+        probe.cancel()
+
+        // Blocks until the counter has actually been updated, so the read below is ordered
+        // after it. The timeout only bounds a failure, it is not a wait for the common case.
+        assert(latch.await(10, TimeUnit.SECONDS), "disconnect counter was never updated")
+        assertEquals(disconnects(registry), 1.0)
+    }
+  }
+
+  test("response entity is a chunked SSE stream") {
+    withFixture(new DefaultRegistry) { (system, registry, source) =>
+      implicit val sys: ActorSystem = system
+      val parts = Await.result(source.runWith(Sink.seq), 30.seconds)
+      assert(parts.nonEmpty)
+      val text = parts.map(_.data().utf8String).mkString
+      assert(text.startsWith("data: "), s"unexpected SSE payload: ${text.take(40)}")
+      assertEquals(disconnects(registry), 0.0)
+    }
+  }
+}