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 index 2a8ce6c5..8d9b28fe 100644 --- 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 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 index fae4979f..a1018746 100644 --- 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 new file mode 100644 index 00000000..16c08e90 --- /dev/null +++ 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; + } +}