Support HitsLoggerPlugin (#695)

* Logging hits feature interface

* Added unit test for hits logger plugin

* Added hits logger creator test

* Use HitsLoggerProvider to initialize HitsLogger with params

* Addressed PR comments

* Added/improved some comments

* Fixed spotless violations

* Replace * imports

* Return empty map instead of null in HitsLoggerPlugin

clientlib/src/main/proto/yelp/nrtsearch/search.proto

@@ -614,6 +614,8 @@ message SearchRequest {
     map<string, InnerHit> inner_hits = 26;
     // Defines runtime fields for this query
     repeated RuntimeField runtimeFields = 27;
+    // Any custom logging that should log hits after ranking
+    LoggingHits loggingHits = 29;
 }
 
 /* Inner Hit search request */
@@ -1065,6 +1067,13 @@ message FilterResult {
     map<string, CollectorResult> nestedCollectorResults = 2;
 }
 
+message LoggingHits {
+    // name of the hits logger to be called, as registered by a HitsLoggerPlugin
+    string name = 1;
+    //Optional logging parameters
+    google.protobuf.Struct params = 2;
+}
+
 // Specify how to highlight matched text in SearchRequest
 message Highlight {
 

src/main/java/com/yelp/nrtsearch/server/grpc/LuceneServer.java

@@ -48,6 +48,7 @@
 import com.yelp.nrtsearch.server.luceneserver.index.handlers.FieldUpdateHandler;
 import com.yelp.nrtsearch.server.luceneserver.index.handlers.LiveSettingsV2Handler;
 import com.yelp.nrtsearch.server.luceneserver.index.handlers.SettingsV2Handler;
+import com.yelp.nrtsearch.server.luceneserver.logging.HitsLoggerCreator;
 import com.yelp.nrtsearch.server.luceneserver.rescore.RescorerCreator;
 import com.yelp.nrtsearch.server.luceneserver.script.ScriptService;
 import com.yelp.nrtsearch.server.luceneserver.search.FetchTaskCreator;
@@ -354,12 +355,14 @@ static void initQueryCache(LuceneServerConfiguration configuration) {
 
     private void initExtendableComponents(
         LuceneServerConfiguration configuration, List<Plugin> plugins) {
+      // this block should be in alphabetical order
       AnalyzerCreator.initialize(configuration, plugins);
       CollectorCreator.initialize(configuration, plugins);
       CustomRequestProcessor.initialize(configuration, plugins);
       FetchTaskCreator.initialize(configuration, plugins);
       FieldDefCreator.initialize(configuration, plugins);
       HighlighterService.initialize(configuration, plugins);
+      HitsLoggerCreator.initialize(configuration, plugins);
       RescorerCreator.initialize(configuration, plugins);
       ScriptService.initialize(configuration, plugins);
       SimilarityCreator.initialize(configuration, plugins);

src/main/java/com/yelp/nrtsearch/server/luceneserver/innerhit/InnerHitContext.java

@@ -85,7 +85,8 @@ private InnerHitContext(InnerHitContextBuilder builder, boolean needValidation)
     this.queryFields = builder.queryFields;
     this.retrieveFields = builder.retrieveFields;
     this.explain = builder.explain;
-    this.fetchTasks = new FetchTasks(Collections.EMPTY_LIST, builder.highlightFetchTask, null);
+    this.fetchTasks =
+        new FetchTasks(Collections.EMPTY_LIST, builder.highlightFetchTask, null, null);
 
     if (builder.querySort == null) {
       // relevance collector

src/main/java/com/yelp/nrtsearch/server/luceneserver/logging/HitsLogger.java

@@ -0,0 +1,35 @@
+/*
+ * Copyright 2024 Yelp Inc.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package com.yelp.nrtsearch.server.luceneserver.logging;
+
+import com.yelp.nrtsearch.server.grpc.SearchResponse;
+import com.yelp.nrtsearch.server.luceneserver.search.SearchContext;
+import java.util.List;
+
+/**
+ * This is the hits logger interface to provide the new logging hits feature with different
+ * implementations. HitLoggers are supposed to be initiated once at the startup time and registered
+ * in the {@link HitsLoggerCreator}.
+ */
+public interface HitsLogger {
+  /**
+   * This will be invoked once as the last fetch task of the search request.
+   *
+   * @param context the {@link SearchContext} to keep the contexts for this search request
+   * @param hits query hits
+   */
+  void log(SearchContext context, List<SearchResponse.Hit.Builder> hits);
+}

src/main/java/com/yelp/nrtsearch/server/luceneserver/logging/HitsLoggerCreator.java

@@ -0,0 +1,87 @@
+/*
+ * Copyright 2024 Yelp Inc.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package com.yelp.nrtsearch.server.luceneserver.logging;
+
+import com.yelp.nrtsearch.server.config.LuceneServerConfiguration;
+import com.yelp.nrtsearch.server.grpc.LoggingHits;
+import com.yelp.nrtsearch.server.plugins.HitsLoggerPlugin;
+import com.yelp.nrtsearch.server.plugins.Plugin;
+import com.yelp.nrtsearch.server.utils.StructValueTransformer;
+import java.util.HashMap;
+import java.util.Map;
+
+/** Factory class that handles registration and creation of {@link HitsLogger}s. */
+public class HitsLoggerCreator {
+  private static HitsLoggerCreator instance;
+  private final Map<String, HitsLoggerProvider<?>> hitsLoggerMap = new HashMap<>();
+
+  /**
+   * Constructor.
+   *
+   * @param configuration server configuration
+   */
+  public HitsLoggerCreator(LuceneServerConfiguration configuration) {}
+
+  private void register(Map<String, HitsLoggerProvider<? extends HitsLogger>> hitsLoggers) {
+    hitsLoggers.forEach(this::register);
+  }
+
+  private void register(String name, HitsLoggerProvider<?> hitsLoggerProvider) {
+    if (hitsLoggerMap.containsKey(name)) {
+      throw new IllegalArgumentException("HitsLogger " + name + " already exists");
+    }
+    hitsLoggerMap.put(name, hitsLoggerProvider);
+  }
+
+  /**
+   * Initialize singleton instance of {@link HitsLoggerCreator}. Registers the hits logger provided
+   * by {@link HitsLoggerPlugin}s.
+   *
+   * @param configuration service configuration
+   * @param plugins list of loaded plugins
+   */
+  public static void initialize(LuceneServerConfiguration configuration, Iterable<Plugin> plugins) {
+    instance = new HitsLoggerCreator(configuration);
+    for (Plugin plugin : plugins) {
+      if (plugin instanceof HitsLoggerPlugin loggerPlugin) {
+        instance.register(loggerPlugin.getHitsLoggers());
+      }
+    }
+  }
+
+  /** Get singleton instance. */
+  public static HitsLoggerCreator getInstance() {
+    return instance;
+  }
+
+  /**
+   * Create a {@link HitsLogger} instance given the {@link LoggingHits} message from the {@link
+   * com.yelp.nrtsearch.server.grpc.SearchRequest}
+   *
+   * @param grpcLoggingHits definition message
+   * @return the corresponding hits logger
+   */
+  public HitsLogger createHitsLogger(LoggingHits grpcLoggingHits) {
+    HitsLoggerProvider<?> provider = hitsLoggerMap.get(grpcLoggingHits.getName());
+    if (provider == null) {
+      throw new IllegalArgumentException(
+          String.format(
+              "Unknown hits logger name [%s] is specified; The available hits loggers are %s",
+              grpcLoggingHits.getName(), hitsLoggerMap.keySet()));
+    }
+    return provider.get(StructValueTransformer.transformStruct(grpcLoggingHits.getParams()));
+  }
+}

src/main/java/com/yelp/nrtsearch/server/luceneserver/logging/HitsLoggerFetchTask.java

@@ -0,0 +1,47 @@
+/*
+ * Copyright 2024 Yelp Inc.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package com.yelp.nrtsearch.server.luceneserver.logging;
+
+import com.yelp.nrtsearch.server.grpc.LoggingHits;
+import com.yelp.nrtsearch.server.grpc.SearchResponse;
+import com.yelp.nrtsearch.server.luceneserver.search.FetchTasks.FetchTask;
+import com.yelp.nrtsearch.server.luceneserver.search.SearchContext;
+import com.yelp.nrtsearch.server.plugins.HitsLoggerPlugin;
+import java.util.List;
+
+/**
+ * Implementation of {@link FetchTask} which holds the required context to be able to log hits for a
+ * search request.
+ */
+public class HitsLoggerFetchTask implements FetchTask {
+  private final HitsLogger hitsLogger;
+
+  public HitsLoggerFetchTask(LoggingHits loggingHits) {
+    this.hitsLogger = HitsLoggerCreator.getInstance().createHitsLogger(loggingHits);
+  }
+
+  /**
+   * Calls {@link HitsLogger} that logs hits. The logic for logging is implemented via {@link
+   * HitsLoggerPlugin}
+   *
+   * @param searchContext search context
+   * @param hits list of hits for query response
+   */
+  @Override
+  public void processAllHits(SearchContext searchContext, List<SearchResponse.Hit.Builder> hits) {
+    hitsLogger.log(searchContext, hits);
+  }
+}

src/main/java/com/yelp/nrtsearch/server/luceneserver/logging/HitsLoggerProvider.java

@@ -0,0 +1,35 @@
+/*
+ * Copyright 2024 Yelp Inc.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package com.yelp.nrtsearch.server.luceneserver.logging;
+
+import com.yelp.nrtsearch.server.grpc.LoggingHits;
+import java.util.Map;
+
+/**
+ * Interface for getting a {@link HitsLogger} implementation initialized with the given parameters
+ * map.
+ *
+ * @param <T> HitsLogger type
+ */
+public interface HitsLoggerProvider<T extends HitsLogger> {
+  /**
+   * Get task instance with the given parameters.
+   *
+   * @param params java native representation of {@link LoggingHits#getParams()}
+   * @return {@link com.yelp.nrtsearch.server.luceneserver.logging.HitsLogger} instance
+   */
+  T get(Map<String, Object> params);
+}

src/main/java/com/yelp/nrtsearch/server/luceneserver/search/FetchTasks.java

@@ -19,6 +19,7 @@
 import com.yelp.nrtsearch.server.grpc.SearchResponse.Hit.Builder;
 import com.yelp.nrtsearch.server.luceneserver.highlights.HighlightFetchTask;
 import com.yelp.nrtsearch.server.luceneserver.innerhit.InnerHitFetchTask;
+import com.yelp.nrtsearch.server.luceneserver.logging.HitsLoggerFetchTask;
 import java.io.IOException;
 import java.util.List;
 import java.util.stream.Collectors;
@@ -78,6 +79,7 @@ default void processHit(
   // highlightFetchTasks only.
   private HighlightFetchTask highlightFetchTask;
   private List<InnerHitFetchTask> innerHitFetchTaskList;
+  private HitsLoggerFetchTask hitsLoggerFetchTask;
 
   public HighlightFetchTask getHighlightFetchTask() {
     return highlightFetchTask;
@@ -87,6 +89,10 @@ public List<InnerHitFetchTask> getInnerHitFetchTaskList() {
     return innerHitFetchTaskList;
   }
 
+  public HitsLoggerFetchTask getHitsLoggerFetchTask() {
+    return hitsLoggerFetchTask;
+  }
+
   public void setHighlightFetchTask(HighlightFetchTask highlightFetchTask) {
     this.highlightFetchTask = highlightFetchTask;
   }
@@ -95,13 +101,17 @@ public void setInnerHitFetchTaskList(List<InnerHitFetchTask> innerHitFetchTaskLi
     this.innerHitFetchTaskList = innerHitFetchTaskList;
   }
 
+  public void setHitsLoggerFetchTask(HitsLoggerFetchTask hitsLoggerFetchTask) {
+    this.hitsLoggerFetchTask = hitsLoggerFetchTask;
+  }
+
   /**
    * Constructor.
    *
    * @param grpcTaskList fetch task definitions from search request
    */
   public FetchTasks(List<com.yelp.nrtsearch.server.grpc.FetchTask> grpcTaskList) {
-    this(grpcTaskList, null, null);
+    this(grpcTaskList, null, null, null);
   }
 
   /**
@@ -110,11 +120,13 @@ public FetchTasks(List<com.yelp.nrtsearch.server.grpc.FetchTask> grpcTaskList) {
    * @param grpcTaskList fetch task definitions from search request
    * @param highlightFetchTask highlight fetch task
    * @param innerHitFetchTaskList innerHit fetch tasks
+   * @param hitsLoggerFetchTask hitsLogger fetch task
    */
   public FetchTasks(
       List<com.yelp.nrtsearch.server.grpc.FetchTask> grpcTaskList,
       HighlightFetchTask highlightFetchTask,
-      List<InnerHitFetchTask> innerHitFetchTaskList) {
+      List<InnerHitFetchTask> innerHitFetchTaskList,
+      HitsLoggerFetchTask hitsLoggerFetchTask) {
     taskList =
         grpcTaskList.stream()
             .map(
@@ -124,6 +136,7 @@ public FetchTasks(
             .collect(Collectors.toList());
     this.highlightFetchTask = highlightFetchTask;
     this.innerHitFetchTaskList = innerHitFetchTaskList;
+    this.hitsLoggerFetchTask = hitsLoggerFetchTask;
   }
 
   /**
@@ -140,6 +153,12 @@ public void processAllHits(SearchContext searchContext, List<SearchResponse.Hit.
       task.processAllHits(searchContext, hits);
     }
     // highlight and innerHit doesn't support processAllHits now
+    // hitsLogger should be the last fetch task to run because it might need shared data from other
+    // plugins, including
+    // other fetch task plugins
+    if (hitsLoggerFetchTask != null) {
+      hitsLoggerFetchTask.processAllHits(searchContext, hits);
+    }
   }
 
   /**

src/main/java/com/yelp/nrtsearch/server/luceneserver/search/SearchRequestProcessor.java

@@ -18,6 +18,7 @@
 import com.yelp.nrtsearch.server.grpc.CollectorResult;
 import com.yelp.nrtsearch.server.grpc.Highlight;
 import com.yelp.nrtsearch.server.grpc.InnerHit;
+import com.yelp.nrtsearch.server.grpc.LoggingHits;
 import com.yelp.nrtsearch.server.grpc.PluginRescorer;
 import com.yelp.nrtsearch.server.grpc.ProfileResult;
 import com.yelp.nrtsearch.server.grpc.QueryRescorer;
@@ -39,6 +40,7 @@
 import com.yelp.nrtsearch.server.luceneserver.innerhit.InnerHitContext;
 import com.yelp.nrtsearch.server.luceneserver.innerhit.InnerHitContext.InnerHitContextBuilder;
 import com.yelp.nrtsearch.server.luceneserver.innerhit.InnerHitFetchTask;
+import com.yelp.nrtsearch.server.luceneserver.logging.HitsLoggerFetchTask;
 import com.yelp.nrtsearch.server.luceneserver.rescore.QueryRescore;
 import com.yelp.nrtsearch.server.luceneserver.rescore.RescoreOperation;
 import com.yelp.nrtsearch.server.luceneserver.rescore.RescoreTask;
@@ -177,6 +179,12 @@ public static SearchContext buildContextForRequest(
           new HighlightFetchTask(indexState, query, HighlighterService.getInstance(), highlight);
     }
 
+    HitsLoggerFetchTask hitsLoggerFetchTask = null;
+    if (searchRequest.hasLoggingHits()) {
+      LoggingHits loggingHits = searchRequest.getLoggingHits();
+      hitsLoggerFetchTask = new HitsLoggerFetchTask(loggingHits);
+    }
+
     List<InnerHitFetchTask> innerHitFetchTasks = null;
     if (searchRequest.getInnerHitsCount() > 0) {
       innerHitFetchTasks = new ArrayList<>(searchRequest.getInnerHitsCount());
@@ -197,7 +205,11 @@ public static SearchContext buildContextForRequest(
     }
 
     contextBuilder.setFetchTasks(
-        new FetchTasks(searchRequest.getFetchTasksList(), highlightFetchTask, innerHitFetchTasks));
+        new FetchTasks(
+            searchRequest.getFetchTasksList(),
+            highlightFetchTask,
+            innerHitFetchTasks,
+            hitsLoggerFetchTask));
 
     contextBuilder.setQuery(query);