/**
* Copyright 2016 LinkedIn Corp. All rights reserved.
*
* 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.
*/
package com.github.ambry.rest;
import com.github.ambry.config.VerifiableProperties;
import com.github.ambry.messageformat.BlobInfo;
import com.github.ambry.messageformat.BlobProperties;
import com.github.ambry.router.ByteBufferRSC;
import com.github.ambry.router.Callback;
import com.github.ambry.router.GetBlobOptions;
import com.github.ambry.router.GetBlobOptionsBuilder;
import com.github.ambry.router.GetBlobResult;
import com.github.ambry.router.ReadableStreamChannel;
import com.github.ambry.router.Router;
import com.github.ambry.router.RouterException;
import com.github.ambry.utils.Utils;
import java.io.IOException;
import java.nio.ByteBuffer;
import java.util.Date;
import java.util.GregorianCalendar;
import java.util.concurrent.CountDownLatch;
/**
* Implementation of the {@link BlobStorageService} that can be used in tests.
*/
public class MockBlobStorageService implements BlobStorageService {
public final static String ECHO_REST_METHOD = "mbssEchoRestMethod";
public final static String THROW_RUNTIME_EXCEPTION = "mbssRuntimeException";
public final static String SEND_RESPONSE_RUNTIME_EXCEPTION = "mbssResponseRuntimeException";
public final static String SEND_RESPONSE_REST_SERVICE_EXCEPTION = "mbssRestServiceException";
public final static String REST_ERROR_CODE = "mock.blob.storage.service.rest.error.code";
private final RestResponseHandler responseHandler;
private final Router router;
private VerifiableProperties verifiableProperties;
private volatile boolean serviceRunning = false;
private volatile boolean blocking = false;
private volatile CountDownLatch blockLatch = new CountDownLatch(0);
/**
* Changes the {@link VerifiableProperties} instance with this instance so that the behaviour can be changed on the
* fly.
* @param verifiableProperties the{@link VerifiableProperties} that will dictate behaviour.
*/
public void setVerifiableProperties(VerifiableProperties verifiableProperties) {
this.verifiableProperties = verifiableProperties;
}
/**
* Creates an instance of {@link MockBlobStorageService} with {@code router} as the backing {@link Router} and
* {@code verifiableProperties} defining the behavior of this instance.
* @param verifiableProperties the {@link VerifiableProperties} that defines the behavior of this instance.
* @param responseHandler the {@link RestResponseHandler} instance to use.
* @param router the {@link Router} that will back this instance.
*/
public MockBlobStorageService(VerifiableProperties verifiableProperties, RestResponseHandler responseHandler,
Router router) {
setVerifiableProperties(verifiableProperties);
this.responseHandler = responseHandler;
this.router = router;
}
@Override
public void start() throws InstantiationException {
serviceRunning = true;
}
@Override
public void shutdown() {
serviceRunning = false;
}
@Override
public void handleGet(RestRequest restRequest, RestResponseChannel restResponseChannel) {
if (shouldProceed(restRequest, restResponseChannel)) {
String blobId = getBlobId(restRequest);
MockGetCallback callback = new MockGetCallback(this, restRequest, restResponseChannel);
router.getBlob(blobId, new GetBlobOptionsBuilder().operationType(GetBlobOptions.OperationType.All).build(),
callback);
}
}
@Override
public void handlePost(RestRequest restRequest, RestResponseChannel restResponseChannel) {
if (shouldProceed(restRequest, restResponseChannel)) {
try {
BlobProperties blobProperties = RestUtils.buildBlobProperties(restRequest.getArgs());
byte[] usermetadata = RestUtils.buildUsermetadata(restRequest.getArgs());
router.putBlob(blobProperties, usermetadata, restRequest,
new MockPostCallback(this, restRequest, restResponseChannel, blobProperties));
} catch (RestServiceException e) {
handleResponse(restRequest, restResponseChannel, null, e);
}
}
}
/**
* {@inheritDoc}
* <p/>
* PUT is not supported by {@link MockBlobStorageService}.
* @param restRequest the {@link RestRequest} that needs to be handled.
* @param restResponseChannel the {@link RestResponseChannel} over which response to {@code restRequest} can be sent.
*/
@Override
public void handlePut(RestRequest restRequest, RestResponseChannel restResponseChannel) {
Exception exception = new RestServiceException("PUT is not supported", RestServiceErrorCode.UnsupportedHttpMethod);
handleResponse(restRequest, restResponseChannel, null, exception);
}
@Override
public void handleDelete(RestRequest restRequest, RestResponseChannel restResponseChannel) {
if (shouldProceed(restRequest, restResponseChannel)) {
String blobId = getBlobId(restRequest);
router.deleteBlob(blobId, null, new MockDeleteCallback(this, restRequest, restResponseChannel));
}
}
@Override
public void handleHead(RestRequest restRequest, RestResponseChannel restResponseChannel) {
if (shouldProceed(restRequest, restResponseChannel)) {
String blobId = getBlobId(restRequest);
router.getBlob(blobId, new GetBlobOptionsBuilder().operationType(GetBlobOptions.OperationType.BlobInfo).build(),
new MockHeadCallback(this, restRequest, restResponseChannel));
}
}
/**
* All operations block until {@link #releaseAllOperations()} is called.
* @throws IllegalStateException each call to this function must (eventually) be followed by a call to
* {@link #releaseAllOperations()}. If this function is invoked more than once before an
* accompanying {@link #releaseAllOperations()} is called, it is illegal state.
*/
public void blockAllOperations() {
if (blocking) {
throw new IllegalStateException("Already in blocking state");
} else {
blocking = true;
blockLatch = new CountDownLatch(1);
}
}
/**
* Releases all blocked operations.
*/
public void releaseAllOperations() {
blockLatch.countDown();
blocking = false;
}
/**
* Handles argument pre-checks and examines the URL to see if any custom operations need to be performed (which might
* involve throwing exceptions).
* <p/>
* Also blocks if required.
* @param restRequest the {@link RestRequest} that needs to be handled.
* @param restResponseChannel the {@link RestResponseChannel} that can be used to set headers.
* @return {@code true} if the pre-checks decided it is OK to continue. Otherwise {@code false}.
*/
private boolean shouldProceed(RestRequest restRequest, RestResponseChannel restResponseChannel) {
if (blocking) {
try {
blockLatch.await();
} catch (InterruptedException e) {
throw new IllegalStateException(e);
}
}
boolean shouldProceed = canHonorRequest(restRequest, restResponseChannel);
if (shouldProceed) {
String uri = restRequest.getUri();
ReadableStreamChannel response = null;
Exception exception = null;
if (uri.startsWith(ECHO_REST_METHOD)) {
String responseStr = restRequest.getRestMethod().toString() + uri.substring(ECHO_REST_METHOD.length());
ByteBuffer buffer = ByteBuffer.wrap(responseStr.getBytes());
response = new ByteBufferRSC(buffer);
shouldProceed = false;
} else if (THROW_RUNTIME_EXCEPTION.equals(uri)) {
throw new RuntimeException(THROW_RUNTIME_EXCEPTION);
} else if (SEND_RESPONSE_RUNTIME_EXCEPTION.equals(uri)) {
shouldProceed = false;
exception = new RuntimeException(SEND_RESPONSE_RUNTIME_EXCEPTION);
} else if (SEND_RESPONSE_REST_SERVICE_EXCEPTION.equals(uri)) {
shouldProceed = false;
RestServiceErrorCode errorCode = RestServiceErrorCode.InternalServerError;
try {
errorCode = RestServiceErrorCode.valueOf(verifiableProperties.getString(REST_ERROR_CODE));
} catch (IllegalArgumentException e) {
// it's alright.
}
exception = new RestServiceException(SEND_RESPONSE_REST_SERVICE_EXCEPTION, errorCode);
}
if (!shouldProceed) {
try {
if (exception == null) {
restResponseChannel.setStatus(ResponseStatus.Ok);
}
} catch (RestServiceException e) {
exception = e;
} finally {
handleResponse(restRequest, restResponseChannel, response, exception);
}
}
}
return shouldProceed;
}
/**
* Performs null pre checks and checks that the service is running.
* @param restRequest the {@link RestRequest} that needs to be handled.
* @param restResponseChannel the {@link RestResponseChannel} that can be used to set headers.
* @throws IllegalArgumentException if either of {@code restRequest} or {@code restResponseChannel}is null.
* @return {@code true} if the the service has the right arguments and is in a state to honor the request. Otherwise
* {@code false}.
*/
private boolean canHonorRequest(RestRequest restRequest, RestResponseChannel restResponseChannel) {
if (restRequest == null || restResponseChannel == null) {
throw new IllegalArgumentException("One of the arguments was null");
} else if (!serviceRunning) {
handleResponse(restRequest, restResponseChannel, null,
new RestServiceException("BlobStorageService not running", RestServiceErrorCode.ServiceUnavailable));
}
return serviceRunning;
}
/**
* Determines the blob ID desired by the request.
* @param restRequest a {@link RestRequest} that represents the request.
* @return the blob ID desired by the request.
*/
protected static String getBlobId(RestRequest restRequest) {
String path = restRequest.getPath();
return path.startsWith("/") ? path.substring(1, path.length()) : path;
}
/**
* Sends out responses immediately. Returns control only after {@code response} has been consumed and sent over the
* {@code restResponseChannel}.
* <p/>
* If the response building was unsuccessful for any reason, the details are included in the {@code exception}.
* <p/>
* The bytes consumed from the {@code response} are streamed out (unmodified) through the {@code restResponseChannel}.
* <p/>
* Assumed that at least one of {@code response} or {@code exception} is null.
* @param restRequest the {@link RestRequest} for which the response has been constructed.
* @param restResponseChannel the {@link RestResponseChannel} to be used to send the response.
* @param response a {@link ReadableStreamChannel} that represents the response to the {@code restRequest}.
* @param exception if the response could not be constructed, the reason for the failure.
*/
protected void handleResponse(RestRequest restRequest, RestResponseChannel restResponseChannel,
ReadableStreamChannel response, Exception exception) {
try {
if (exception != null && exception instanceof RouterException) {
exception = new RestServiceException(exception,
RestServiceErrorCode.getRestServiceErrorCode(((RouterException) exception).getErrorCode()));
}
responseHandler.handleResponse(restRequest, restResponseChannel, response, exception);
} catch (RestServiceException e) {
exception = exception == null ? e : exception;
restResponseChannel.onResponseComplete(exception);
if (response != null) {
try {
response.close();
} catch (IOException ioe) {
throw new IllegalStateException(ioe);
}
}
}
}
}
/**
* Callback for GET operations. Updates headers and submits response.
*/
class MockGetCallback implements Callback<GetBlobResult> {
private final MockBlobStorageService mockBlobStorageService;
private final RestRequest restRequest;
private final RestResponseChannel restResponseChannel;
/**
* Create a GET callback.
* @param mockBlobStorageService the {@link MockBlobStorageService} to use to submit responses.
* @param restRequest the {@link RestRequest} for whose response this is a callback.
* @param restResponseChannel the {@link RestResponseChannel} to set headers on.
*/
public MockGetCallback(MockBlobStorageService mockBlobStorageService, RestRequest restRequest,
RestResponseChannel restResponseChannel) {
this.mockBlobStorageService = mockBlobStorageService;
this.restRequest = restRequest;
this.restResponseChannel = restResponseChannel;
}
/**
* If there was no exception, sets headers and submits response.
* @param result The result of the request - a {@link GetBlobResult} object with the {@link BlobInfo} containing the
* blob properties and other headers of the blob, and the {@link ReadableStreamChannel} of blob data.
* This is non null if the request executed successfully.
* @param exception The exception that was reported on execution of the request (if any).
*/
@Override
public void onCompletion(GetBlobResult result, Exception exception) {
try {
restResponseChannel.setHeader(RestUtils.Headers.DATE, new GregorianCalendar().getTime());
if (exception == null && result != null) {
setResponseHeaders(result.getBlobInfo());
} else if (exception != null && exception instanceof RouterException) {
exception = new RestServiceException(exception,
RestServiceErrorCode.getRestServiceErrorCode(((RouterException) exception).getErrorCode()));
}
} catch (Exception e) {
exception = exception == null ? e : exception;
} finally {
ReadableStreamChannel channel = result != null ? result.getBlobDataChannel() : null;
mockBlobStorageService.handleResponse(restRequest, restResponseChannel, channel, exception);
}
}
/**
* Sets the required headers in the response.
* @param blobInfo the {@link BlobInfo} to refer to while setting headers.
* @throws RestServiceException if there was any problem setting the headers.
*/
private void setResponseHeaders(BlobInfo blobInfo) throws RestServiceException {
BlobProperties blobProperties = blobInfo.getBlobProperties();
restResponseChannel.setHeader(RestUtils.Headers.LAST_MODIFIED, new Date(blobProperties.getCreationTimeInMs()));
restResponseChannel.setHeader(RestUtils.Headers.BLOB_SIZE, blobProperties.getBlobSize());
if (blobProperties.getContentType() != null) {
restResponseChannel.setHeader(RestUtils.Headers.CONTENT_TYPE, blobProperties.getContentType());
if (blobProperties.getContentType().equals("text/html")) {
restResponseChannel.setHeader("Content-Disposition", "attachment");
}
}
}
}
/**
* Callback for POST operations. Sends the blob ID of the created blob to the client.
*/
class MockPostCallback implements Callback<String> {
private final MockBlobStorageService mockBlobStorageService;
private final RestRequest restRequest;
private final RestResponseChannel restResponseChannel;
private final BlobProperties blobProperties;
/**
* Create a POST callback.
* @param mockBlobStorageService the {@link MockBlobStorageService} to use to submit responses.
* @param restRequest the {@link RestRequest} for whose response this is a callback.
* @param restResponseChannel the {@link RestResponseChannel} over which response to {@code restRequest} can be sent.
* @param createdBlobProperties the {@link BlobProperties} of the blob that was asked to be POSTed.
*/
public MockPostCallback(MockBlobStorageService mockBlobStorageService, RestRequest restRequest,
RestResponseChannel restResponseChannel, BlobProperties createdBlobProperties) {
this.mockBlobStorageService = mockBlobStorageService;
this.restRequest = restRequest;
this.restResponseChannel = restResponseChannel;
this.blobProperties = createdBlobProperties;
}
/**
* If there was no exception, updates the header with the location of the object.
* @param result The result of the request. This is the blob ID of the blob. This is non null if the request executed
* successfully.
* @param exception The exception that was reported on execution of the request (if any).
*/
@Override
public void onCompletion(String result, Exception exception) {
try {
restResponseChannel.setHeader(RestUtils.Headers.DATE, new GregorianCalendar().getTime());
if (exception == null && result != null) {
setResponseHeaders(result);
} else if (exception != null && exception instanceof RouterException) {
exception = new RestServiceException(exception,
RestServiceErrorCode.getRestServiceErrorCode(((RouterException) exception).getErrorCode()));
}
} catch (Exception e) {
exception = exception == null ? e : exception;
} finally {
mockBlobStorageService.handleResponse(restRequest, restResponseChannel, null, exception);
}
}
/**
* Sets the required headers in the response.
* @param location the location of the created resource.
* @throws RestServiceException if there was any problem setting the headers.
*/
private void setResponseHeaders(String location) throws RestServiceException {
restResponseChannel.setStatus(ResponseStatus.Created);
restResponseChannel.setHeader(RestUtils.Headers.LOCATION, location);
restResponseChannel.setHeader(RestUtils.Headers.CONTENT_LENGTH, 0);
restResponseChannel.setHeader(RestUtils.Headers.CREATION_TIME, new Date(blobProperties.getCreationTimeInMs()));
}
}
/**
* Callback for DELETE operations. Sends an ACCEPTED response to the client if operation is successful.
*/
class MockDeleteCallback implements Callback<Void> {
private final MockBlobStorageService mockBlobStorageService;
private final RestRequest restRequest;
private final RestResponseChannel restResponseChannel;
/**
* Create a DELETE callback.
* @param mockBlobStorageService the {@link MockBlobStorageService} to use to submit responses.
* @param restRequest the {@link RestRequest} for whose response this is a callback.
* @param restResponseChannel the {@link RestResponseChannel} over which response to {@code restRequest} can be sent.
*/
public MockDeleteCallback(MockBlobStorageService mockBlobStorageService, RestRequest restRequest,
RestResponseChannel restResponseChannel) {
this.mockBlobStorageService = mockBlobStorageService;
this.restRequest = restRequest;
this.restResponseChannel = restResponseChannel;
}
/**
* If there was no exception, updates the header with the acceptance of the request.
* @param result The result of the request. This is always null.
* @param exception The exception that was reported on execution of the request (if any).
*/
@Override
public void onCompletion(Void result, Exception exception) {
try {
restResponseChannel.setHeader(RestUtils.Headers.DATE, new GregorianCalendar().getTime());
if (exception == null) {
restResponseChannel.setStatus(ResponseStatus.Accepted);
restResponseChannel.setHeader(RestUtils.Headers.CONTENT_LENGTH, 0);
} else if (exception instanceof RouterException) {
exception = new RestServiceException(exception,
RestServiceErrorCode.getRestServiceErrorCode(((RouterException) exception).getErrorCode()));
}
} catch (Exception e) {
exception = exception == null ? e : exception;
} finally {
mockBlobStorageService.handleResponse(restRequest, restResponseChannel, null, exception);
}
}
}
/**
* Callback for HEAD operations. Sends the headers to the client if operation is successful.
*/
class MockHeadCallback implements Callback<GetBlobResult> {
private final MockBlobStorageService mockBlobStorageService;
private final RestRequest restRequest;
private final RestResponseChannel restResponseChannel;
/**
* Create a HEAD callback.
* @param mockBlobStorageService the {@link MockBlobStorageService} to use to submit responses.
* @param restRequest the {@link RestRequest} for whose response this is a callback.
* @param restResponseChannel the {@link RestResponseChannel} over which response to {@code restRequest} can be sent.
*/
public MockHeadCallback(MockBlobStorageService mockBlobStorageService, RestRequest restRequest,
RestResponseChannel restResponseChannel) {
this.mockBlobStorageService = mockBlobStorageService;
this.restRequest = restRequest;
this.restResponseChannel = restResponseChannel;
}
/**
* If there was no exception, updates the header with the properties. Exceptions, if any, will be handled upon
* submission.
* @param result The result of the request i.e a {@link GetBlobResult} object with the properties of the blob. This is
* non null if the request executed successfully.
* @param exception The exception that was reported on execution of the request (if any).
*/
@Override
public void onCompletion(GetBlobResult result, Exception exception) {
try {
restResponseChannel.setHeader(RestUtils.Headers.DATE, new GregorianCalendar().getTime());
if (exception == null && result != null) {
setBlobPropertiesResponseHeaders(result.getBlobInfo());
} else if (exception != null && exception instanceof RouterException) {
exception = new RestServiceException(exception,
RestServiceErrorCode.getRestServiceErrorCode(((RouterException) exception).getErrorCode()));
}
} catch (Exception e) {
exception = exception == null ? e : exception;
} finally {
mockBlobStorageService.handleResponse(restRequest, restResponseChannel, null, exception);
}
}
/**
* Sets the required blob properties headers in the response.
* @param blobInfo the {@link BlobInfo} to refer to while setting headers.
* @throws RestServiceException if there was any problem setting the headers.
*/
private void setBlobPropertiesResponseHeaders(BlobInfo blobInfo) throws RestServiceException {
BlobProperties blobProperties = blobInfo.getBlobProperties();
restResponseChannel.setHeader(RestUtils.Headers.LAST_MODIFIED, new Date(blobProperties.getCreationTimeInMs()));
restResponseChannel.setHeader(RestUtils.Headers.CONTENT_LENGTH, blobProperties.getBlobSize());
// Blob props
restResponseChannel.setHeader(RestUtils.Headers.BLOB_SIZE, blobProperties.getBlobSize());
restResponseChannel.setHeader(RestUtils.Headers.SERVICE_ID, blobProperties.getServiceId());
restResponseChannel.setHeader(RestUtils.Headers.CREATION_TIME, new Date(blobProperties.getCreationTimeInMs()));
restResponseChannel.setHeader(RestUtils.Headers.PRIVATE, blobProperties.isPrivate());
if (blobProperties.getTimeToLiveInSeconds() != Utils.Infinite_Time) {
restResponseChannel.setHeader(RestUtils.Headers.TTL, Long.toString(blobProperties.getTimeToLiveInSeconds()));
}
if (blobProperties.getContentType() != null) {
restResponseChannel.setHeader(RestUtils.Headers.AMBRY_CONTENT_TYPE, blobProperties.getContentType());
restResponseChannel.setHeader(RestUtils.Headers.CONTENT_TYPE, blobProperties.getContentType());
}
if (blobProperties.getOwnerId() != null) {
restResponseChannel.setHeader(RestUtils.Headers.OWNER_ID, blobProperties.getOwnerId());
}
}
}