Notes on a pipeline service - krickert/search-api GitHub Wiki

package com.krickert.search.pipeline.config;

import io.micronaut.context.annotation.ConfigurationProperties;
import jakarta.inject.Singleton;
import lombok.Data;
import lombok.Getter;

import java.util.List;
import java.util.Map;

@Data
@Singleton
@ConfigurationProperties("pipeline")
public class PipelineConfig {

    private final Map<String, ServiceConfiguration> service;

    public PipelineConfig(Map<String, ServiceConfiguration> service) {
        this.service = service;
    }

    public ServiceConfiguration getRouteForService(String serviceName) {
        return service.get(serviceName);
    }

    public String getServiceGrpcEndpoint(String serviceName) {
        return getRouteForService(serviceName).getGrpcEndpoint();
    }
    public List<String> getServiceKafkaListenTopics(String serviceName) {
        return getRouteForService(serviceName).getKafkaListenTopics();
    }
    public List<String> getServiceKafkaPublishTopics(String serviceName) {
        return getRouteForService(serviceName).getKafkaPublishTopics();
    }

    public List<String> getServiceGrpcForwardTo(String serviceName) {
        return getRouteForService(serviceName).getGrpcForwardTo();
    }
}

package com.krickert.search.pipeline.config;

import io.micronaut.context.annotation.EachProperty;
import io.micronaut.context.annotation.Parameter;
import jakarta.inject.Singleton;
import lombok.Data;
import lombok.Getter;
import lombok.Setter;

import java.util.List;

@Data
@EachProperty("pipeline.service")
@Singleton
public class ServiceConfiguration {

    private final String name;
    private String grpcEndpoint;
    private List<String> kafkaListenTopics;
    private List<String> kafkaPublishTopics;
    private List<String> grpcForwardTo;

    public ServiceConfiguration(@Parameter String name) {
        this.name = name;
    }
}

package com.krickert.search.pipeline;

import io.micronaut.configuration.kafka.annotation.KafkaClient;
import io.micronaut.context.annotation.Value;
import io.micronaut.grpc.annotation.GrpcService;
import io.micronaut.http.annotation.Controller;
import io.micronaut.messaging.annotation.Topic;
import io.reactivex.Single;
import jakarta.inject.Inject;
import lombok.extern.slf4j.Slf4j;
import reactor.core.publisher.Mono;
import reactor.core.publisher.Flux;
import com.krickert.search.model.pipe.*;

import java.util.List;
import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;

@Slf4j
@GrpcService
public abstract class AbstractPipelineService extends PipelineServiceGrpc.PipelineServiceImplBase {

    @Inject private KafkaProducer kafkaProducer;
    @Inject private GrpcForwarder grpcForwarder;
    @Inject private PipelineConfig pipelineConfig;

    @Value("${pipeline.service.name}")
    private String serviceName; // Inject the service name for self-registration

    private static final Map<String, RouteConfig> routeRegistry = new ConcurrentHashMap<>();

    public AbstractPipelineService() {
        log.info("Initializing pipeline service: {}", serviceName);
    }

    /**
     * The only method users need to implement to process documents.
     */
    protected abstract PipeDoc processPipe(PipeDoc doc);

    @Override
    public void forward(PipeStream request, io.grpc.stub.StreamObserver<com.google.protobuf.Empty> responseObserver) {
        log.info("Received request for forward processing, processing {} documents", request.getRequestCount());

        Flux.fromIterable(request.getPipeRepliesList())
            .flatMap(reply -> {
                PipeDoc enhancedDoc;
                try {
                    enhancedDoc = processPipe(reply.getRequest().getDoc());
                } catch (Exception e) {
                    log.error("Processing failed: {}", e.getMessage(), e);
                    return handleError(reply.getRequest(), e.getMessage());
                }
                return sendToDestinations(reply.getRequest(), enhancedDoc);
            })
            .doOnTerminate(() -> responseObserver.onNext(com.google.protobuf.Empty.getDefaultInstance()))
            .subscribe();
    }

    private Mono<PipeStream> sendToDestinations(PipeRequest originalRequest, PipeDoc updatedDoc) {
        RouteConfig routes = pipelineConfig.getRouteForService(serviceName);

        if (routes == null) {
            log.error("No configured routes found for service: {}", serviceName);
            return Mono.empty();
        }

        log.info("Routing document {} to destinations", updatedDoc.getId());

        // Send to Kafka
        routes.getKafkaTopics().forEach(topic -> kafkaProducer.send(topic, updatedDoc));

        // Send to gRPC destinations
        routes.getGrpcEndpoints().forEach(endpoint -> grpcForwarder.send(endpoint, updatedDoc));

        return Mono.just(PipeStream.newBuilder()
                .addPipeReplies(PipeResponse.newBuilder().setSuccess(true).build())
                .build());
    }

    private Mono<PipeStream> handleError(PipeRequest request, String errorMessage) {
        RouteConfig routes = pipelineConfig.getRouteForService(serviceName);

        ErrorData errorData = ErrorData.newBuilder()
                .setErrorMessage(errorMessage)
                .addAllFailedRoutes(request.getDestinationsList())
                .setErrorRequest(request)
                .build();

        PipeResponse errorResponse = PipeResponse.newBuilder()
                .setSuccess(false)
                .setErrorDate(errorData)
                .build();

        log.error("Routing error to failure queue");
        routes.getKafkaFailureTopics().forEach(topic -> kafkaProducer.send(topic, request.getDoc()));

        return Mono.just(PipeStream.newBuilder()
                .addPipeReplies(errorResponse)
                .build());
    }
}

// PipelineServiceImpl.java
package com.example.pipeline;

import io.grpc.stub.StreamObserver;
import io.micronaut.configuration.grpc.server.GrpcService;
import jakarta.inject.Inject;
import java.util.List;

@GrpcService
public class PipelineServiceImpl extends PipelineServiceGrpc.PipelineServiceImplBase {

    @Inject
    KafkaForwarder kafkaForwarder;

    @Inject
    GrpcForwarder grpcForwarder;

    @Override
    public StreamObserver<Pipe> forward(StreamObserver<PipeResponse> responseObserver) {
        return new StreamObserver<Pipe>() {
            @Override
            public void onNext(Pipe pipe) {
                PipeResponse response = processPipe(pipe);
                responseObserver.onNext(response);
            }

            @Override
            public void onError(Throwable t) {
                // Log error as needed and propagate the error
                responseObserver.onError(t);
            }

            @Override
            public void onCompleted() {
                responseObserver.onCompleted();
            }
        };
    }

    @Override
    public void getOutput(Pipe request, StreamObserver<PipeResponse> responseObserver) {
        PipeResponse response = processPipe(request);
        responseObserver.onNext(response);
        responseObserver.onCompleted();
    }

    private PipeResponse processPipe(Pipe pipe) {
        PipeResponse.Builder responseBuilder = PipeResponse.newBuilder();
        List<Route> routes = pipe.getRoutesList();

        for (Route route : routes) {
            try {
                switch (route.getRouteType()) {
                    case KAFKA:
                        kafkaForwarder.forwardToKafka(pipe, route);
                        break;
                    case GRPC:
                        grpcForwarder.forwardToGrpc(pipe, route);
                        break;
                    case NULL_TERMINATION:
                        // Do nothing: this route terminates forwarding.
                        break;
                    default:
                        throw new IllegalArgumentException("Unsupported route type: " + route.getRouteType());
                }
            } catch (Exception e) {
                // Record the failed route (and log the error)
                responseBuilder.addFailedRoutes(route);
                // Optionally: forward the same pipe to a backup Kafka topic for reprocessing
                try {
                    kafkaForwarder.forwardToBackup(pipe, route);
                } catch (Exception backupEx) {
                    // Log backup failure if needed
                }
            }
        }

        boolean success = responseBuilder.getFailedRoutesCount() == 0;
        responseBuilder.setSuccess(success);
        if (!success) {
            responseBuilder.setErrorMessage("Some routes failed during processing.");
        }
        return responseBuilder.build();
    }
}

// KafkaForwarder.java
package com.example.pipeline;

import jakarta.inject.Inject;
import jakarta.inject.Singleton;

@Singleton
public class KafkaForwarder {

    @Inject
    KafkaForwarderClient kafkaForwarderClient;

    public void forwardToKafka(Pipe pipe, Route route) {
        // The 'destination' field contains the Kafka topic name.
        String topic = route.getDestination();
        kafkaForwarderClient.send(topic, pipe);
    }

    public void forwardToBackup(Pipe pipe, Route route) {
        // Use a backup topic (e.g. prefix with "backup-") for reprocessing failed messages.
        String backupTopic = "backup-" + route.getDestination();
        kafkaForwarderClient.send(backupTopic, pipe);
    }
}

// KafkaForwarderClient.java
package com.example.pipeline;

import io.micronaut.configuration.kafka.annotation.KafkaClient;

@KafkaClient
public interface KafkaForwarderClient {
    // This method sends a Pipe message to the given topic.
    void send(String topic, Pipe pipe);
}

// GrpcForwarder.java
package com.example.pipeline;

import io.grpc.ManagedChannel;
import io.grpc.ManagedChannelBuilder;
import jakarta.inject.Singleton;
import jakarta.inject.Inject;

@Singleton
public class GrpcForwarder {

    // For demonstration we create a single stub.
    // In practice, you may want to select the stub based on route.getDestination().
    private final PipelineServiceGrpc.PipelineServiceBlockingStub stub;

    @Inject
    public GrpcForwarder() {
        // Create a channel to the destination gRPC service (address could be externalized to config)
        ManagedChannel channel = ManagedChannelBuilder.forAddress("localhost", 50051)
                .usePlaintext()
                .build();
        this.stub = PipelineServiceGrpc.newBlockingStub(channel);
    }

    public void forwardToGrpc(Pipe pipe, Route route) {
        // In a real-world scenario, use route.getDestination() to choose the correct stub.
        // Here we simply call the getOutput method and ignore the response.
        stub.getOutput(pipe);
    }
}

// PipelineKafkaListener.java
package com.example.pipeline;

import io.grpc.stub.StreamObserver;
import io.micronaut.configuration.kafka.annotation.KafkaListener;
import io.micronaut.configuration.kafka.annotation.Topic;
import jakarta.inject.Inject;
import jakarta.inject.Singleton;

@KafkaListener(groupId = "pipe-group")
@Singleton
public class PipelineKafkaListener {

    @Inject
    PipelineServiceImpl pipelineService;

    @Topic("input-pipe")
    public void receive(Pipe pipe) {
        // Process the pipe using the same logic as in the gRPC service.
        pipelineService.getOutput(pipe, new StreamObserver<PipeResponse>() {
            @Override
            public void onNext(PipeResponse response) {
                // For example, log the response or perform additional processing.
            }

            @Override
            public void onError(Throwable t) {
                // Log error as needed.
            }

            @Override
            public void onCompleted() {
                // Completion logic if needed.
            }
        });
    }
}