Semantic Query may throw "too many rewrite rounds"

[benwtrent]

Elasticsearch Version

current

Installed Plugins

No response

Java Version

bundled

OS Version

any

Problem Description

Seen in production, SemanticQueryBuilder will sometimes throw IllegalStateException due to too many rewrite rounds

Steps to Reproduce

N/A

Logs (if relevant)

Here is the full trace:

java.lang.IllegalStateException: too many rewrite rounds, rewriteable might return new objects even if they are not rewritten
	at org.elasticsearch.server@9.3.0/org.elasticsearch.index.query.Rewriteable.rewriteAndFetch(Rewriteable.java:96)
	at org.elasticsearch.server@9.3.0/org.elasticsearch.index.query.Rewriteable.lambda$rewriteAndFetch$0(Rewriteable.java:104)
	at org.elasticsearch.server@9.3.0/org.elasticsearch.action.ActionListenerImplementations$ResponseWrappingActionListener.onResponse(ActionListenerImplementations.java:261)
	at org.elasticsearch.server@9.3.0/org.elasticsearch.index.query.QueryRewriteContext$1.onResponse(QueryRewriteContext.java:364)
	at org.elasticsearch.inference@9.3.0/org.elasticsearch.xpack.inference.queries.SemanticQueryBuilder.lambda$registerInferenceAsyncAction$8(SemanticQueryBuilder.java:343)
	at org.elasticsearch.server@9.3.0/org.elasticsearch.action.ActionListenerImplementations$ResponseWrappingActionListener.onResponse(ActionListenerImplementations.java:261)
	at org.elasticsearch.server@9.3.0/org.elasticsearch.action.support.ContextPreservingActionListener.onResponse(ContextPreservingActionListener.java:33)
	at org.elasticsearch.server@9.3.0/org.elasticsearch.tasks.TaskManager$1.onResponse(TaskManager.java:228)
	at org.elasticsearch.server@9.3.0/org.elasticsearch.tasks.TaskManager$1.onResponse(TaskManager.java:222)
	at org.elasticsearch.server@9.3.0/org.elasticsearch.action.ActionListenerImplementations$RunBeforeActionListener.onResponse(ActionListenerImplementations.java:350)
	at org.elasticsearch.server@9.3.0/org.elasticsearch.action.support.ContextPreservingActionListener.onResponse(ContextPreservingActionListener.java:33)
	at org.elasticsearch.server@9.3.0/org.elasticsearch.action.ActionListenerImplementations$MappedActionListener.onResponse(ActionListenerImplementations.java:111)
	at org.elasticsearch.inference@9.3.0/org.elasticsearch.xpack.inference.action.BaseTransportInferenceAction.lambda$inferOnServiceWithMetrics$7(BaseTransportInferenceAction.java:214)
	at org.elasticsearch.server@9.3.0/org.elasticsearch.action.ActionListener$2.onResponse(ActionListener.java:258)
	at org.elasticsearch.inference@9.3.0/org.elasticsearch.xpack.inference.services.elasticsearch.ElasticsearchInternalService.lambda$inferTextEmbedding$5(ElasticsearchInternalService.java:648)
	at org.elasticsearch.server@9.3.0/org.elasticsearch.action.ActionListenerImplementations$ResponseWrappingActionListener.onResponse(ActionListenerImplementations.java:261)
	at org.elasticsearch.server@9.3.0/org.elasticsearch.action.ActionListenerImplementations$DelegatingResponseActionListener.onResponse(ActionListenerImplementations.java:198)
	at org.elasticsearch.server@9.3.0/org.elasticsearch.action.support.ContextPreservingActionListener.onResponse(ContextPreservingActionListener.java:33)
	at org.elasticsearch.server@9.3.0/org.elasticsearch.tasks.TaskManager$1.onResponse(TaskManager.java:228)
	at org.elasticsearch.server@9.3.0/org.elasticsearch.tasks.TaskManager$1.onResponse(TaskManager.java:222)
	at org.elasticsearch.server@9.3.0/org.elasticsearch.action.ActionListenerImplementations$RunBeforeActionListener.onResponse(ActionListenerImplementations.java:350)
	at org.elasticsearch.server@9.3.0/org.elasticsearch.action.support.ContextPreservingActionListener.onResponse(ContextPreservingActionListener.java:33)
	at org.elasticsearch.server@9.3.0/org.elasticsearch.action.ActionListenerImplementations$MappedActionListener.onResponse(ActionListenerImplementations.java:111)
	at org.elasticsearch.ml@9.3.0/org.elasticsearch.xpack.ml.action.TransportInternalInferModelAction$1.sendResponse(TransportInternalInferModelAction.java:449)
	at org.elasticsearch.ml@9.3.0/org.elasticsearch.xpack.ml.action.TransportInternalInferModelAction$1.onResponse(TransportInternalInferModelAction.java:417)
	at org.elasticsearch.ml@9.3.0/org.elasticsearch.xpack.ml.action.TransportInternalInferModelAction$1.onResponse(TransportInternalInferModelAction.java:412)
	at org.elasticsearch.server@9.3.0/org.elasticsearch.action.support.ContextPreservingActionListener.onResponse(ContextPreservingActionListener.java:33)
	at org.elasticsearch.server@9.3.0/org.elasticsearch.tasks.TaskManager$1.onResponse(TaskManager.java:228)
	at org.elasticsearch.server@9.3.0/org.elasticsearch.tasks.TaskManager$1.onResponse(TaskManager.java:222)
	at org.elasticsearch.server@9.3.0/org.elasticsearch.action.ActionListenerImplementations$RunBeforeActionListener.onResponse(ActionListenerImplementations.java:350)
	at org.elasticsearch.server@9.3.0/org.elasticsearch.action.support.ContextPreservingActionListener.onResponse(ContextPreservingActionListener.java:33)
	at org.elasticsearch.server@9.3.0/org.elasticsearch.action.ActionListenerImplementations$MappedActionListener.onResponse(ActionListenerImplementations.java:111)
	at org.elasticsearch.server@9.3.0/org.elasticsearch.action.support.SubscribableListener$SuccessResult.complete(SubscribableListener.java:406)
	at org.elasticsearch.server@9.3.0/org.elasticsearch.action.support.SubscribableListener.tryComplete(SubscribableListener.java:326)
	at org.elasticsearch.server@9.3.0/org.elasticsearch.action.support.SubscribableListener.addListener(SubscribableListener.java:222)
	at org.elasticsearch.server@9.3.0/org.elasticsearch.action.support.SubscribableListener.addListener(SubscribableListener.java:180)
	at org.elasticsearch.server@9.3.0/org.elasticsearch.action.support.CancellableFanOut$SubtasksCompletionHandler.run(CancellableFanOut.java:202)
	at org.elasticsearch.base@9.3.0/org.elasticsearch.core.AbstractRefCounted$1.closeInternal(AbstractRefCounted.java:125)
	at org.elasticsearch.base@9.3.0/org.elasticsearch.core.AbstractRefCounted.decRef(AbstractRefCounted.java:77)
	at org.elasticsearch.server@9.3.0/org.elasticsearch.action.support.RefCountingRunnable.close(RefCountingRunnable.java:113)
	at org.elasticsearch.base@9.3.0/org.elasticsearch.core.Releasables.close(Releasables.java:39)
	at org.elasticsearch.base@9.3.0/org.elasticsearch.core.Releasables.closeExpectNoException(Releasables.java:72)
	at org.elasticsearch.server@9.3.0/org.elasticsearch.action.ActionListenerImplementations$2.run(ActionListenerImplementations.java:51)
	at org.elasticsearch.server@9.3.0/org.elasticsearch.action.ActionListenerImplementations$RunAfterActionListener.onResponse(ActionListenerImplementations.java:314)
	at org.elasticsearch.server@9.3.0/org.elasticsearch.action.ActionListenerResponseHandler.handleResponse(ActionListenerResponseHandler.java:49)
	at org.elasticsearch.server@9.3.0/org.elasticsearch.transport.TransportService$UnregisterChildTransportResponseHandler.handleResponse(TransportService.java:1726)
	at org.elasticsearch.server@9.3.0/org.elasticsearch.transport.TransportService$ContextRestoreResponseHandler.handleResponse(TransportService.java:1504)
	at org.elasticsearch.server@9.3.0/org.elasticsearch.transport.InboundHandler.doHandleResponse(InboundHandler.java:465)
	at org.elasticsearch.server@9.3.0/org.elasticsearch.transport.InboundHandler.handleResponse(InboundHandler.java:419)
	at org.elasticsearch.server@9.3.0/org.elasticsearch.transport.InboundHandler.executeResponseHandler(InboundHandler.java:166)
	at org.elasticsearch.server@9.3.0/org.elasticsearch.transport.InboundHandler.messageReceived(InboundHandler.java:134)
	at org.elasticsearch.server@9.3.0/org.elasticsearch.transport.InboundHandler.inboundMessage(InboundHandler.java:102)
	at org.elasticsearch.server@9.3.0/org.elasticsearch.transport.TcpTransport.inboundMessage(TcpTransport.java:830)
	at org.elasticsearch.server@9.3.0/org.elasticsearch.transport.InboundPipeline.forwardFragment(InboundPipeline.java:116)
	at org.elasticsearch.server@9.3.0/org.elasticsearch.transport.InboundPipeline.lambda$doHandleBytes$1(InboundPipeline.java:78)
	at org.elasticsearch.server@9.3.0/org.elasticsearch.transport.InboundDecoder.finishMessage(InboundDecoder.java:144)
	at org.elasticsearch.server@9.3.0/org.elasticsearch.transport.InboundDecoder.internalDecode(InboundDecoder.java:129)
	at org.elasticsearch.server@9.3.0/org.elasticsearch.transport.InboundDecoder.decode(InboundDecoder.java:58)
	at org.elasticsearch.server@9.3.0/org.elasticsearch.transport.InboundPipeline.doHandleBytes(InboundPipeline.java:79)
	at org.elasticsearch.server@9.3.0/org.elasticsearch.transport.InboundPipeline.handleBytes(InboundPipeline.java:69)
	at org.elasticsearch.transport.netty4@9.3.0/org.elasticsearch.transport.netty4.Netty4MessageInboundHandler.channelRead(Netty4MessageInboundHandler.java:54)
	at io.netty.transport@4.1.126.Final/io.netty.channel.AbstractChannelHandlerContext.invokeChannelRead(AbstractChannelHandlerContext.java:444)
	at io.netty.transport@4.1.126.Final/io.netty.channel.AbstractChannelHandlerContext.invokeChannelRead(AbstractChannelHandlerContext.java:420)
	at io.netty.transport@4.1.126.Final/io.netty.channel.AbstractChannelHandlerContext.fireChannelRead(AbstractChannelHandlerContext.java:412)
	at org.elasticsearch.transport.netty4@9.3.0/org.elasticsearch.transport.netty4.NettyByteBufSizer.channelRead(NettyByteBufSizer.java:33)
	at io.netty.transport@4.1.126.Final/io.netty.channel.AbstractChannelHandlerContext.invokeChannelRead(AbstractChannelHandlerContext.java:444)
	at io.netty.transport@4.1.126.Final/io.netty.channel.AbstractChannelHandlerContext.invokeChannelRead(AbstractChannelHandlerContext.java:420)
	at io.netty.transport@4.1.126.Final/io.netty.channel.AbstractChannelHandlerContext.fireChannelRead(AbstractChannelHandlerContext.java:412)
	at io.netty.handler@4.1.126.Final/io.netty.handler.ssl.SslHandler.unwrap(SslHandler.java:1519)
	at io.netty.handler@4.1.126.Final/io.netty.handler.ssl.SslHandler.decodeJdkCompatible(SslHandler.java:1377)
	at io.netty.handler@4.1.126.Final/io.netty.handler.ssl.SslHandler.decode(SslHandler.java:1428)
	at io.netty.codec@4.1.126.Final/io.netty.handler.codec.ByteToMessageDecoder.decodeRemovalReentryProtection(ByteToMessageDecoder.java:530)
	at io.netty.codec@4.1.126.Final/io.netty.handler.codec.ByteToMessageDecoder.callDecode(ByteToMessageDecoder.java:469)
	at io.netty.codec@4.1.126.Final/io.netty.handler.codec.ByteToMessageDecoder.channelRead(ByteToMessageDecoder.java:290)
	at io.netty.transport@4.1.126.Final/io.netty.channel.AbstractChannelHandlerContext.invokeChannelRead(AbstractChannelHandlerContext.java:444)
	at io.netty.transport@4.1.126.Final/io.netty.channel.AbstractChannelHandlerContext.invokeChannelRead(AbstractChannelHandlerContext.java:420)
	at io.netty.transport@4.1.126.Final/io.netty.channel.AbstractChannelHandlerContext.fireChannelRead(AbstractChannelHandlerContext.java:412)
	at io.netty.transport@4.1.126.Final/io.netty.channel.DefaultChannelPipeline$HeadContext.channelRead(DefaultChannelPipeline.java:1357)
	at io.netty.transport@4.1.126.Final/io.netty.channel.AbstractChannelHandlerContext.invokeChannelRead(AbstractChannelHandlerContext.java:440)
	at io.netty.transport@4.1.126.Final/io.netty.channel.AbstractChannelHandlerContext.invokeChannelRead(AbstractChannelHandlerContext.java:420)
	at io.netty.transport@4.1.126.Final/io.netty.channel.DefaultChannelPipeline.fireChannelRead(DefaultChannelPipeline.java:868)
	at io.netty.transport@4.1.126.Final/io.netty.channel.nio.AbstractNioByteChannel$NioByteUnsafe.read(AbstractNioByteChannel.java:166)
	at io.netty.transport@4.1.126.Final/io.netty.channel.nio.NioEventLoop.processSelectedKey(NioEventLoop.java:796)
	at io.netty.transport@4.1.126.Final/io.netty.channel.nio.NioEventLoop.processSelectedKeysPlain(NioEventLoop.java:697)
	at io.netty.transport@4.1.126.Final/io.netty.channel.nio.NioEventLoop.processSelectedKeys(NioEventLoop.java:660)
	at io.netty.transport@4.1.126.Final/io.netty.channel.nio.NioEventLoop.run(NioEventLoop.java:562)
	at io.netty.common@4.1.126.Final/io.netty.util.concurrent.SingleThreadEventExecutor$4.run(SingleThreadEventExecutor.java:998)
	at io.netty.common@4.1.126.Final/io.netty.util.internal.ThreadExecutorMap$2.run(ThreadExecutorMap.java:74)
	at java.base/java.lang.Thread.run(Thread.java:1474)

[elasticsearchmachine]

Pinging @elastic/search-relevance (Team:Search - Relevance)

[benwtrent]

@Mikep86 looking at the code, these couple of paths are problematic:

doRewriteGetInferenceResults

If at all possible, we should actually return a supplier that we can call get on to actually populate the results in the caller. This will return a new instance of the query object that has a "supplier" populated. This pattern is well worn, works, and is the most common way to do these async actions.

getInferenceResults

We are registering too many async actions (it should just be a single action that does multiple inference calls).

I would rewrite the inference action itself to use a GroupedActionListener

Something like, but with returning a supplier that will provide the map:

    static void registerInferenceAsyncInferenceActions(
        QueryRewriteContext queryRewriteContext,
        ConcurrentHashMap<FullyQualifiedInferenceId, InferenceResults> inferenceResultsMap,
        String query,
        Set<FullyQualifiedInferenceId> fullyQualifiedInferenceIds
    ) {
        List<InferenceAction.Request> inferenceRequests = new ArrayList<>(fullyQualifiedInferenceIds.size());
        for (FullyQualifiedInferenceId fullyQualifiedInferenceId : fullyQualifiedInferenceIds) {
            inferenceRequests.add(
                new InferenceAction.Request(
                    TaskType.ANY,
                    fullyQualifiedInferenceId.inferenceId(),
                    null,
                    null,
                    null,
                    List.of(query),
                    Map.of(),
                    InputType.INTERNAL_SEARCH,
                    null,
                    false
                )
            );
        }
        queryRewriteContext.registerAsyncAction(
            (client, listener) -> {
                GroupedActionListener<Tuple<InferenceAction.Response, String>> gal = new GroupedActionListener<>(
                    inferenceRequests.size(),
                    ActionListener.wrap(
                        responses -> {
                            for (Tuple<InferenceAction.Response, String> response : responses) {
                                InferenceAction.Response inferenceResponse = response.v1();
                                String inferenceId = response.v2();
                                inferenceResultsMap.put(
                                    new FullyQualifiedInferenceId(queryRewriteContext.getLocalClusterAlias(), inferenceId),
                                    validateAndConvertInferenceResults(inferenceResponse.getResults(), inferenceId)
                                );
                            }
                            listener.onResponse(null);
                        },
                        listener::onFailure
                    )
                );
                for (InferenceAction.Request inferenceRequest : inferenceRequests) {
                    String inferenceId = inferenceRequest.getInferenceEntityId();
                    executeAsyncWithOrigin(
                        client,
                        ML_ORIGIN,
                        InferenceAction.INSTANCE,
                        inferenceRequest,
                        gal.delegateFailureAndWrap((l, inferenceResponse) -> l.onResponse(Tuple.tuple(inferenceResponse, inferenceId)))
                    );
                }
            }
        );
    }