Fix ProfileScorer.getChildren() to expose wrapped scorer in scorer tree (opensearch-project#20607)

* Fix ProfileScorer.getChildren() to expose wrapped scorer in scorer tree

Signed-off-by: Martin Gaievski <gaievski@amazon.com>

* Changed approach to use new ScoreWrapper interface

Signed-off-by: Martin Gaievski <gaievski@amazon.com>

* Renamed interface to WrappedScorerAccessor, added unit test

Signed-off-by: Martin Gaievski <gaievski@amazon.com>

* Improving tests, corrected changelog

Signed-off-by: Martin Gaievski <gaievski@amazon.com>

* Changed to generic interface ProfilingWrapper

Signed-off-by: Martin Gaievski <gaievski@amazon.com>

* Adding ProfilingWrapper to ProfilingAggregator and ProfilingLeafBucketCollector

Signed-off-by: Martin Gaievski <gaievski@amazon.com>

---------

Signed-off-by: Martin Gaievski <gaievski@amazon.com>

Author: martin-gaievski

CHANGELOG.md

@@ -5,7 +5,7 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/),
 
 ## [Unreleased 3.x]
 ### Added
-- Add getWrappedScorer method to ProfileScorer for plugin access to wrapped scorers ([#20548](https://github.com/opensearch-project/OpenSearch/issues/20548))
+- Add ProfilingWrapper interface for plugin access to delegates in profiling decorators ([#20607](https://github.com/opensearch-project/OpenSearch/pull/20607))
 - Support expected cluster name with validation in CCS Sniff mode ([#20532](https://github.com/opensearch-project/OpenSearch/pull/20532))
 - Choose the best performing node when writing with append-only index ([#20065](https://github.com/opensearch-project/OpenSearch/pull/20065))
 - Add security policy to allow `accessUnixDomainSocket` in `transport-grpc` module ([#20463](https://github.com/opensearch-project/OpenSearch/pull/20463), [#20649](https://github.com/opensearch-project/OpenSearch/pull/20649))

server/src/main/java/org/opensearch/search/profile/ProfilingWrapper.java

@@ -0,0 +1,44 @@
+/*
+ * SPDX-License-Identifier: Apache-2.0
+ *
+ * The OpenSearch Contributors require contributions made to
+ * this file be licensed under the Apache-2.0 license or a
+ * compatible open source license.
+ */
+
+package org.opensearch.search.profile;
+
+import org.opensearch.common.annotation.PublicApi;
+
+/**
+ * Generic interface for profiling wrappers that decorate search components with timing instrumentation
+ * while preserving access to the original (delegate) component.
+ * <p>
+ * When search profiling is enabled, OpenSearch wraps scorers, collectors, aggregators, and other components
+ * in profiling decorators that are package-private. This interface provides a public contract for plugins
+ * to detect and unwrap such wrappers, accessing the underlying component without requiring reflection
+ * or knowledge of package-private implementation classes.
+ * <p>
+ * Example usage in a plugin to unwrap a profiled scorer:
+ * <pre>{@code
+ * if (scorer instanceof ProfilingWrapper) {
+ *     Scorer delegate = ((ProfilingWrapper<Scorer>) scorer).getDelegate();
+ *     // access custom scorer methods on the delegate
+ * }
+ * }</pre>
+ *
+ * @param <T> the type of the wrapped component (e.g., {@link org.apache.lucene.search.Scorer},
+ *            {@link org.apache.lucene.search.Collector},
+ *            {@link org.opensearch.search.aggregations.Aggregator})
+ * @opensearch.api
+ */
+@PublicApi(since = "3.6.0")
+public interface ProfilingWrapper<T> {
+
+    /**
+     * Returns the underlying delegate component that this profiling wrapper decorates.
+     *
+     * @return the original component that was wrapped for profiling
+     */
+    T getDelegate();
+}

server/src/main/java/org/opensearch/search/profile/aggregation/ProfilingAggregator.java

@@ -39,6 +39,7 @@
 import org.opensearch.search.aggregations.LeafBucketCollector;
 import org.opensearch.search.aggregations.support.AggregationPath.PathElement;
 import org.opensearch.search.internal.SearchContext;
+import org.opensearch.search.profile.ProfilingWrapper;
 import org.opensearch.search.profile.Timer;
 import org.opensearch.search.sort.SortOrder;
 
@@ -48,7 +49,7 @@
 /**
  * An aggregator that aggregates the performance profiling of other aggregations
  */
-public class ProfilingAggregator extends Aggregator {
+public class ProfilingAggregator extends Aggregator implements ProfilingWrapper<Aggregator> {
 
     private final Aggregator delegate;
     private final AggregationProfiler profiler;
@@ -162,6 +163,11 @@ public String toString() {
         return delegate.toString();
     }
 
+    @Override
+    public Aggregator getDelegate() {
+        return delegate;
+    }
+
     @Override
     public Aggregator unwrapAggregator() {
         return delegate;

server/src/main/java/org/opensearch/search/profile/aggregation/ProfilingLeafBucketCollector.java

@@ -34,6 +34,7 @@
 
 import org.apache.lucene.search.Scorable;
 import org.opensearch.search.aggregations.LeafBucketCollector;
+import org.opensearch.search.profile.ProfilingWrapper;
 import org.opensearch.search.profile.Timer;
 
 import java.io.IOException;
@@ -43,9 +44,9 @@
  *
  * @opensearch.internal
  */
-public class ProfilingLeafBucketCollector extends LeafBucketCollector {
+public class ProfilingLeafBucketCollector extends LeafBucketCollector implements ProfilingWrapper<LeafBucketCollector> {
 
-    private LeafBucketCollector delegate;
+    private final LeafBucketCollector delegate;
     private Timer collectTimer;
 
     public ProfilingLeafBucketCollector(LeafBucketCollector delegate, AggregationProfileBreakdown profileBreakdown) {
@@ -63,6 +64,11 @@ public void collect(int doc, long bucket) throws IOException {
         }
     }
 
+    @Override
+    public LeafBucketCollector getDelegate() {
+        return delegate;
+    }
+
     @Override
     public void setScorer(Scorable scorer) throws IOException {
         delegate.setScorer(scorer);

server/src/main/java/org/opensearch/search/profile/query/ProfileCollector.java

@@ -39,6 +39,7 @@
 import org.apache.lucene.search.LeafCollector;
 import org.apache.lucene.search.Scorable;
 import org.apache.lucene.search.ScoreMode;
+import org.opensearch.search.profile.ProfilingWrapper;
 
 import java.io.IOException;
 
@@ -47,7 +48,7 @@
  *
  * @opensearch.internal
  */
-final class ProfileCollector extends FilterCollector {
+final class ProfileCollector extends FilterCollector implements ProfilingWrapper<Collector> {
 
     private long time;
     private long sliceStartTime;

server/src/main/java/org/opensearch/search/profile/query/ProfileScorer.java

@@ -36,6 +36,7 @@
 import org.apache.lucene.search.Scorer;
 import org.apache.lucene.search.TwoPhaseIterator;
 import org.opensearch.search.profile.AbstractProfileBreakdown;
+import org.opensearch.search.profile.ProfilingWrapper;
 import org.opensearch.search.profile.Timer;
 
 import java.io.IOException;
@@ -47,7 +48,7 @@
  *
  * @opensearch.internal
  */
-final class ProfileScorer extends Scorer {
+final class ProfileScorer extends Scorer implements ProfilingWrapper<Scorer> {
 
     private final Scorer scorer;
 
@@ -81,7 +82,8 @@ final class ProfileScorer extends Scorer {
      * @return the underlying wrapped scorer
      * @see ProfileCollector#getDelegate()
      */
-    public Scorer getWrappedScorer() {
+    @Override
+    public Scorer getDelegate() {
         return scorer;
     }
 

server/src/test/java/org/opensearch/search/profile/aggregation/ProfilingWrapperAggregationTests.java

@@ -0,0 +1,51 @@
+/*
+ * SPDX-License-Identifier: Apache-2.0
+ *
+ * The OpenSearch Contributors require contributions made to
+ * this file be licensed under the Apache-2.0 license or a
+ * compatible open source license.
+ */
+
+package org.opensearch.search.profile.aggregation;
+
+import org.opensearch.search.aggregations.Aggregator;
+import org.opensearch.search.aggregations.LeafBucketCollector;
+import org.opensearch.search.profile.ProfilingWrapper;
+import org.opensearch.test.OpenSearchTestCase;
+
+import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.when;
+
+/**
+ * Tests that ProfilingAggregator and ProfilingLeafBucketCollector
+ * correctly implement the ProfilingWrapper interface.
+ */
+public class ProfilingWrapperAggregationTests extends OpenSearchTestCase {
+
+    public void testProfilingAggregatorImplementsProfilingWrapper() {
+        Aggregator mockAggregator = mock(Aggregator.class);
+        when(mockAggregator.name()).thenReturn("test_agg");
+        AggregationProfiler profiler = new AggregationProfiler();
+
+        ProfilingAggregator profilingAggregator = new ProfilingAggregator(mockAggregator, profiler);
+
+        // Verify implements ProfilingWrapper
+        assertTrue("ProfilingAggregator should implement ProfilingWrapper", profilingAggregator instanceof ProfilingWrapper);
+
+        // Verify getDelegate returns the original aggregator
+        assertSame("getDelegate() should return the original aggregator", mockAggregator, profilingAggregator.getDelegate());
+    }
+
+    public void testProfilingLeafBucketCollectorImplementsProfilingWrapper() {
+        LeafBucketCollector mockCollector = mock(LeafBucketCollector.class);
+        AggregationProfileBreakdown profileBreakdown = new AggregationProfileBreakdown();
+
+        ProfilingLeafBucketCollector profilingCollector = new ProfilingLeafBucketCollector(mockCollector, profileBreakdown);
+
+        // Verify implements ProfilingWrapper
+        assertTrue("ProfilingLeafBucketCollector should implement ProfilingWrapper", profilingCollector instanceof ProfilingWrapper);
+
+        // Verify getDelegate returns the original collector
+        assertSame("getDelegate() should return the original collector", mockCollector, profilingCollector.getDelegate());
+    }
+}

server/src/test/java/org/opensearch/search/profile/query/ProfileScorerTests.java

@@ -37,13 +37,16 @@
 import org.apache.lucene.search.IndexSearcher;
 import org.apache.lucene.search.MatchAllDocsQuery;
 import org.apache.lucene.search.Query;
+import org.apache.lucene.search.Scorable;
 import org.apache.lucene.search.ScoreMode;
 import org.apache.lucene.search.Scorer;
 import org.apache.lucene.search.Weight;
 import org.opensearch.search.profile.ProfileMetricUtil;
+import org.opensearch.search.profile.ProfilingWrapper;
 import org.opensearch.test.OpenSearchTestCase;
 
 import java.io.IOException;
+import java.util.Collection;
 
 public class ProfileScorerTests extends OpenSearchTestCase {
 
@@ -102,14 +105,71 @@ public void testPropagateMaxScore() throws IOException {
         assertEquals(42f, profileScorer.getMaxScore(DocIdSetIterator.NO_MORE_DOCS), 0f);
     }
 
-    public void testGetWrappedScorer() throws IOException {
+    public void testGetDelegate() throws IOException {
         Query query = new MatchAllDocsQuery();
         Weight weight = query.createWeight(new IndexSearcher(new MultiReader()), ScoreMode.TOP_SCORES, 1f);
         FakeScorer fakeScorer = new FakeScorer(weight);
         QueryProfileBreakdown profile = new QueryProfileBreakdown(ProfileMetricUtil.getDefaultQueryProfileMetrics());
         ProfileScorer profileScorer = new ProfileScorer(fakeScorer, profile);
 
-        // Verify getWrappedScorer returns the original scorer
-        assertSame(fakeScorer, profileScorer.getWrappedScorer());
+        // Verify getDelegate returns the original scorer
+        assertSame(fakeScorer, profileScorer.getDelegate());
+    }
+
+    @SuppressWarnings("unchecked")
+    public void testImplementsProfilingWrapper() throws IOException {
+        Query query = new MatchAllDocsQuery();
+        Weight weight = query.createWeight(new IndexSearcher(new MultiReader()), ScoreMode.TOP_SCORES, 1f);
+        FakeScorer fakeScorer = new FakeScorer(weight);
+        QueryProfileBreakdown profile = new QueryProfileBreakdown(ProfileMetricUtil.getDefaultQueryProfileMetrics());
+        ProfileScorer profileScorer = new ProfileScorer(fakeScorer, profile);
+
+        // ProfileScorer implements ProfilingWrapper<Scorer>, allowing plugins to detect and unwrap
+        // profiling wrappers using instanceof without reflection
+        assertTrue("ProfileScorer should implement ProfilingWrapper", profileScorer instanceof ProfilingWrapper);
+        ProfilingWrapper<Scorer> wrapper = (ProfilingWrapper<Scorer>) profileScorer;
+        assertSame("ProfilingWrapper.getDelegate() should return the original scorer", fakeScorer, wrapper.getDelegate());
+    }
+
+    @SuppressWarnings("unchecked")
+    public void testUnwrapNestedProfilingScorerFromScorableReference() throws IOException {
+        // Simulate nested profiling: a scorer wrapped in two ProfileScorer layers.
+        // This tests that a plugin can iteratively unwrap through multiple profiling
+        // layers via a generic Scorable reference (as received in LeafCollector.setScorer())
+        Query query = new MatchAllDocsQuery();
+        Weight weight = query.createWeight(new IndexSearcher(new MultiReader()), ScoreMode.TOP_SCORES, 1f);
+        FakeScorer fakeScorer = new FakeScorer(weight);
+        QueryProfileBreakdown innerProfile = new QueryProfileBreakdown(ProfileMetricUtil.getDefaultQueryProfileMetrics());
+        QueryProfileBreakdown outerProfile = new QueryProfileBreakdown(ProfileMetricUtil.getDefaultQueryProfileMetrics());
+
+        // Create nested profiling: FakeScorer -> innerProfileScorer -> outerProfileScorer
+        ProfileScorer innerProfileScorer = new ProfileScorer(fakeScorer, innerProfile);
+        ProfileScorer outerProfileScorer = new ProfileScorer(innerProfileScorer, outerProfile);
+
+        // Plugin only sees a Scorable reference
+        Scorable scorable = outerProfileScorer;
+
+        // Unwrap first layer
+        assertTrue("Outer Scorable should be detected as ProfilingWrapper", scorable instanceof ProfilingWrapper);
+        Scorer firstUnwrap = ((ProfilingWrapper<Scorer>) scorable).getDelegate();
+        assertNotSame("First unwrap should not be the original scorer", fakeScorer, firstUnwrap);
+        assertTrue("First unwrap should still be a ProfilingWrapper (inner ProfileScorer)", firstUnwrap instanceof ProfilingWrapper);
+
+        // Unwrap second layer to reach the original scorer
+        Scorer secondUnwrap = ((ProfilingWrapper<Scorer>) firstUnwrap).getDelegate();
+        assertSame("Second unwrap should be the original FakeScorer", fakeScorer, secondUnwrap);
+        assertFalse("Original scorer should not implement ProfilingWrapper", secondUnwrap instanceof ProfilingWrapper);
+    }
+
+    public void testGetChildren_delegatesToWrappedScorer() throws IOException {
+        Query query = new MatchAllDocsQuery();
+        Weight weight = query.createWeight(new IndexSearcher(new MultiReader()), ScoreMode.TOP_SCORES, 1f);
+        FakeScorer fakeScorer = new FakeScorer(weight);
+        QueryProfileBreakdown profile = new QueryProfileBreakdown(ProfileMetricUtil.getDefaultQueryProfileMetrics());
+        ProfileScorer profileScorer = new ProfileScorer(fakeScorer, profile);
+
+        // getChildren() should delegate to the wrapped scorer's getChildren()
+        Collection<Scorer.ChildScorable> children = profileScorer.getChildren();
+        assertEquals("FakeScorer has no children, so ProfileScorer should return empty collection", 0, children.size());
     }
 }