Merge pull request #15 from hynix/master

yarn: use delegation token
exacaster · Jan 28, 2022 · 5afc5a7 · 5afc5a7
2 parents 82ccbd5 + e884017
commit 5afc5a7
Show file tree

Hide file tree

Showing 3 changed files with 33 additions and 3 deletions.
diff --git a/server/src/main/java/com/exacaster/lighter/backend/yarn/YarnBackend.java b/server/src/main/java/com/exacaster/lighter/backend/yarn/YarnBackend.java
@@ -7,6 +7,7 @@
 import com.exacaster.lighter.application.ApplicationState;
 import com.exacaster.lighter.backend.Backend;
 import com.exacaster.lighter.backend.yarn.resources.State;
+import com.exacaster.lighter.backend.yarn.resources.Token;
 import com.exacaster.lighter.backend.yarn.resources.YarnApplication;
 import com.exacaster.lighter.backend.yarn.resources.YarnApplicationListResponse;
 import com.exacaster.lighter.backend.yarn.resources.YarnApplicationResponse;
@@ -21,11 +22,12 @@
 import java.util.Map;
 import java.util.Objects;
 import java.util.Optional;
+import org.springframework.http.HttpEntity;
 import org.springframework.security.kerberos.client.KerberosRestTemplate;
 
 public class YarnBackend implements Backend {
 
-    private static final String STATE_ENDPOINT = "/ws/v1/cluster/apps/{appId}/state";
+    private static final String TOKEN_ENDPOINT = "/ws/v1/cluster/delegation-token";
 
     private final YarnProperties yarnProperties;
     private final YarnClient client;
@@ -80,8 +82,17 @@ public String getSessionJobResources() {
     @Override
     public void kill(Application application) {
         var state = new State("KILLED");
-        getYarnApplicationId(application).ifPresent(id -> kerberosRestTemplate.ifPresentOrElse(
-                it -> it.put(yarnProperties.getUrl() + STATE_ENDPOINT, state, id), () -> client.setState(id, state)));
+        getYarnApplicationId(application).ifPresent(
+                id -> getToken().ifPresentOrElse(t -> client.setState(id, state, t), () -> client.setState(id, state)));
+    }
+
+    private Optional<String> getToken() {
+        var url = yarnProperties.getUrl() + TOKEN_ENDPOINT;
+        var body = Map.of("renewer", "lighter");
+        return kerberosRestTemplate
+                .map(it -> it.postForEntity(url, body, Token.class))
+                .map(HttpEntity::getBody)
+                .map(Token::getToken);
     }
 
     @Override

diff --git a/server/src/main/java/com/exacaster/lighter/backend/yarn/YarnClient.java b/server/src/main/java/com/exacaster/lighter/backend/yarn/YarnClient.java
@@ -21,6 +21,9 @@ public interface YarnClient {
     @Put("/apps/{appId}/state")
     State setState(String appId, @Body State state);
 
+    @Put("/apps/{appId}/state")
+    State setState(String appId, @Body State state, @Header("X-Hadoop-Delegation-Token") String token);
+
     @Get("/apps/{appId}")
     YarnApplicationResponse getApplication(String appId);
 

diff --git a/server/src/main/java/com/exacaster/lighter/backend/yarn/resources/Token.java b/server/src/main/java/com/exacaster/lighter/backend/yarn/resources/Token.java
@@ -0,0 +1,16 @@
+package com.exacaster.lighter.backend.yarn.resources;
+
+import io.micronaut.core.annotation.Introspected;
+
+@Introspected
+public class Token {
+    private final String token;
+
+    public Token(String token) {
+        this.token = token;
+    }
+
+    public String getToken() {
+        return token;
+    }
+}