Update the WorkRequestHandler to use callbacks of type: BiFunction<WorkRequest, PrintWriter, Integer>:

 - Mark constructors that use BiFunction<List<String>, PrintWriter, Integer> callback as deprecated.
 - Use a wrapper class for the  BiFunction<WorkRequest, PrintWriter, Integer>. Suggesting this to avoid having two constructors that takes a BiFunction, as it creates a confusion between the deprecated and new constructor when given a lambda expressions.

RELNOTES: None.
PiperOrigin-RevId: 384643302

Author: Googler

src/main/java/com/google/devtools/build/lib/worker/WorkRequestHandler.java

@@ -106,7 +106,7 @@ synchronized void addOutput(String s) {
   final ConcurrentMap<Integer, RequestInfo> activeRequests = new ConcurrentHashMap<>();
 
   /** The function to be called after each {@link WorkRequest} is read. */
-  private final BiFunction<List<String>, PrintWriter, Integer> callback;
+  private final WorkRequestCallback callback;
 
   /** This worker's stderr. */
   private final PrintStream stderr;
@@ -129,6 +129,7 @@ synchronized void addOutput(String s) {
    * @param messageProcessor Object responsible for parsing {@code WorkRequest}s from the server and
    *     writing {@code WorkResponses} to the server.
    */
+  @Deprecated
   public WorkRequestHandler(
       BiFunction<List<String>, PrintWriter, Integer> callback,
       PrintStream stderr,
@@ -163,23 +164,69 @@ public WorkRequestHandler(
   /**
    * Creates a {@code WorkRequestHandler} that will call {@code callback} for each WorkRequest
    * received. Only used for the Builder.
+   *
+   * @deprecated Use WorkRequestHandlerBuilder instead.
    */
+  @Deprecated
   private WorkRequestHandler(
       BiFunction<List<String>, PrintWriter, Integer> callback,
       PrintStream stderr,
       WorkerMessageProcessor messageProcessor,
       Duration cpuUsageBeforeGc,
       BiConsumer<Integer, Thread> cancelCallback) {
+    this(
+        new WorkRequestCallback((request, pw) -> callback.apply(request.getArgumentsList(), pw)),
+        stderr,
+        messageProcessor,
+        cpuUsageBeforeGc,
+        cancelCallback);
+  }
+
+  /**
+   * Creates a {@code WorkRequestHandler} that will call {@code callback} for each WorkRequest
+   * received. Only used for the Builder.
+   *
+   * @param callback WorkRequestCallback object with Callback method for executing a single
+   *     WorkRequest in a thread. The first argument to {@code callback} is the WorkRequest, the
+   *     second is where all error messages and other user-oriented messages should be written to.
+   *     The callback must return an exit code indicating success (zero) or failure (nonzero).
+   */
+  private WorkRequestHandler(
+      WorkRequestCallback callback,
+      PrintStream stderr,
+      WorkerMessageProcessor messageProcessor,
+      Duration cpuUsageBeforeGc,
+      BiConsumer<Integer, Thread> cancelCallback) {
     this.callback = callback;
     this.stderr = stderr;
     this.messageProcessor = messageProcessor;
     this.gcScheduler = new CpuTimeBasedGcScheduler(cpuUsageBeforeGc);
     this.cancelCallback = cancelCallback;
   }
 
+  /** A wrapper class for the callback BiFunction */
+  public static class WorkRequestCallback {
+
+    /**
+     * Callback method for executing a single WorkRequest in a thread. The first argument to {@code
+     * callback} is the WorkRequest, the second is where all error messages and other user-oriented
+     * messages should be written to. The callback must return an exit code indicating success
+     * (zero) or failure (nonzero).
+     */
+    private final BiFunction<WorkRequest, PrintWriter, Integer> callback;
+
+    public WorkRequestCallback(BiFunction<WorkRequest, PrintWriter, Integer> callback) {
+      this.callback = callback;
+    }
+
+    public Integer apply(WorkRequest workRequest, PrintWriter printWriter) {
+      return callback.apply(workRequest, printWriter);
+    }
+  }
+
   /** Builder class for WorkRequestHandler. Required parameters are passed to the constructor. */
   public static class WorkRequestHandlerBuilder {
-    private final BiFunction<List<String>, PrintWriter, Integer> callback;
+    private final WorkRequestCallback callback;
     private final PrintStream stderr;
     private final WorkerMessageProcessor messageProcessor;
     private Duration cpuUsageBeforeGc = Duration.ZERO;
@@ -195,11 +242,32 @@ public static class WorkRequestHandlerBuilder {
      * @param stderr Stream that log messages should be written to, typically the process' stderr.
      * @param messageProcessor Object responsible for parsing {@code WorkRequest}s from the server
      *     and writing {@code WorkResponses} to the server.
+     * @deprecated use WorkRequestHandlerBuilder with WorkRequestCallback instead
      */
+    @Deprecated
     public WorkRequestHandlerBuilder(
         BiFunction<List<String>, PrintWriter, Integer> callback,
         PrintStream stderr,
         WorkerMessageProcessor messageProcessor) {
+      this(
+          new WorkRequestCallback((request, pw) -> callback.apply(request.getArgumentsList(), pw)),
+          stderr,
+          messageProcessor);
+    }
+
+    /**
+     * Creates a {@code WorkRequestHandlerBuilder}.
+     *
+     * @param callback WorkRequestCallback object with Callback method for executing a single
+     *     WorkRequest in a thread. The first argument to {@code callback} is the WorkRequest, the
+     *     second is where all error messages and other user-oriented messages should be written to.
+     *     The callback must return an exit code indicating success (zero) or failure (nonzero).
+     * @param stderr Stream that log messages should be written to, typically the process' stderr.
+     * @param messageProcessor Object responsible for parsing {@code WorkRequest}s from the server
+     *     and writing {@code WorkResponses} to the server.
+     */
+    public WorkRequestHandlerBuilder(
+        WorkRequestCallback callback, PrintStream stderr, WorkerMessageProcessor messageProcessor) {
       this.callback = callback;
       this.stderr = stderr;
       this.messageProcessor = messageProcessor;
@@ -287,7 +355,7 @@ void respondToRequest(WorkRequest request, RequestInfo requestInfo) throws IOExc
         PrintWriter pw = new PrintWriter(sw)) {
       int exitCode;
       try {
-        exitCode = callback.apply(request.getArgumentsList(), pw);
+        exitCode = callback.apply(request, pw);
       } catch (RuntimeException e) {
         e.printStackTrace(pw);
         exitCode = 1;

src/test/java/com/google/devtools/build/lib/worker/WorkRequestHandlerTest.java

@@ -17,6 +17,7 @@
 import static com.google.common.truth.Truth.assertThat;
 
 import com.google.devtools.build.lib.worker.WorkRequestHandler.RequestInfo;
+import com.google.devtools.build.lib.worker.WorkRequestHandler.WorkRequestCallback;
 import com.google.devtools.build.lib.worker.WorkRequestHandler.WorkRequestHandlerBuilder;
 import com.google.devtools.build.lib.worker.WorkRequestHandler.WorkerMessageProcessor;
 import com.google.devtools.build.lib.worker.WorkerProtocol.WorkRequest;
@@ -407,4 +408,27 @@ public void close() throws IOException {
       delegate.close();
     }
   }
+
+  @Test
+  public void testWorkRequestHandler_withWorkRequestCallback() throws IOException {
+    ByteArrayOutputStream out = new ByteArrayOutputStream();
+    WorkRequestCallback callback =
+        new WorkRequestCallback((request, err) -> request.getArgumentsCount());
+    WorkRequestHandler handler =
+        new WorkRequestHandlerBuilder(
+                callback,
+                new PrintStream(new ByteArrayOutputStream()),
+                new ProtoWorkerMessageProcessor(new ByteArrayInputStream(new byte[0]), out))
+            .build();
+
+    List<String> args = Arrays.asList("--sources", "B.java");
+    WorkRequest request = WorkRequest.newBuilder().addAllArguments(args).build();
+    handler.respondToRequest(request, new RequestInfo(null));
+
+    WorkResponse response =
+        WorkResponse.parseDelimitedFrom(new ByteArrayInputStream(out.toByteArray()));
+    assertThat(response.getRequestId()).isEqualTo(0);
+    assertThat(response.getExitCode()).isEqualTo(2);
+    assertThat(response.getOutput()).isEmpty();
+  }
 }