io.axoniq.flowcontrol.producer.grpc.FlowControlledOutgoingStream Maven / Gradle / Ivy
/*
* Copyright (c) 2021. AxonIQ
*
* 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 io.axoniq.flowcontrol.producer.grpc;
import io.axoniq.flowcontrol.OutgoingStream;
import io.axoniq.flowcontrol.producer.grpc.ActiveSubscriptions.ActiveSubscription;
import io.axoniq.flowcontrol.producer.grpc.shutdown.ProducerCompletedShutDownPolicy;
import io.axoniq.flowcontrol.producer.grpc.subscriptions.RoundRobinSubscriptions;
import io.grpc.stub.CallStreamObserver;
import org.reactivestreams.Publisher;
import org.reactivestreams.Subscriber;
import org.reactivestreams.Subscription;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import java.util.Arrays;
import java.util.Optional;
import java.util.concurrent.Executor;
import java.util.concurrent.atomic.AtomicBoolean;
import java.util.concurrent.atomic.AtomicReference;
import java.util.stream.Collectors;
/**
* @author Sara Pellegrini
* @author Milan Savic
* @since 1.0
*/
public class FlowControlledOutgoingStream implements OutgoingStream {
private final ShutdownPolicy shutdownPolicy;
private final ActiveSubscriptions activeSubscriptions;
private final CallStreamObserver outgoingStream;
private final AtomicBoolean flowGate = new AtomicBoolean();
private final Logger logger = LoggerFactory.getLogger(FlowControlledOutgoingStream.class);
public FlowControlledOutgoingStream(CallStreamObserver outgoingStream,
Executor executor) {
this(outgoingStream,
new ProducerCompletedShutDownPolicy<>(),
new RoundRobinSubscriptions(),
executor);
}
public FlowControlledOutgoingStream(CallStreamObserver outgoingStream,
ShutdownPolicy shutdownPolicy,
ActiveSubscriptions activeSubscriptions,
Executor executor) {
this.outgoingStream = outgoingStream;
this.outgoingStream.setOnReadyHandler(() -> executor.execute(this::flow));
this.shutdownPolicy = shutdownPolicy;
this.shutdownPolicy.initialize(this::shutdown, this::shutdown);
this.activeSubscriptions = activeSubscriptions;
}
@Override
public void accept(Publisher publisher) {
publisher.subscribe(new FlowControlledSubscriber());
}
private void flow() {
if (activeSubscriptions.hasNext()
&& outgoingStream.isReady()) {
if (flowGate.compareAndSet(false, true)) {
request();
logger.debug("Flow control gate opened.");
} else {
logger.debug("Flow control gate already open.");
}
} else {
logger.debug("Flow control gate still closed.");
}
}
private void request() {
boolean hasNext = activeSubscriptions.hasNext();
if (hasNext && outgoingStream.isReady()) {
Optional subscription = activeSubscriptions.next();
if (subscription.isPresent()) {
subscription.get().request();
return;
}
}
flowGate.set(false);
String cause = hasNext ? "stream is not ready" : "no message available yet";
logger.debug("Flow control gate closed. Cause: {}.", cause);
flow();
}
private void shutdown() {
if (logger.isDebugEnabled()) {
logger.debug("Shutdown call: {}", Arrays.stream(Thread.currentThread().getStackTrace())
.map(StackTraceElement::toString)
.collect(Collectors.joining("\n")));
}
shutdownPolicy.shutdown();
outgoingStream.onCompleted();
}
private void shutdown(Throwable t) {
if (logger.isDebugEnabled()) {
logger.debug("Shutdown call exceptionally: " +
Arrays.stream(Thread.currentThread().getStackTrace())
.map(StackTraceElement::toString)
.collect(Collectors.joining("\n")), t);
}
shutdownPolicy.shutdown();
outgoingStream.onError(t);
}
protected class FlowControlledSubscriber implements Subscriber, ActiveSubscription {
private final AtomicReference subscription = new AtomicReference<>();
private volatile boolean pending = false;
@Override
public void onSubscribe(Subscription subscription) {
if (subscription == null) {
throw new NullPointerException("Subscription must be null!");
}
if (!this.subscription.compareAndSet(null, subscription)) {
subscription.cancel();
return;
}
activeSubscriptions.add(this);
flow();
}
@Override
public void onNext(Message message) {
if (message == null) {
throw new NullPointerException();
}
pending = false;
outgoingStream.onNext(message);
shutdownPolicy.onMessageConsumed(message);
if (outgoingStream.isReady()) {
FlowControlledOutgoingStream.this.request();
} else {
flowGate.set(false);
logger.debug("Flow control gate closed. Cause: stream is not ready.");
flow();
}
}
@Override
public void onComplete() {
shutdownPolicy.onPublisherComplete();
dispose();
}
@Override
public void onError(Throwable t) {
if (t == null) {
throw new NullPointerException();
}
shutdownPolicy.onPublisherError(t);
dispose();
}
@Override
public void request() {
pending = true;
subscription.get().request(1);
}
private void dispose() {
logger.debug("Disposing subscription.");
activeSubscriptions.remove(this);
// if the activeSubscription is removed after one message is requested on this subscription,
// we need to notify the outgoing stream to continue requesting the next subscription because
// this one will never invoke onNext as it is ended.
// This is thread safe because when pending is true, the flowGate is opened (true).
if (pending) {
FlowControlledOutgoingStream.this.request();
}
}
}
}