From 9e24ed2196e062289e1ad81a90b0cc1a4ad45c32 Mon Sep 17 00:00:00 2001
From: Christos Bisias <christosbis@gmail.com>
Date: Thu, 20 Apr 2023 20:11:25 +0300
Subject: [PATCH] HADOOP-18691. Add a CallerContext getter on the Schedulable
 interface (#5540)

---
 .../org/apache/hadoop/ipc/Schedulable.java    | 14 ++++++++++++++
 .../java/org/apache/hadoop/ipc/Server.java    |  5 +++++
 .../hadoop/ipc/TestIdentityProviders.java     | 19 ++++++++++++++-----
 3 files changed, 33 insertions(+), 5 deletions(-)

diff --git a/hadoop-common-project/hadoop-common/src/main/java/org/apache/hadoop/ipc/Schedulable.java b/hadoop-common-project/hadoop-common/src/main/java/org/apache/hadoop/ipc/Schedulable.java
index 3b28d85428..00c9994e2a 100644
--- a/hadoop-common-project/hadoop-common/src/main/java/org/apache/hadoop/ipc/Schedulable.java
+++ b/hadoop-common-project/hadoop-common/src/main/java/org/apache/hadoop/ipc/Schedulable.java
@@ -29,5 +29,19 @@
 public interface Schedulable {
   public UserGroupInformation getUserGroupInformation();
 
+  /**
+   * This is overridden only in {@link Server.Call}.
+   * The CallerContext field will be used to carry information
+   * about the user in cases where UGI proves insufficient.
+   * Any other classes that might try to use this method,
+   * will get an UnsupportedOperationException.
+   *
+   * @return an instance of CallerContext if method
+   * is overridden else get an UnsupportedOperationException
+   */
+  default CallerContext getCallerContext() {
+    throw new UnsupportedOperationException("Invalid operation.");
+  }
+
   int getPriorityLevel();
 }
diff --git a/hadoop-common-project/hadoop-common/src/main/java/org/apache/hadoop/ipc/Server.java b/hadoop-common-project/hadoop-common/src/main/java/org/apache/hadoop/ipc/Server.java
index 26f519716f..9851f9b09f 100644
--- a/hadoop-common-project/hadoop-common/src/main/java/org/apache/hadoop/ipc/Server.java
+++ b/hadoop-common-project/hadoop-common/src/main/java/org/apache/hadoop/ipc/Server.java
@@ -1086,6 +1086,11 @@ public UserGroupInformation getUserGroupInformation() {
       return getRemoteUser();
     }
 
+    @Override
+    public CallerContext getCallerContext() {
+      return this.callerContext;
+    }
+
     @Override
     public int getPriorityLevel() {
       return this.priorityLevel;
diff --git a/hadoop-common-project/hadoop-common/src/test/java/org/apache/hadoop/ipc/TestIdentityProviders.java b/hadoop-common-project/hadoop-common/src/test/java/org/apache/hadoop/ipc/TestIdentityProviders.java
index 263841246b..b528186ad2 100644
--- a/hadoop-common-project/hadoop-common/src/test/java/org/apache/hadoop/ipc/TestIdentityProviders.java
+++ b/hadoop-common-project/hadoop-common/src/test/java/org/apache/hadoop/ipc/TestIdentityProviders.java
@@ -20,8 +20,9 @@
 
 import static org.junit.Assert.assertEquals;
 import static org.junit.Assert.assertNotNull;
-import static org.junit.Assert.assertTrue;
+import static org.assertj.core.api.Assertions.assertThat;
 
+import org.apache.hadoop.test.LambdaTestUtils;
 import org.junit.Test;
 
 import java.util.List;
@@ -33,7 +34,7 @@
 import org.apache.hadoop.conf.Configuration;
 
 public class TestIdentityProviders {
-  public class FakeSchedulable implements Schedulable {
+  public static class FakeSchedulable implements Schedulable {
     public FakeSchedulable() {
     }
 
@@ -61,7 +62,9 @@ public void testPluggableIdentityProvider() {
       CommonConfigurationKeys.IPC_IDENTITY_PROVIDER_KEY,
       IdentityProvider.class);
 
-    assertTrue(providers.size() == 1);
+    assertThat(providers)
+        .describedAs("provider list")
+        .hasSize(1);
 
     IdentityProvider ip = providers.get(0);
     assertNotNull(ip);
@@ -69,14 +72,20 @@ public void testPluggableIdentityProvider() {
   }
 
   @Test
-  public void testUserIdentityProvider() throws IOException {
+  public void testUserIdentityProvider() throws Exception {
     UserIdentityProvider uip = new UserIdentityProvider();
-    String identity = uip.makeIdentity(new FakeSchedulable());
+    FakeSchedulable fakeSchedulable = new FakeSchedulable();
+    String identity = uip.makeIdentity(fakeSchedulable);
 
     // Get our username
     UserGroupInformation ugi = UserGroupInformation.getCurrentUser();
     String username = ugi.getUserName();
 
     assertEquals(username, identity);
+
+    // FakeSchedulable doesn't override getCallerContext()
+    // accessing it should throw an UnsupportedOperationException
+    LambdaTestUtils.intercept(UnsupportedOperationException.class,
+        "Invalid operation.", fakeSchedulable::getCallerContext);
   }
 }