Skip to content

Commit

Permalink
Add gRPC topic producer endpoint - Java side (#693)
Browse files Browse the repository at this point in the history 
* Add gRPC topic producer endpoint - Java side

* Changes following review
  • Loading branch information
@cbornet
cbornet authored Nov 14, 2023
1 parent 63b3c05 commit b16e58d
Show file tree
Hide file tree
Showing 11 changed files with 577 additions and 67 deletions.
111 changes: 111 additions & 0 deletions ...ents/langstream-agent-grpc/src/main/java/ai/langstream/agents/grpc/AbstractGrpcAgent.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -17,20 +17,24 @@

import ai.langstream.api.runner.code.AbstractAgentCode;
import ai.langstream.api.runner.code.SimpleRecord;
import ai.langstream.api.runner.topics.TopicProducer;
import ai.langstream.api.util.ConfigurationUtils;
import com.fasterxml.jackson.core.JsonProcessingException;
import com.fasterxml.jackson.databind.JsonNode;
import com.fasterxml.jackson.databind.ObjectMapper;
import com.google.protobuf.ByteString;
import com.google.protobuf.Empty;
import io.grpc.ManagedChannel;
import io.grpc.stub.StreamObserver;
import java.io.ByteArrayOutputStream;
import java.io.IOException;
import java.util.ArrayList;
import java.util.Collection;
import java.util.List;
import java.util.Map;
import java.util.concurrent.CompletableFuture;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.ExecutionException;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.atomic.AtomicBoolean;
import java.util.concurrent.atomic.AtomicInteger;
Expand Down Expand Up @@ -65,6 +69,12 @@ abstract class AbstractGrpcAgent extends AbstractAgentCode {
protected final AtomicBoolean restarting = new AtomicBoolean(false);

@Getter protected volatile boolean startFailedButDevelopmentMode = false;
protected AgentServiceGrpc.AgentServiceStub asyncStub;

protected CompletableFuture<StreamObserver<TopicProducerWriteResult>>
topicProducerWriteResults = CompletableFuture.completedFuture(null);

private final Map<String, TopicProducer> topicProducers = new ConcurrentHashMap<>();

protected record GrpcAgentRecord(
Long id,
Expand All @@ -90,6 +100,93 @@ public void start() throws Exception {
}
blockingStub =
AgentServiceGrpc.newBlockingStub(channel).withDeadlineAfter(30, TimeUnit.SECONDS);
asyncStub = AgentServiceGrpc.newStub(channel).withWaitForReady();

topicProducerWriteResults = new CompletableFuture<>();
topicProducerWriteResults.complete(
asyncStub.getTopicProducerRecords(
new StreamObserver<>() {
@Override
public void onNext(TopicProducerRecord topicProducerRecord) {
TopicProducer topicProducer =
topicProducers.computeIfAbsent(
topicProducerRecord.getTopic(),
topic -> {
TopicProducer tp =
agentContext
.getTopicConnectionProvider()
.createProducer(
agentContext
.getGlobalAgentId(),
topic,
Map.of());
tp.start();
return tp;
});
try {
topicProducer
.write(fromGrpc(topicProducerRecord.getRecord()))
.whenComplete(
(r, e) -> {
if (e != null) {
log.error("Error writing record", e);
sendTopicProducerWriteResult(
TopicProducerWriteResult
.newBuilder()
.setError(
e
.getMessage()));
} else {
sendTopicProducerWriteResult(
TopicProducerWriteResult
.newBuilder()
.setRecordId(
topicProducerRecord
.getRecord()
.getRecordId()));
}
});
} catch (IOException e) {
agentContext.criticalFailure(e);
}
}

@Override
public void onError(Throwable throwable) {
if (!restarting.get()) {
agentContext.criticalFailure(
new RuntimeException(
"getTopicProducerRecords: gRPC server sent error: %s"
.formatted(throwable.getMessage()),
throwable));
} else {
log.info(
"getTopicProducerRecords: ignoring error during restart {}",
throwable + "");
}
}

@Override
public void onCompleted() {
if (!restarting.get()) {
agentContext.criticalFailure(
new RuntimeException(
"getTopicProducerRecords: gRPC server completed the stream unexpectedly"));
} else {
log.info(
"getTopicProducerRecords: ignoring error server stop during restart");
}
}
}));
}

private synchronized void sendTopicProducerWriteResult(
TopicProducerWriteResult.Builder result) {
try {
topicProducerWriteResults.get().onNext(result.build());
} catch (InterruptedException | ExecutionException e) {
throw new RuntimeException(e);
}
}

@Override
Expand All @@ -103,6 +200,16 @@ protected Map<String, Object> buildAdditionalInfo() {
}

protected synchronized void stopBeforeRestart() throws Exception {
restarting.set(true);
StreamObserver<TopicProducerWriteResult> topicProducerWriteResultStreamObserver =
topicProducerWriteResults.get();
if (topicProducerWriteResultStreamObserver != null) {
try {
topicProducerWriteResultStreamObserver.onCompleted();
} catch (IllegalStateException e) {
log.info("Ignoring error while stopping {}", e + "");
}
}
stopChannel(false);
}

Expand All @@ -123,6 +230,10 @@ public void stopChannel(boolean wait) throws Exception {
public synchronized void close() throws Exception {
stopBeforeRestart();
stopChannel(true);
for (TopicProducer topicProducer : topicProducers.values()) {
topicProducer.close();
}
topicProducers.clear();
}

protected Object fromGrpc(Value value) throws IOException {
Expand Down
11 changes: 11 additions & 0 deletions langstream-agents/langstream-agent-grpc/src/main/proto/langstream_grpc/proto/agent.proto
View file
Original file line number Diff line number Diff line change
Expand Up @@ -26,6 +26,7 @@ service AgentService {
rpc read(stream SourceRequest) returns (stream SourceResponse) {}
rpc process(stream ProcessorRequest) returns (stream ProcessorResponse) {}
rpc write(stream SinkRequest) returns (stream SinkResponse) {}
rpc get_topic_producer_records(stream TopicProducerWriteResult) returns (stream TopicProducerRecord) {}
}

message InfoResponse {
Expand Down Expand Up @@ -68,6 +69,16 @@ message Record {
optional int64 timestamp = 6;
}

message TopicProducerWriteResult {
int64 record_id = 1;
optional string error = 2;
}

message TopicProducerRecord {
string topic = 1;
Record record = 3;
}

message PermanentFailure {
int64 record_id = 1;
string error_message = 2;
Expand Down
Loading

0 comments on commit b16e58d

Please sign in to comment.