rebase

common/src/main/java/org/apache/uniffle/common/audit/RpcAuditContext.java

@@ -29,6 +29,7 @@ public abstract class RpcAuditContext implements Closeable {
   private String command;
   private String statusCode;
   private String args;
+  private String returnValue;
   private String from;
   private long creationTimeNs;
   protected long executionTimeNs;
@@ -103,12 +104,17 @@ public RpcAuditContext withStatusCode(String statusCode) {
     return this;
   }
 
-  public RpcAuditContext setArgs(String args) {
+  public RpcAuditContext withArgs(String args) {
     this.args = args;
     return this;
   }
 
-  public RpcAuditContext setFrom(String from) {
+  public RpcAuditContext withReturnValue(String returnValue) {
+    this.returnValue = returnValue;
+    return this;
+  }
+
+  public RpcAuditContext withFrom(String from) {
     this.from = from;
     return this;
   }
@@ -131,6 +137,9 @@ public String toString() {
     if (args != null) {
       line += String.format("\targs{%s}", args);
     }
+    if (returnValue != null) {
+      line += String.format("\treturn{%s}", returnValue);
+    }
     return line;
   }
 }

coordinator/src/main/java/org/apache/uniffle/coordinator/CoordinatorGrpcService.java

@@ -136,7 +136,7 @@ public void getShuffleAssignments(
       final int estimateTaskConcurrency = request.getEstimateTaskConcurrency();
       final Set<String> faultyServerIds = new HashSet<>(request.getFaultyServerIdsList());
 
-      auditContext.withAppId(appId).withShuffleId(shuffleId);
+      auditContext.withAppId(appId);
       auditContext.withArgs(
           String.format(
               "shuffleId=%d, partitionNum=%d, partitionNumPerRange=%d, replica=%d, requiredTags=%s, "
@@ -540,7 +540,7 @@ private CoordinatorRpcAuditContext createAuditContext(String command) {
     if (auditLogger != null) {
       auditContext
           .withCommand(command)
-          .setFrom(ClientContextServerInterceptor.getIpAddress())
+          .withFrom(ClientContextServerInterceptor.getIpAddress())
           .withCreationTimeNs(System.nanoTime());
     }
     return auditContext;

coordinator/src/main/java/org/apache/uniffle/coordinator/audit/CoordinatorRpcAuditContext.java

@@ -39,7 +39,7 @@ protected String content() {
     return String.format("appId=%s", appId);
   }
 
-  public CoordinatorRpcAuditContext setAppId(String appId) {
+  public CoordinatorRpcAuditContext withAppId(String appId) {
     this.appId = appId;
     return this;
   }

server/src/main/java/org/apache/uniffle/server/ShuffleServerGrpcService.java

@@ -1466,7 +1466,7 @@ private ServerRpcAuditContext createAuditContext(String command) {
     if (auditLogger != null) {
       auditContext
           .withCommand(command)
-          .setFrom(ClientContextServerInterceptor.getIpAddress())
+          .withFrom(ClientContextServerInterceptor.getIpAddress())
           .withCreationTimeNs(System.nanoTime());
     }
     return auditContext;

server/src/main/java/org/apache/uniffle/server/audit/ServerRpcAuditContext.java

@@ -0,0 +1,52 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You 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 org.apache.uniffle.server.audit;
+
+import org.slf4j.Logger;
+
+import org.apache.uniffle.common.audit.RpcAuditContext;
+
+/** An audit context for shuffle server rpc. */
+public class ServerRpcAuditContext extends RpcAuditContext {
+  private String appId = "N/A";
+  private int shuffleId = -1;
+
+  /**
+   * Constructor of {@link ServerRpcAuditContext}.
+   *
+   * @param log the logger to log the audit information
+   */
+  public ServerRpcAuditContext(Logger log) {
+    super(log);
+  }
+
+  @Override
+  protected String content() {
+    return String.format("appId=%s\tshuffleId=%s", appId, shuffleId);
+  }
+
+  public ServerRpcAuditContext withAppId(String appId) {
+    this.appId = appId;
+    return this;
+  }
+
+  public ServerRpcAuditContext withShuffleId(int shuffleId) {
+    this.shuffleId = shuffleId;
+    return this;
+  }
+}

server/src/main/java/org/apache/uniffle/server/netty/ShuffleServerNettyHandler.java

@@ -895,7 +895,7 @@ private ServerRpcAuditContext createAuditContext(
     if (auditLogger != null) {
       auditContext
           .withCommand(command)
-          .setFrom(transportClient.getSocketAddress().toString())
+          .withFrom(transportClient.getSocketAddress().toString())
           .withCreationTimeNs(System.nanoTime());
     }
     return auditContext;