Skip to content

Commit

Permalink
Support SKU spec upgrade (#459)
Browse files Browse the repository at this point in the history
* Support sku spec upgrade

* separate UpgradeReq to workflow

* annotate nullable fields
  • Loading branch information
Andyz26 authored Jun 13, 2023
1 parent 0936beb commit 16ae78d
Show file tree
Hide file tree
Showing 8 changed files with 153 additions and 11 deletions.
Original file line number Diff line number Diff line change
Expand Up @@ -38,8 +38,10 @@
import io.mantisrx.master.resourcecluster.proto.ResourceClusterScaleSpec;
import io.mantisrx.master.resourcecluster.proto.ScaleResourceRequest;
import io.mantisrx.master.resourcecluster.proto.UpgradeClusterContainersRequest;
import io.mantisrx.master.resourcecluster.proto.UpgradeClusterContainersResponse;
import io.mantisrx.master.resourcecluster.resourceprovider.InMemoryOnlyResourceClusterStorageProvider;
import io.mantisrx.master.resourcecluster.resourceprovider.ResourceClusterProvider;
import io.mantisrx.master.resourcecluster.resourceprovider.ResourceClusterProviderUpgradeRequest;
import io.mantisrx.master.resourcecluster.resourceprovider.ResourceClusterStorageProvider;
import io.mantisrx.master.resourcecluster.writable.ResourceClusterScaleRulesWritable;
import io.mantisrx.master.resourcecluster.writable.ResourceClusterScaleRulesWritable.ResourceClusterScaleRulesWritableBuilder;
Expand Down Expand Up @@ -328,10 +330,37 @@ private void onScaleResourceClusterRequest(ScaleResourceRequest req) {
private void onUpgradeClusterContainersRequest(UpgradeClusterContainersRequest req) {
log.info("Entering onScaleResourceClusterRequest: " + req);
// [Notes] for scaling-up the request can go straight into provider to increase desire size.
// FOr scaling-down the decision requires getting idle hosts first.

pipe(this.resourceClusterProvider.upgradeContainerResource(req), getContext().dispatcher()).to(getSender());
// For scaling-down the decision requires getting idle hosts first.
// if enableSkuSpecUpgrade is true, first fetch the latest spec to override the sku spec during upgrade
// workflow.

CompletionStage<UpgradeClusterContainersResponse> upgradeFut;
if (req.isEnableSkuSpecUpgrade()) {
upgradeFut = this.resourceClusterStorageProvider.getResourceClusterSpecWritable(req.getClusterId())
.thenCompose(specW -> {
if (specW == null) {
return CompletableFuture.completedFuture(UpgradeClusterContainersResponse.builder()
.responseCode(ResponseCode.CLIENT_ERROR_NOT_FOUND)
.build());
}

ResourceClusterProviderUpgradeRequest enrichedReq =
ResourceClusterProviderUpgradeRequest.from(req, specW.getClusterSpec());
return this.resourceClusterProvider.upgradeContainerResource(enrichedReq);
})
.exceptionally(err ->
UpgradeClusterContainersResponse.builder()
.responseCode(ResponseCode.SERVER_ERROR)
.message(err.getMessage())
.build());
}
else {
log.info("Upgrading cluster image only: {}", req.getClusterId());
upgradeFut =
this.resourceClusterProvider.upgradeContainerResource(ResourceClusterProviderUpgradeRequest.from(req));
}

pipe(upgradeFut, getContext().dispatcher()).to(getSender());
}

private static boolean validateClusterSpec(ProvisionResourceClusterRequest req) {
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -21,7 +21,7 @@
import lombok.Value;

@Value
@Builder
@Builder(toBuilder = true)
public class UpgradeClusterContainersRequest {
ClusterID clusterId;

Expand All @@ -36,4 +36,6 @@ public class UpgradeClusterContainersRequest {
int optionalBatchMaxSize;

boolean forceUpgradeOnSameImage;

boolean enableSkuSpecUpgrade;
}
Original file line number Diff line number Diff line change
Expand Up @@ -20,7 +20,6 @@
import io.mantisrx.master.resourcecluster.proto.ResourceClusterProvisionSubmissionResponse;
import io.mantisrx.master.resourcecluster.proto.ScaleResourceRequest;
import io.mantisrx.master.resourcecluster.proto.ScaleResourceResponse;
import io.mantisrx.master.resourcecluster.proto.UpgradeClusterContainersRequest;
import io.mantisrx.master.resourcecluster.proto.UpgradeClusterContainersResponse;
import java.util.concurrent.CompletableFuture;
import java.util.concurrent.CompletionStage;
Expand All @@ -40,7 +39,8 @@ public CompletionStage<ScaleResourceResponse> scaleResource(ScaleResourceRequest
}

@Override
public CompletionStage<UpgradeClusterContainersResponse> upgradeContainerResource(UpgradeClusterContainersRequest request) {
public CompletionStage<UpgradeClusterContainersResponse> upgradeContainerResource(
ResourceClusterProviderUpgradeRequest request) {
return CompletableFuture.completedFuture(null);
}

Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -20,7 +20,6 @@
import io.mantisrx.master.resourcecluster.proto.ResourceClusterProvisionSubmissionResponse;
import io.mantisrx.master.resourcecluster.proto.ScaleResourceRequest;
import io.mantisrx.master.resourcecluster.proto.ScaleResourceResponse;
import io.mantisrx.master.resourcecluster.proto.UpgradeClusterContainersRequest;
import io.mantisrx.master.resourcecluster.proto.UpgradeClusterContainersResponse;
import java.util.concurrent.CompletionStage;

Expand Down Expand Up @@ -59,7 +58,7 @@ CompletionStage<ResourceClusterProvisionSubmissionResponse> provisionClusterIfNo
* If multiple image digest versions need to be ran/hosted at the same time, it is recommended to create a separate
* sku id in addition to the existing sku(s).
*/
CompletionStage<UpgradeClusterContainersResponse> upgradeContainerResource(UpgradeClusterContainersRequest request);
CompletionStage<UpgradeClusterContainersResponse> upgradeContainerResource(ResourceClusterProviderUpgradeRequest request);

ResourceClusterResponseHandler getResponseHandler();
}
Original file line number Diff line number Diff line change
Expand Up @@ -21,7 +21,6 @@
import io.mantisrx.master.resourcecluster.proto.ResourceClusterProvisionSubmissionResponse;
import io.mantisrx.master.resourcecluster.proto.ScaleResourceRequest;
import io.mantisrx.master.resourcecluster.proto.ScaleResourceResponse;
import io.mantisrx.master.resourcecluster.proto.UpgradeClusterContainersRequest;
import io.mantisrx.master.resourcecluster.proto.UpgradeClusterContainersResponse;
import java.util.concurrent.CompletionStage;
import lombok.extern.slf4j.Slf4j;
Expand Down Expand Up @@ -76,7 +75,7 @@ public CompletionStage<ScaleResourceResponse> scaleResource(ScaleResourceRequest

@Override
public CompletionStage<UpgradeClusterContainersResponse> upgradeContainerResource(
UpgradeClusterContainersRequest request) {
ResourceClusterProviderUpgradeRequest request) {
return providerImpl.upgradeContainerResource(request);
}

Expand Down
Original file line number Diff line number Diff line change
@@ -0,0 +1,71 @@
/*
* Copyright 2023 Netflix, Inc.
*
* Licensed 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 io.mantisrx.master.resourcecluster.resourceprovider;

import io.mantisrx.master.resourcecluster.proto.MantisResourceClusterEnvType;
import io.mantisrx.master.resourcecluster.proto.MantisResourceClusterSpec;
import io.mantisrx.master.resourcecluster.proto.UpgradeClusterContainersRequest;
import io.mantisrx.server.master.resourcecluster.ClusterID;
import javax.annotation.Nullable;
import lombok.Builder;
import lombok.Value;

@Value
@Builder(toBuilder = true)
public class ResourceClusterProviderUpgradeRequest {
ClusterID clusterId;

String region;

@Nullable
String optionalImageId;

@Nullable
String optionalSkuId;

MantisResourceClusterEnvType optionalEnvType;

int optionalBatchMaxSize;

boolean forceUpgradeOnSameImage;

boolean enableSkuSpecUpgrade;

@Nullable
MantisResourceClusterSpec resourceClusterSpec;

public static ResourceClusterProviderUpgradeRequest from(
UpgradeClusterContainersRequest req) {
return from(req, null);
}

public static ResourceClusterProviderUpgradeRequest from(
UpgradeClusterContainersRequest req,
MantisResourceClusterSpec resourceClusterSpec) {
return ResourceClusterProviderUpgradeRequest.builder()
.clusterId(req.getClusterId())
.region(req.getRegion())
.optionalImageId(req.getOptionalImageId())
.optionalSkuId(req.getOptionalSkuId())
.optionalEnvType(req.getOptionalEnvType())
.optionalBatchMaxSize(req.getOptionalBatchMaxSize())
.forceUpgradeOnSameImage(req.isForceUpgradeOnSameImage())
.enableSkuSpecUpgrade(req.isEnableSkuSpecUpgrade())
.resourceClusterSpec(resourceClusterSpec)
.build();
}
}
Original file line number Diff line number Diff line change
Expand Up @@ -57,6 +57,7 @@
import io.mantisrx.master.resourcecluster.resourceprovider.NoopResourceClusterResponseHandler;
import io.mantisrx.master.resourcecluster.resourceprovider.ResourceClusterProvider;
import io.mantisrx.master.resourcecluster.resourceprovider.ResourceClusterProviderAdapter;
import io.mantisrx.master.resourcecluster.resourceprovider.ResourceClusterProviderUpgradeRequest;
import io.mantisrx.master.resourcecluster.resourceprovider.ResourceClusterResponseHandler;
import io.mantisrx.runtime.MachineDefinition;
import io.mantisrx.server.master.config.ConfigurationProvider;
Expand Down Expand Up @@ -468,7 +469,7 @@ public CompletionStage<ScaleResourceResponse> scaleResource(ScaleResourceRequest

@Override
public CompletionStage<UpgradeClusterContainersResponse> upgradeContainerResource(
UpgradeClusterContainersRequest request) {
ResourceClusterProviderUpgradeRequest request) {
return CompletableFuture.completedFuture(
UpgradeClusterContainersResponse.builder()
.message("test scale resp")
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -44,6 +44,7 @@
import io.mantisrx.master.resourcecluster.resourceprovider.ResourceClusterProvider;
import io.mantisrx.master.resourcecluster.resourceprovider.ResourceClusterResponseHandler;
import io.mantisrx.master.resourcecluster.resourceprovider.ResourceClusterStorageProvider;
import io.mantisrx.master.resourcecluster.writable.ResourceClusterSpecWritable;
import io.mantisrx.server.master.resourcecluster.ClusterID;
import io.mantisrx.server.master.resourcecluster.ContainerSkuID;
import java.util.concurrent.CompletableFuture;
Expand Down Expand Up @@ -468,6 +469,46 @@ public void testUpgradeRequest() {
probe.getSystem().stop(resourceClusterActor);
}

@Test
public void testUpgradeRequestEnableSkuSpecUpgrade() {
TestKit probe = new TestKit(system);
ResourceClusterStorageProvider resStorageProvider = mock(ResourceClusterStorageProvider.class);
ResourceClusterProvider resProvider = mock(ResourceClusterProvider.class);
ResourceClusterResponseHandler responseHandler = mock(ResourceClusterResponseHandler.class);

UpgradeClusterContainersResponse upgradeRes =
UpgradeClusterContainersResponse.builder().responseCode(ResponseCode.SUCCESS).build();
when(resProvider.upgradeContainerResource(any())).thenReturn(CompletableFuture.completedFuture(
upgradeRes
));

ProvisionResourceClusterRequest provisionReq = buildProvisionRequest();
when(resStorageProvider.getResourceClusterSpecWritable(any()))
.thenReturn(CompletableFuture.completedFuture(
ResourceClusterSpecWritable.builder()
.clusterSpec(provisionReq.getClusterSpec())
.id(provisionReq.getClusterId())
.build()));

when(resProvider.getResponseHandler()).thenReturn(responseHandler);

ActorRef resourceClusterActor = system.actorOf(
ResourceClustersHostManagerActor.props(resProvider, resStorageProvider));

UpgradeClusterContainersRequest request = UpgradeClusterContainersRequest.builder()
.clusterId(provisionReq.getClusterId())
.enableSkuSpecUpgrade(true)
.build();

resourceClusterActor.tell(request, probe.getRef());
UpgradeClusterContainersResponse createResp = probe.expectMsgClass(UpgradeClusterContainersResponse.class);

assertEquals(ResponseCode.SUCCESS, createResp.responseCode);

verify(resProvider, times(1)).upgradeContainerResource(any());
probe.getSystem().stop(resourceClusterActor);
}

private ProvisionResourceClusterRequest buildProvisionRequest() {
return buildProvisionRequest("mantisTestResCluster1", "[email protected]");
}
Expand Down

0 comments on commit 16ae78d

Please sign in to comment.