Skip to content

RATIS-2421. Gracefully cancel stream after complete in GrpcLogAppender #1363

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
symious wants to merge 5 commits into
base: master
Choose a base branch
from
Open

RATIS-2421. Gracefully cancel stream after complete in GrpcLogAppender #1363

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
5 commits
Select commit Hold shift + click to select a range
54af1e1
RATIS-2421. Gracefully cancel stream after complete
symious Mar 6, 2026
03d189d
RATIS-2421. Fix checkstyle
symious Mar 6, 2026
6726e78
RATIS-2421. Fix unit test
symious Mar 6, 2026
30b53c2
RATIS-2421. Revert exception message
symious Mar 6, 2026
16297f7
RATIS-2421. Address comments
symious Mar 11, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
96 changes: 87 additions & 9 deletions ratis-grpc/src/main/java/org/apache/ratis/grpc/server/GrpcLogAppender.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -35,6 +35,7 @@
import org.apache.ratis.server.util.ServerStringUtils;
import org.apache.ratis.thirdparty.io.grpc.StatusRuntimeException;
import org.apache.ratis.thirdparty.io.grpc.stub.CallStreamObserver;
import org.apache.ratis.thirdparty.io.grpc.stub.ClientCallStreamObserver;
import org.apache.ratis.thirdparty.io.grpc.stub.StreamObserver;
import org.apache.ratis.proto.RaftProtos.AppendEntriesReplyProto;
import org.apache.ratis.proto.RaftProtos.AppendEntriesReplyProto.AppendResult;
Expand All @@ -60,7 +61,10 @@
import java.util.concurrent.CompletableFuture;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.ExecutionException;
import java.util.concurrent.Executors;
import java.util.concurrent.ScheduledExecutorService;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.atomic.AtomicBoolean;
import java.util.concurrent.atomic.AtomicLong;

/**
Expand Down Expand Up @@ -341,20 +345,36 @@ private boolean haveTooManyPendingRequests() {
}

static class StreamObservers {
private final CallStreamObserver<AppendEntriesRequestProto> appendLog;
private final CallStreamObserver<AppendEntriesRequestProto> heartbeat;
private final ClientCallStreamObserver<AppendEntriesRequestProto> appendLog;
private final ClientCallStreamObserver<AppendEntriesRequestProto> heartbeat;
private final TimeDuration waitForReady;
private final TimeDuration completeGracePeriod;
private volatile boolean running = true;

private final ScheduledExecutorService closer =
Executors.newSingleThreadScheduledExecutor(r -> {
Thread t = new Thread(r, "grpc-log-appender-stream-closer");
t.setDaemon(true);
return t;
});

private final AtomicBoolean completed = new AtomicBoolean(false);
private final AtomicBoolean cancelled = new AtomicBoolean(false);

StreamObservers(GrpcServerProtocolClient client, AppendLogResponseHandler handler, boolean separateHeartbeat,
TimeDuration waitTimeMin) {
this.appendLog = client.appendEntries(handler, false);
this.heartbeat = separateHeartbeat? client.appendEntries(handler, true): null;
TimeDuration waitTimeMin, TimeDuration completeGracePeriod) {
this.appendLog = (ClientCallStreamObserver<AppendEntriesRequestProto>) client.appendEntries(handler, false);
this.heartbeat = separateHeartbeat?
(ClientCallStreamObserver<AppendEntriesRequestProto>) client.appendEntries(handler, true): null;
this.waitForReady = waitTimeMin.isPositive()? waitTimeMin: TimeDuration.ONE_MILLISECOND;
this.completeGracePeriod = completeGracePeriod.isPositive()? completeGracePeriod : TimeDuration.ONE_SECOND;
}

void onNext(AppendEntriesRequestProto proto)
throws InterruptedIOException {
if (!running) {
throw new InterruptedIOException("StreamObservers is stopping/closing");
}
CallStreamObserver<AppendEntriesRequestProto> stream;
boolean isHeartBeat = heartbeat != null && proto.getEntriesCount() == 0;
if (isHeartBeat) {
Expand All @@ -366,16 +386,73 @@ void onNext(AppendEntriesRequestProto proto)
while (!stream.isReady() && running) {
sleep(waitForReady, isHeartBeat);
}
stream.onNext(proto);
try {
stream.onNext(proto);
} catch (Exception e) {
InterruptedIOException ioe =
new InterruptedIOException("Failed to send request via stream");
ioe.initCause(e);
throw ioe;
}
}

void stop() {
running = false;
}

void onCompleted() {
appendLog.onCompleted();
Optional.ofNullable(heartbeat).ifPresent(StreamObserver::onCompleted);
if (completed.compareAndSet(false, true)) {
completeStreamGracefully(appendLog, "appendLog");
Optional.ofNullable(heartbeat)
.ifPresent(s -> completeStreamGracefully(s, "heartbeat"));
}
final long delayMs = Math.max(1L, completeGracePeriod.toLong(TimeUnit.MILLISECONDS));
closer.schedule(this::cancelIfStillNeeded, delayMs, TimeUnit.MILLISECONDS);
Copy link
Copy Markdown
Contributor

@spacemonkd spacemonkd Mar 10, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This schedules cancelIfStillNeeded() outside the completed.compareAndSet(...) block. Repeated onCompleted() calls can queue redundant cancel tasks.

Can this be avoided?
Also I think if it’s invoked again after shutdownCloser() it may throw exception.

Copy link
Copy Markdown
Contributor Author

@symious symious Mar 11, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The cancel and shutdown parts can be called multiple times.

Copy link
Copy Markdown
Contributor

@spacemonkd spacemonkd Mar 11, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

But if it is completed already aren't subsequent cancel calls queuing tasks which are no longer required?

Copy link
Copy Markdown
Contributor Author

@symious symious Mar 11, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

If onComplete() is invoked, the scheudler should be called and shutdown later, later requests should be rejected then.

}

void cancelNow(String reason, Throwable cause) {
if (cancelled.compareAndSet(false, true)) {
running = false;
cancelStream(appendLog, "appendLog", reason, cause);
Optional.ofNullable(heartbeat)
.ifPresent(s -> cancelStream(s, "heartbeat", reason, cause));
shutdownCloser();
}
}

private void cancelIfStillNeeded() {
if (cancelled.compareAndSet(false, true)) {
cancelStream(appendLog, "appendLog", "Stream completion timeout", null);
Optional.ofNullable(heartbeat)
.ifPresent(s -> cancelStream(s, "heartbeat", "Stream completion timeout", null));
}
shutdownCloser();
}

private void completeStreamGracefully(
ClientCallStreamObserver<AppendEntriesRequestProto> stream,
String name) {
try {
stream.onCompleted();
} catch (Exception e) {
LOG.warn("Failed to call onCompleted on {}", name, e);
}
}

private void cancelStream(
ClientCallStreamObserver<AppendEntriesRequestProto> stream,
Copy link
Copy Markdown
Contributor

@spacemonkd spacemonkd Mar 10, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

We are doing explicit cast to ClientCallStreamObserver<>, but I am thinking can there be a case where appendEntries() can return only CallStreamObserver?
That is the method signature in GrpcServerProtocolClient.

In such a case this can throw error.
Better to accept CallStreamObserver<> and then do an instanceof check?

try {
    if (stream instanceof ClientCallStreamObserver) {
      ((ClientCallStreamObserver<AppendEntriesRequestProto>) stream).cancel(reason, cause);
    } else {
      ...handle cancel when it is not clientcallstreamobserver, something like stream.onError(...)?
    }
  } catch (Exception e) {
    LOG.warn("Failed to cancel {}", name, e);
  }

Copy link
Copy Markdown
Contributor Author

@symious symious Mar 11, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

As a client, this one should always be a ClientCallStreamObserver.

String name,
String reason,
Throwable cause) {
try {
stream.cancel(reason, cause);
} catch (Exception e) {
LOG.warn("Failed to cancel {}", name, e);
}
}

private void shutdownCloser() {
closer.shutdown();
}
}

Expand Down Expand Up @@ -404,7 +481,8 @@ private void appendLog(boolean heartbeat) throws IOException {
increaseNextIndex(pending);
if (appendLogRequestObserver == null) {
appendLogRequestObserver = new StreamObservers(
getClient(), new AppendLogResponseHandler(), useSeparateHBChannel, getWaitTimeMin());
getClient(), new AppendLogResponseHandler(), useSeparateHBChannel, getWaitTimeMin(),
getCompleteGracePeriod());
}
}

Expand Down
10 changes: 10 additions & 0 deletions ratis-server-api/src/main/java/org/apache/ratis/server/RaftServerConfigKeys.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -686,6 +686,16 @@ static void setWaitTimeMin(RaftProperties properties, TimeDuration minDuration)
setTimeDuration(properties::setTimeDuration, WAIT_TIME_MIN_KEY, minDuration);
}

String COMPLETE_GRACE_PERIOD_KEY = PREFIX + ".complete.grace.period";
TimeDuration COMPLETE_GRACE_PERIOD_DEFAULT = TimeDuration.ONE_SECOND;
static TimeDuration completeGracePeriod(RaftProperties properties) {
return getTimeDuration(properties.getTimeDuration(COMPLETE_GRACE_PERIOD_DEFAULT.getUnit()),
COMPLETE_GRACE_PERIOD_KEY, COMPLETE_GRACE_PERIOD_DEFAULT, getDefaultLog());
}
static void setCompleteGracePeriod(RaftProperties properties, TimeDuration duration) {
setTimeDuration(properties::setTimeDuration, COMPLETE_GRACE_PERIOD_KEY, duration);
}

String RETRY_POLICY_KEY = PREFIX + ".retry.policy";
/**
* The min wait time as 1ms (0 is not allowed) for first 10,
Expand Down
6 changes: 6 additions & 0 deletions ratis-server/src/main/java/org/apache/ratis/server/leader/LogAppenderBase.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -61,6 +61,7 @@ public abstract class LogAppenderBase implements LogAppender {

private final AtomicBoolean heartbeatTrigger = new AtomicBoolean();
private final TimeDuration waitTimeMin;
private final TimeDuration completeGracePeriod;

protected LogAppenderBase(RaftServer.Division server, LeaderState leaderState, FollowerInfo f) {
this.follower = f;
Expand All @@ -78,6 +79,7 @@ protected LogAppenderBase(RaftServer.Division server, LeaderState leaderState, F
this.eventAwaitForSignal = new AwaitForSignal(name);

this.waitTimeMin = RaftServerConfigKeys.Log.Appender.waitTimeMin(properties);
this.completeGracePeriod = RaftServerConfigKeys.Log.Appender.completeGracePeriod(properties);
}

@Override
Expand Down Expand Up @@ -144,6 +146,10 @@ protected TimeDuration getWaitTimeMin() {
return waitTimeMin;
}

protected TimeDuration getCompleteGracePeriod() {
return completeGracePeriod;
}

protected TimeDuration getRemainingWaitTime() {
return waitTimeMin.add(getFollower().getLastRpcSendTime().elapsedTime().negate());
}
Expand Down
Loading