Suppress capturing responses

TrafficCapture/nettyWireLogging/src/main/java/org/opensearch/migrations/trafficcapture/netty/ConditionallyReliableLoggingHttpRequestHandler.java → TrafficCapture/nettyWireLogging/src/main/java/org/opensearch/migrations/trafficcapture/netty/ConditionallyReliableLoggingHttpHandler.java

@@ -3,31 +3,23 @@
 import io.netty.channel.ChannelHandlerContext;
 import io.netty.handler.codec.http.HttpRequest;
 import io.netty.util.ReferenceCountUtil;
-import io.opentelemetry.api.trace.Span;
-import lombok.Getter;
 import lombok.Lombok;
 import lombok.NonNull;
 import lombok.extern.slf4j.Slf4j;
-import org.opensearch.migrations.tracing.IWithStartTimeAndAttributes;
-import org.opensearch.migrations.tracing.commoncontexts.IConnectionContext;
-import org.opensearch.migrations.trafficcapture.IChannelConnectionCaptureSerializer;
 import org.opensearch.migrations.trafficcapture.IConnectionCaptureFactory;
 import org.opensearch.migrations.trafficcapture.netty.tracing.HttpMessageContext;
-import org.opensearch.migrations.trafficcapture.tracing.ConnectionContext;
 
 import java.io.IOException;
-import java.time.Instant;
-import java.util.function.Function;
 import java.util.function.Predicate;
 
 @Slf4j
-public class ConditionallyReliableLoggingHttpRequestHandler<T> extends LoggingHttpRequestHandler<T> {
+public class ConditionallyReliableLoggingHttpHandler<T> extends LoggingHttpHandler<T> {
     private final Predicate<HttpRequest> shouldBlockPredicate;
 
-    public ConditionallyReliableLoggingHttpRequestHandler(@NonNull String nodeId, String connectionId,
-                                                          @NonNull IConnectionCaptureFactory<T> trafficOffloaderFactory,
-                                                          @NonNull RequestCapturePredicate requestCapturePredicate,
-                                                          @NonNull Predicate<HttpRequest> headerPredicateForWhenToBlock)
+    public ConditionallyReliableLoggingHttpHandler(@NonNull String nodeId, String connectionId,
+                                                   @NonNull IConnectionCaptureFactory<T> trafficOffloaderFactory,
+                                                   @NonNull RequestCapturePredicate requestCapturePredicate,
+                                                   @NonNull Predicate<HttpRequest> headerPredicateForWhenToBlock)
     throws IOException {
         super(nodeId, connectionId, trafficOffloaderFactory, requestCapturePredicate);
         this.shouldBlockPredicate = headerPredicateForWhenToBlock;
@@ -52,7 +44,7 @@
                     // This is a spot where we would benefit from having a behavioral policy that different users
                     // could set as needed. Some users may be fine with just logging a failed offloading of a request
                     // where other users may want to stop entirely. JIRA here: https://opensearch.atlassian.net/browse/MIGRATIONS-1276
                     log.atWarn().setCause(t).setMessage("Dropping request - Got error").log();
                     ReferenceCountUtil.release(msg);
                 } else {
                     try {

TrafficCapture/nettyWireLogging/src/main/java/org/opensearch/migrations/trafficcapture/netty/LoggingHttpRequestHandler.java → TrafficCapture/nettyWireLogging/src/main/java/org/opensearch/migrations/trafficcapture/netty/LoggingHttpHandler.java

@@ -32,7 +32,7 @@
 import java.time.Instant;
 
 @Slf4j
-public class LoggingHttpRequestHandler<T> extends ChannelDuplexHandler {
+public class LoggingHttpHandler<T> extends ChannelDuplexHandler {
     public static final String TELEMETRY_SCOPE_NAME = "CapturingHttpHandler";
     public static final SimpleMeteringClosure METERING_CLOSURE = new SimpleMeteringClosure(TELEMETRY_SCOPE_NAME);
     private static final MetricsLogger metricsLogger = new MetricsLogger("LoggingHttpRequestHandler");
@@ -41,11 +41,40 @@
     public static final String GATHERING_RESPONSE = "gatheringResponse";
     public static final String BLOCKED = "blocked";
 
+    static class CaptureIgnoreState {
+        static final byte CAPTURE = 0;
+        static final byte IGNORE_REQUEST = 1;
+        static final byte IGNORE_RESPONSE = 2;
+        private CaptureIgnoreState() {}
+    }
+
+    static class CaptureState {
+        byte captureIgnoreState = CaptureIgnoreState.CAPTURE;
+        boolean liveReadObservationsInOffloader = false;
+
+        boolean shouldCapture() {
+            return captureIgnoreState == CaptureIgnoreState.CAPTURE;
+        }
+
+        public void setShouldCaptureForRequest(boolean b) {
+            captureIgnoreState = b ? CaptureIgnoreState.CAPTURE : CaptureIgnoreState.IGNORE_REQUEST;
+        }
+
+        public void advanceStateModelIntoResponseGather() {
+            if (CaptureIgnoreState.CAPTURE != captureIgnoreState) {
+                captureIgnoreState = CaptureIgnoreState.IGNORE_RESPONSE;
+            }
+        }
+    }
+
     static class SimpleHttpRequestDecoder extends HttpRequestDecoder {
         private final PassThruHttpHeaders.HttpHeadersToPreserve headersToPreserve;
+        private final CaptureState captureState;
 
-        public SimpleHttpRequestDecoder(@NonNull PassThruHttpHeaders.HttpHeadersToPreserve headersToPreserve) {
+        public SimpleHttpRequestDecoder(@NonNull PassThruHttpHeaders.HttpHeadersToPreserve headersToPreserve,
+                                        CaptureState captureState) {
             this.headersToPreserve = headersToPreserve;
+            this.captureState = captureState;
         }
 
         /**
@@ -60,46 +89,50 @@
                     , new PassThruHttpHeaders(headersToPreserve)
             );
         }
-    }
 
+        @Override
+        public void channelRead(ChannelHandlerContext ctx, Object msg) throws Exception {
+            if (captureState.captureIgnoreState == CaptureIgnoreState.IGNORE_RESPONSE) {
+                captureState.captureIgnoreState = CaptureIgnoreState.CAPTURE;
+            }
+            super.channelRead(ctx, msg);
+        }
+    }
+
     static class SimpleDecodedHttpRequestHandler extends ChannelInboundHandlerAdapter {
         @Getter
         private HttpRequest currentRequest;
         final RequestCapturePredicate requestCapturePredicate;
-        boolean isDone;
-        boolean shouldCapture;
-        boolean liveReadObservationsInOffloader;
-
-        SimpleDecodedHttpRequestHandler(RequestCapturePredicate requestCapturePredicate) {
+        boolean haveParsedFullRequest;
+        final CaptureState captureState;
+
+        SimpleDecodedHttpRequestHandler(RequestCapturePredicate requestCapturePredicate, CaptureState captureState) {
             this.requestCapturePredicate = requestCapturePredicate;
             this.currentRequest = null;
-            this.isDone = false;
-            this.shouldCapture = true;
-            liveReadObservationsInOffloader = false;
+            this.haveParsedFullRequest = false;
+            this.captureState = captureState;
         }
 
         @Override
         public void channelRead(@NonNull ChannelHandlerContext ctx, @NonNull Object msg) throws Exception {
             if (msg instanceof HttpRequest) {
                 currentRequest = (HttpRequest) msg;
-                shouldCapture = RequestCapturePredicate.CaptureDirective.CAPTURE ==
-                        requestCapturePredicate.apply((HttpRequest) msg);
+                captureState.setShouldCaptureForRequest(RequestCapturePredicate.CaptureDirective.CAPTURE ==
+                        requestCapturePredicate.apply((HttpRequest) msg));
             } else if (msg instanceof HttpContent) {
                 ((HttpContent)msg).release();
                 if (msg instanceof LastHttpContent) {
-                    isDone = true;
+                    haveParsedFullRequest = true;
                 }
             } else {
                 super.channelRead(ctx, msg);
             }
         }
 
         public HttpRequest resetCurrentRequest() {
-            this.shouldCapture = true;
-            this.isDone = false;
+            this.haveParsedFullRequest = false;
             var old = currentRequest;
             this.currentRequest = null;
-            this.liveReadObservationsInOffloader = false;
             return old;
         }
     }
@@ -110,9 +143,9 @@
 
     protected HttpMessageContext messageContext;
 
-    public LoggingHttpRequestHandler(String nodeId, String channelKey,
-                                     @NonNull IConnectionCaptureFactory<T> trafficOffloaderFactory,
-                                     @NonNull RequestCapturePredicate httpHeadersCapturePredicate)
+    public LoggingHttpHandler(String nodeId, String channelKey,
+                              @NonNull IConnectionCaptureFactory<T> trafficOffloaderFactory,
+                              @NonNull RequestCapturePredicate httpHeadersCapturePredicate)
     throws IOException {
         var parentContext = new ConnectionContext(channelKey, nodeId,
                 METERING_CLOSURE.makeSpanContinuation("connectionLifetime", null));
@@ -122,13 +155,13 @@
         METERING_CLOSURE.meterIncrementEvent(messageContext, "requestStarted");
 
         this.trafficOffloader = trafficOffloaderFactory.createOffloader(parentContext, channelKey);
+        var captureState = new CaptureState();
         httpDecoderChannel = new EmbeddedChannel(
-                new SimpleHttpRequestDecoder(httpHeadersCapturePredicate.getHeadersRequiredForMatcher()),
-                new SimpleDecodedHttpRequestHandler(httpHeadersCapturePredicate)
+                new SimpleHttpRequestDecoder(httpHeadersCapturePredicate.getHeadersRequiredForMatcher(), captureState),
+                new SimpleDecodedHttpRequestHandler(httpHeadersCapturePredicate, captureState)
         );
     }
 
-
     static String getSpanLabelForState(HttpMessageContext.HttpTransactionState state) {
         switch (state) {
             case REQUEST:
@@ -140,7 +173,7 @@
             case RESPONSE:
                 return GATHERING_RESPONSE;
             default:
                 throw new IllegalStateException("Unknown enum value: "+state);
         }
     }
 
@@ -162,13 +195,13 @@
         METERING_CLOSURE.meterIncrementEvent(messageContext, "unregistered");
         trafficOffloader.flushCommitAndResetStream(true).whenComplete((result, t) -> {
             if (t != null) {
                 log.warn("Got error: " + t.getMessage());
                 ctx.close();
             } else {
                 try {
                     super.channelUnregistered(ctx);
                 } catch (Exception e) {
                     throw Lombok.sneakyThrow(e);
                 }
             }
         });
@@ -182,12 +215,12 @@
 
         trafficOffloader.flushCommitAndResetStream(true).whenComplete((result, t) -> {
             if (t != null) {
                 log.warn("Got error: " + t.getMessage());
             }
             try {
                 super.channelUnregistered(ctx);
             } catch (Exception e) {
                 throw Lombok.sneakyThrow(e);
             }
         });
         super.handlerRemoved(ctx);
@@ -217,16 +250,16 @@
         httpDecoderChannel.writeInbound(bb.retainedDuplicate()); // the ByteBuf is consumed/release by this method
 
         METERING_CLOSURE.meterIncrementEvent(messageContext,
-                getHandlerThatHoldsParsedHttpRequest().isDone ? "requestFullyParsed" : "requestPartiallyParsed");
+                getHandlerThatHoldsParsedHttpRequest().haveParsedFullRequest ? "requestFullyParsed" : "requestPartiallyParsed");
 
-        var shouldCapture = requestParsingHandler.shouldCapture;
+        var captureState = requestParsingHandler.captureState;
+        var shouldCapture = captureState.shouldCapture();
         if (shouldCapture) {
-            requestParsingHandler.liveReadObservationsInOffloader = true;
+            captureState.liveReadObservationsInOffloader = true;
             trafficOffloader.addReadEvent(timestamp, bb);
-
-        } else if (requestParsingHandler.liveReadObservationsInOffloader) {
+        } else if (captureState.liveReadObservationsInOffloader) {
             trafficOffloader.cancelCaptureForCurrentRequest(timestamp);
-            requestParsingHandler.liveReadObservationsInOffloader = false;
+            captureState.liveReadObservationsInOffloader = false;
         }
 
         metricsLogger.atSuccess(MetricsEvent.RECEIVED_REQUEST_COMPONENT)
@@ -235,9 +268,12 @@
         METERING_CLOSURE.meterIncrementEvent(messageContext, "readBytes", bb.readableBytes());
 
 
-        if (requestParsingHandler.isDone) {
+        if (requestParsingHandler.haveParsedFullRequest) {
             messageContext.getCurrentSpan().end();
             var httpRequest = requestParsingHandler.resetCurrentRequest();
+            captureState.liveReadObservationsInOffloader = false;
+            captureState.advanceStateModelIntoResponseGather();
+
             if (shouldCapture) {
                 var decoderResultLoose = httpRequest.decoderResult();
                 if (decoderResultLoose instanceof HttpMessageDecoderResult) {
@@ -260,7 +296,9 @@
             rotateNextMessageContext(HttpMessageContext.HttpTransactionState.RESPONSE);
         }
         var bb = (ByteBuf) msg;
-        trafficOffloader.addWriteEvent(Instant.now(), bb);
+        if (getHandlerThatHoldsParsedHttpRequest().captureState.shouldCapture()) {
+            trafficOffloader.addWriteEvent(Instant.now(), bb);
+        }
         metricsLogger.atSuccess(MetricsEvent.RECEIVED_RESPONSE_COMPONENT)
                 .setAttribute(MetricsAttributeKey.CHANNEL_ID, ctx.channel().id().asLongText()).emit();
         METERING_CLOSURE.meterIncrementEvent(messageContext, "write");
@@ -271,10 +309,10 @@
 
     @Override
     public void exceptionCaught(ChannelHandlerContext ctx, Throwable cause) throws Exception {
         trafficOffloader.addExceptionCaughtEvent(Instant.now(), cause);
         METERING_CLOSURE.meterIncrementEvent(messageContext, "exception");
         httpDecoderChannel.close();
         super.exceptionCaught(ctx, cause);
     }
 
 }