org.reactivestreams.Subscription Scala Examples

// Copyright (c) 2020 Digital Asset (Switzerland) GmbH and/or its affiliates. All rights reserved.
// SPDX-License-Identifier: Apache-2.0

package com.daml.grpc.adapter.server.rs

import com.daml.grpc.adapter.TestExecutionSequencerFactory
import org.reactivestreams.tck.SubscriberWhiteboxVerification.SubscriberPuppet
import org.reactivestreams.tck.{SubscriberWhiteboxVerification, TestEnvironment}
import org.reactivestreams.{Subscriber, Subscription}
import org.scalatest.testng.TestNGSuiteLike

// This suite uses Integer instead of Int because if we'd use Int, scala would interpret the `null` element from
// test required_spec213_onNext_mustThrowNullPointerExceptionWhenParametersAreNull as numerical 0.
class ServerSubscriberWhiteboxTest
    extends SubscriberWhiteboxVerification[Integer](new TestEnvironment(500, 500, false))
    with TestNGSuiteLike {

  override def createSubscriber(
      probe: SubscriberWhiteboxVerification.WhiteboxSubscriberProbe[Integer])
    : Subscriber[Integer] = {
    val so = new MockServerCallStreamObserver[Integer]
    val sub = new ServerSubscriber[Integer](
      so,
      TestExecutionSequencerFactory.instance.getExecutionSequencer) {
      override def onSubscribe(subscription: Subscription): Unit = {
        super.onSubscribe(subscription)
        probe.registerOnSubscribe(new SubscriberPuppet {
          override def triggerRequest(elements: Long): Unit = {
            for (i <- 1l to elements) so.demandResponse()
          }

          override def signalCancel(): Unit = {
            so.signalCancellation()
          }
        })
      }

      override def onNext(response: Integer): Unit = {
        super.onNext(response)
        probe.registerOnNext(response)
      }

      override def onError(throwable: Throwable): Unit = {
        super.onError(throwable)
        probe.registerOnError(throwable)
      }

      override def onComplete(): Unit = {
        super.onComplete()
        probe.registerOnComplete()
      }
    }
    so.demandResponse()
    sub
  }

  override def createElement(element: Int): Integer = element
} 

package play.api.libs.ws.ahc

import java.net.URI

import org.reactivestreams.Subscriber
import org.reactivestreams.Subscription
import org.reactivestreams.Publisher
import play.shaded.ahc.io.netty.handler.codec.http.HttpHeaders
import akka.Done
import play.shaded.ahc.org.asynchttpclient.AsyncHandler.State
import play.shaded.ahc.org.asynchttpclient._
import play.shaded.ahc.org.asynchttpclient.handler.StreamedAsyncHandler

import scala.concurrent.Promise

case class StreamedState(
    statusCode: Int = -1,
    statusText: String = "",
    uriOption: Option[URI] = None,
    responseHeaders: Map[String, scala.collection.Seq[String]] = Map.empty,
    publisher: Publisher[HttpResponseBodyPart] = EmptyPublisher
)

class DefaultStreamedAsyncHandler[T](
    f: java.util.function.Function[StreamedState, T],
    streamStarted: Promise[T],
    streamDone: Promise[Done]
) extends StreamedAsyncHandler[Unit]
    with AhcUtilities {
  private var state = StreamedState()

  def onStream(publisher: Publisher[HttpResponseBodyPart]): State = {
    if (this.state.publisher != EmptyPublisher) State.ABORT
    else {
      this.state = state.copy(publisher = publisher)
      streamStarted.success(f(state))
      State.CONTINUE
    }
  }

  override def onStatusReceived(status: HttpResponseStatus): State = {
    if (this.state.publisher != EmptyPublisher) State.ABORT
    else {
      state = state.copy(
        statusCode = status.getStatusCode,
        statusText = status.getStatusText,
        uriOption = Option(status.getUri.toJavaNetURI)
      )
      State.CONTINUE
    }
  }

  override def onHeadersReceived(h: HttpHeaders): State = {
    if (this.state.publisher != EmptyPublisher) State.ABORT
    else {
      state = state.copy(responseHeaders = headersToMap(h))
      State.CONTINUE
    }
  }

  override def onBodyPartReceived(bodyPart: HttpResponseBodyPart): State =
    throw new IllegalStateException("Should not have received bodypart")

  override def onCompleted(): Unit = {
    // EmptyPublisher can be replaces with `Source.empty` when we carry out the refactoring
    // mentioned in the `execute2` method.
    streamStarted.trySuccess(f(state.copy(publisher = EmptyPublisher)))
    streamDone.trySuccess(Done)
  }

  override def onThrowable(t: Throwable): Unit = {
    streamStarted.tryFailure(t)
    streamDone.tryFailure(t)
  }
}

private case object EmptyPublisher extends Publisher[HttpResponseBodyPart] {
  def subscribe(s: Subscriber[_ >: HttpResponseBodyPart]): Unit = {
    if (s eq null)
      throw new NullPointerException("Subscriber must not be null, rule 1.9")
    s.onSubscribe(CancelledSubscription)
    s.onComplete()
  }
  private case object CancelledSubscription extends Subscription {
    override def request(elements: Long): Unit = ()
    override def cancel(): Unit                = ()
  }
} 

package org.apache.gearpump.akkastream.task

import java.time.Instant
import java.util
import java.util.concurrent.TimeUnit

import akka.actor.Actor.Receive
import akka.actor.{Actor, ActorRef, ActorSystem, Props}
import akka.util.Timeout
import org.apache.gearpump.Message
import org.apache.gearpump.akkastream.task.SinkBridgeTask.RequestMessage
import org.apache.gearpump.cluster.UserConfig
import org.apache.gearpump.cluster.client.ClientContext
import org.apache.gearpump.streaming.ProcessorId
import org.apache.gearpump.streaming.appmaster.AppMaster.{LookupTaskActorRef, TaskActorRef}
import org.apache.gearpump.streaming.task.{Task, TaskContext, TaskId}
import org.apache.gearpump.util.LogUtil
import org.reactivestreams.{Publisher, Subscriber, Subscription}


class SinkBridgeTask(taskContext : TaskContext, userConf : UserConfig)
  extends Task(taskContext, userConf) {
  import taskContext.taskId

  val queue = new util.LinkedList[Message]()
  var subscriber: ActorRef = _

  var request: Int = 0

  override def onStart(startTime : Instant) : Unit = {}

  override def onNext(msg: Message) : Unit = {
    queue.add(msg)
    trySendingData()
  }

  override def onStop() : Unit = {}

  private def trySendingData(): Unit = {
    if (subscriber != null) {
      (0 to request).map(_ => queue.poll()).filter(_ != null).foreach { msg =>
        subscriber ! msg.value
        request -= 1
      }
    }
  }

  override def receiveUnManagedMessage: Receive = {
    case RequestMessage(n) =>
      this.subscriber = sender
      LOG.info("the downstream has requested " + n + " messages from " + subscriber)
      request += n.toInt
      trySendingData()
    case msg =>
      LOG.error("Failed! Received unknown message " + "taskId: " + taskId + ", " + msg.toString)
  }
}

object SinkBridgeTask {

  case class RequestMessage(number: Int)

  class SinkBridgeTaskClient(system: ActorSystem, context: ClientContext, appId: Int,
      processorId: ProcessorId) extends Publisher[AnyRef] with Subscription {
    private val taskId = TaskId(processorId, index = 0)
    private val LOG = LogUtil.getLogger(getClass)

    private var actor: ActorRef = _
    import system.dispatcher

    private val task =
      context.askAppMaster[TaskActorRef](appId, LookupTaskActorRef(taskId)).map{container =>
      // println("Successfully resolved taskRef for taskId " + taskId + ", " + container.task)
      container.task
    }

    override def subscribe(subscriber: Subscriber[_ >: AnyRef]): Unit = {
      this.actor = system.actorOf(Props(new ClientActor(subscriber)))
      subscriber.onSubscribe(this)
    }

    override def cancel(): Unit = Unit

    private implicit val timeout = Timeout(5, TimeUnit.SECONDS)

    override def request(l: Long): Unit = {
      task.foreach{ task =>
        task.tell(RequestMessage(l.toInt), actor)
      }
    }
  }

  class ClientActor(subscriber: Subscriber[_ >: AnyRef]) extends Actor {
    def receive: Receive = {
      case result: AnyRef =>
        subscriber.onNext(result)
    }
  }
} 

package org.apache.gearpump.akkastream.task

import java.time.Instant

import akka.actor.Actor.Receive
import org.apache.gearpump.Message
import org.apache.gearpump.akkastream.task.SourceBridgeTask.{AkkaStreamMessage, Complete, Error}
import org.apache.gearpump.cluster.UserConfig
import org.apache.gearpump.cluster.client.ClientContext
import org.apache.gearpump.streaming.ProcessorId
import org.apache.gearpump.streaming.appmaster.AppMaster.{LookupTaskActorRef, TaskActorRef}
import org.apache.gearpump.streaming.task.{Task, TaskContext, TaskId}
import org.reactivestreams.{Subscriber, Subscription}

import scala.concurrent.ExecutionContext


class SourceBridgeTask(taskContext : TaskContext, userConf : UserConfig)
  extends Task(taskContext, userConf) {
  import taskContext.taskId

  override def onStart(startTime : Instant) : Unit = {}

  override def onNext(msg: Message) : Unit = {
    LOG.info("AkkaStreamSource receiving message " + msg)
  }

  override def onStop() : Unit = {}

  override def receiveUnManagedMessage: Receive = {
    case Error(ex) =>
      LOG.error("the stream has error", ex)
    case AkkaStreamMessage(msg) =>
      LOG.info("we have received message from akka stream source: " + msg)
      taskContext.output(Message(msg, Instant.now()))
    case Complete(description) =>
      LOG.info("the stream is completed: " + description)
    case msg =>
      LOG.error("Failed! Received unknown message " + "taskId: " + taskId + ", " + msg.toString)
  }
}


object SourceBridgeTask {
  case class Error(ex: java.lang.Throwable)

  case class Complete(description: String)

  case class AkkaStreamMessage[T >: AnyRef](msg: T)

  class SourceBridgeTaskClient[T >: AnyRef](ec: ExecutionContext,
      context: ClientContext, appId: Int, processorId: ProcessorId) extends Subscriber[T] {
    val taskId = TaskId(processorId, 0)
    var subscription: Subscription = _
    implicit val dispatcher = ec

    val task = context.askAppMaster[TaskActorRef](appId,
      LookupTaskActorRef(taskId)).map{container =>
      // println("Successfully resolved taskRef for taskId " + taskId + ", " + container.task)
      container.task
    }

    override def onError(throwable: Throwable): Unit = {
      task.map(task => task ! Error(throwable))
    }

    override def onSubscribe(subscription: Subscription): Unit = {
      // when taskActorRef is resolved, request message from upstream
      this.subscription = subscription
      task.map(task => subscription.request(1))
    }

    override def onComplete(): Unit = {
      task.map(task => task ! Complete("the upstream is completed"))
    }

    override def onNext(t: T): Unit = {
      task.map {task =>
        task ! AkkaStreamMessage(t)
      }
      subscription.request(1)
    }
  }
} 

package engine

import akka.actor.{Actor, ActorRef, Props}
import com.typesafe.scalalogging.LazyLogging
import engine.BackPressuredWebSocketActor._
import monix.execution.Scheduler
import monix.execution.rstreams.SingleAssignmentSubscription
import monix.reactive.Observable
import org.reactivestreams.{Subscriber, Subscription}
import play.api.libs.json._

import scala.concurrent.duration._
import scala.util.Try


class BackPressuredWebSocketActor[T: Writes]
  (producer: Observable[T], out: ActorRef)(implicit s: Scheduler)
  extends Actor with LazyLogging {

  def receive: Receive = {
    case JsNumber(nr) if nr > 0 =>
      Try(nr.toLongExact).foreach(subscription.request)
  }

  private[this] val subscription = SingleAssignmentSubscription()

  def now(): Long =
    System.currentTimeMillis()

  override def preStart(): Unit = {
    super.preStart()

    val source = {
      val initial = Observable.evalOnce(initMessage(now()))
      val obs = initial ++ producer.map(x => Json.toJson(x))
      val timeout = obs.debounceRepeated(5.seconds).map(_ => keepAliveMessage(now()))

      Observable
        .merge(obs, timeout)
        .whileBusyDropEventsAndSignal(nr => onOverflow(nr, now()))
    }

    source.toReactivePublisher.subscribe(new Subscriber[JsValue] {
      def onSubscribe(s: Subscription): Unit = {
        subscription := s
      }

      def onNext(json: JsValue): Unit = {
        out ! json
      }

      def onError(t: Throwable): Unit = {
        logger.warn(s"Error while serving a web-socket stream", t)
        out ! Json.obj(
          "event" -> "error",
          "type" -> t.getClass.getName,
          "message" -> t.getMessage,
          "timestamp" -> now())

        context.stop(self)
      }

      def onComplete(): Unit = {
        out ! Json.obj("event" -> "complete", "timestamp" -> now())
        context.stop(self)
      }
    })
  }

  override def postStop(): Unit = {
    subscription.cancel()
    super.postStop()
  }
}

object BackPressuredWebSocketActor {
  
  def initMessage(now: Long) = {
    Json.obj("event" -> "init", "timestamp" -> now)
  }
} 

package swave.core.impl.stages.spout

import org.reactivestreams.{Publisher, Subscriber, Subscription}
import swave.core.Stage
import swave.core.impl.Outport
import swave.core.impl.rs.RSCompliance
import swave.core.impl.stages.SpoutStage
import swave.core.macros.StageImplementation
import swave.core.util._

// format: OFF
@StageImplementation
private[core] final class PublisherSpoutStage(publisher: Publisher[AnyRef]) extends SpoutStage  { stage =>

  def kind = Stage.Kind.Spout.FromPublisher(publisher)

  connectOutAndSealWith { out ⇒
    region.impl.requestDispatcherAssignment()
    region.impl.registerForXStart(this)
    awaitingXStart(out)
  }

  def awaitingXStart(out: Outport) = state(
    xStart = () => {
      publisher.subscribe {
        new Subscriber[AnyRef] {
          def onSubscribe(s: Subscription) = {
            RSCompliance.verifyNonNull(s, "Subscription", "2.13")
            region.enqueueXEvent(stage, s)
          }
          def onNext(elem: AnyRef) = {
            RSCompliance.verifyNonNull(elem, "Element", "2.13")
            region.enqueueOnNext(stage, elem, stage)
          }
          def onComplete() = region.enqueueOnComplete(stage, stage)
          def onError(e: Throwable) = {
            RSCompliance.verifyNonNull(e, "Throwable", "2.13")
            region.enqueueOnError(stage, e, stage)
          }
        }
      }
      awaitingSubscription(out, 0L)
    })

  def awaitingSubscription(out: Outport, requested: Long): State = state(
    request = (n, _) ⇒ awaitingSubscription(out, requested ⊹ n),
    cancel = _ => awaitingSubscriptionDownstreamCancelled(),

    xEvent = {
      case s: Subscription =>
        if (requested > 0) s.request(requested)
        running(out, s)
    })

  def awaitingSubscriptionDownstreamCancelled(): State = state(
    request = (_, _) ⇒ stay(),
    cancel = _ => stay(),

    xEvent = {
      case s: Subscription =>
        s.cancel()
        stop()
    })

  def running(out: Outport, subscription: Subscription) = state(
    intercept = false,

    request = (n, _) ⇒ {
      subscription.request(n.toLong)
      stay()
    },

    cancel = _ => {
      subscription.cancel()
      stop()
    },

    onNext = onNextF(out),
    onComplete = stopCompleteF(out),
    onError = stopErrorF(out),

    xEvent = { case s: Subscription =>
      s.cancel()
      stay()
    })
} 

package korolev.akka.util

import korolev.effect.{Effect, Stream}
import org.reactivestreams.{Subscriber, Subscription}

final class KorolevStreamSubscriber[F[_]: Effect,T] extends Stream[F, T] with Subscriber[T] {

  private var subscription: Subscription = _

  private var pullCallback: Either[Throwable, Option[T]] => Unit = _

  private var consumedCallback: Either[Throwable, Unit] => Unit = _

  private var completeValue: Either[Throwable, Unit] = _

  def onSubscribe(subscription: Subscription): Unit = {
    this.subscription = subscription
    if (pullCallback != null)
      subscription.request(1)
  }

  def onNext(value: T): Unit = {
    val cb = pullCallback
    pullCallback = null
    cb(Right(Some(value)))
  }

  def onError(error: Throwable): Unit = {
    completeWith(Left(error))
    val cb = pullCallback
    pullCallback = null
    cb(Left(error))
  }

  def onComplete(): Unit = {
    completeWith(Right(()))
    val cb = pullCallback
    pullCallback = null
    cb(Right(None))
  }

  private def completeWith(that: Either[Throwable, Unit]): Unit = {
    completeValue = that
    if (consumedCallback != null) {
      val cb = consumedCallback
      consumedCallback = null
      cb(that)
    }
  }

  def pull(): F[Option[T]] = Effect[F].promise { cb =>
    if (completeValue == null) {
      pullCallback = cb
      if (subscription != null) {
        subscription.request(1)
      }
    } else {
      cb(completeValue.map(_ => None))
    }
  }

  def cancel(): F[Unit] = Effect[F].delay(subscription.cancel())

  val consumed: F[Unit] = Effect[F].promise { cb =>
    if (completeValue != null) cb(completeValue)
    else consumedCallback = cb
  }

  val size: Option[Long] = None
}