|
|
@@ -19,14 +19,20 @@
|
|
|
package org.apache.hadoop.yarn.server.timelineservice.documentstore.writer.cosmosdb;
|
|
|
|
|
|
|
|
|
-import com.microsoft.azure.documentdb.AccessCondition;
|
|
|
-import com.microsoft.azure.documentdb.AccessConditionType;
|
|
|
-import com.microsoft.azure.documentdb.Database;
|
|
|
-import com.microsoft.azure.documentdb.Document;
|
|
|
-import com.microsoft.azure.documentdb.DocumentClient;
|
|
|
-import com.microsoft.azure.documentdb.DocumentClientException;
|
|
|
-import com.microsoft.azure.documentdb.DocumentCollection;
|
|
|
-import com.microsoft.azure.documentdb.RequestOptions;
|
|
|
+import com.google.common.annotations.VisibleForTesting;
|
|
|
+import com.microsoft.azure.cosmosdb.AccessCondition;
|
|
|
+import com.microsoft.azure.cosmosdb.AccessConditionType;
|
|
|
+import com.microsoft.azure.cosmosdb.Database;
|
|
|
+import com.microsoft.azure.cosmosdb.Document;
|
|
|
+import com.microsoft.azure.cosmosdb.DocumentClientException;
|
|
|
+import com.microsoft.azure.cosmosdb.DocumentCollection;
|
|
|
+import com.microsoft.azure.cosmosdb.FeedResponse;
|
|
|
+import com.microsoft.azure.cosmosdb.RequestOptions;
|
|
|
+import com.microsoft.azure.cosmosdb.ResourceResponse;
|
|
|
+import com.microsoft.azure.cosmosdb.SqlParameter;
|
|
|
+import com.microsoft.azure.cosmosdb.SqlParameterCollection;
|
|
|
+import com.microsoft.azure.cosmosdb.SqlQuerySpec;
|
|
|
+import com.microsoft.azure.cosmosdb.rx.AsyncDocumentClient;
|
|
|
import org.apache.hadoop.conf.Configuration;
|
|
|
import org.apache.hadoop.util.Time;
|
|
|
import org.apache.hadoop.yarn.server.timelineservice.metrics.PerNodeAggTimelineCollectorMetrics;
|
|
|
@@ -40,6 +46,13 @@ import org.apache.hadoop.yarn.server.timelineservice.documentstore.lib.DocumentS
|
|
|
import org.apache.hadoop.yarn.server.timelineservice.documentstore.writer.DocumentStoreWriter;
|
|
|
import org.slf4j.Logger;
|
|
|
import org.slf4j.LoggerFactory;
|
|
|
+import rx.Observable;
|
|
|
+import rx.Scheduler;
|
|
|
+import rx.functions.Func1;
|
|
|
+import rx.schedulers.Schedulers;
|
|
|
+
|
|
|
+import java.util.concurrent.ExecutorService;
|
|
|
+import java.util.concurrent.Executors;
|
|
|
|
|
|
/**
|
|
|
* This is the Document Store Writer implementation for
|
|
|
@@ -51,79 +64,102 @@ public class CosmosDBDocumentStoreWriter<TimelineDoc extends TimelineDocument>
|
|
|
private static final Logger LOG = LoggerFactory
|
|
|
.getLogger(CosmosDBDocumentStoreWriter.class);
|
|
|
|
|
|
- private static volatile DocumentClient client;
|
|
|
private final String databaseName;
|
|
|
private static final PerNodeAggTimelineCollectorMetrics METRICS =
|
|
|
PerNodeAggTimelineCollectorMetrics.getInstance();
|
|
|
+
|
|
|
+ private static AsyncDocumentClient client;
|
|
|
+ // creating thread pool of size equal to number of collection types
|
|
|
+ private ExecutorService executorService =
|
|
|
+ Executors.newFixedThreadPool(CollectionType.values().length);
|
|
|
+ private Scheduler schedulerForBlockingWork =
|
|
|
+ Schedulers.from(executorService);
|
|
|
+
|
|
|
private static final String DATABASE_LINK = "/dbs/%s";
|
|
|
private static final String COLLECTION_LINK = DATABASE_LINK + "/colls/%s";
|
|
|
private static final String DOCUMENT_LINK = COLLECTION_LINK + "/docs/%s";
|
|
|
+ private static final String ID = "@id";
|
|
|
+ private static final String QUERY_COLLECTION_IF_EXISTS = "SELECT * FROM r " +
|
|
|
+ "where r.id = " + ID;
|
|
|
|
|
|
public CosmosDBDocumentStoreWriter(Configuration conf) {
|
|
|
LOG.info("Initializing Cosmos DB DocumentStoreWriter...");
|
|
|
databaseName = DocumentStoreUtils.getCosmosDBDatabaseName(conf);
|
|
|
- // making CosmosDB Client Singleton
|
|
|
+ initCosmosDBClient(conf);
|
|
|
+ }
|
|
|
+
|
|
|
+ private synchronized void initCosmosDBClient(Configuration conf) {
|
|
|
+ // making CosmosDB Async Client Singleton
|
|
|
if (client == null) {
|
|
|
- synchronized (this) {
|
|
|
- if (client == null) {
|
|
|
- LOG.info("Creating Cosmos DB Client...");
|
|
|
- client = DocumentStoreUtils.createCosmosDBClient(conf);
|
|
|
- }
|
|
|
- }
|
|
|
+ LOG.info("Creating Cosmos DB Writer Async Client...");
|
|
|
+ client = DocumentStoreUtils.createCosmosDBAsyncClient(conf);
|
|
|
+ addShutdownHook();
|
|
|
}
|
|
|
}
|
|
|
|
|
|
@Override
|
|
|
public void createDatabase() {
|
|
|
- try {
|
|
|
- client.readDatabase(String.format(
|
|
|
- DATABASE_LINK, databaseName), new RequestOptions());
|
|
|
- LOG.info("Database {} already exists.", databaseName);
|
|
|
- } catch (DocumentClientException docExceptionOnRead) {
|
|
|
- if (docExceptionOnRead.getStatusCode() == 404) {
|
|
|
- LOG.info("Creating new Database : {}", databaseName);
|
|
|
- Database databaseDefinition = new Database();
|
|
|
- databaseDefinition.setId(databaseName);
|
|
|
- try {
|
|
|
- client.createDatabase(databaseDefinition, new RequestOptions());
|
|
|
- } catch (DocumentClientException docExceptionOnCreate) {
|
|
|
- LOG.error("Unable to create new Database : {}", databaseName,
|
|
|
- docExceptionOnCreate);
|
|
|
- }
|
|
|
- } else {
|
|
|
- LOG.error("Error while reading Database : {}", databaseName,
|
|
|
- docExceptionOnRead);
|
|
|
- }
|
|
|
- }
|
|
|
+ Observable<ResourceResponse<Database>> databaseReadObs =
|
|
|
+ client.readDatabase(String.format(DATABASE_LINK, databaseName), null);
|
|
|
+
|
|
|
+ Observable<ResourceResponse<Database>> databaseExistenceObs =
|
|
|
+ databaseReadObs
|
|
|
+ .doOnNext(databaseResourceResponse ->
|
|
|
+ LOG.info("Database {} already exists.", databaseName))
|
|
|
+ .onErrorResumeNext(throwable -> {
|
|
|
+ // if the database doesn't exists
|
|
|
+ // readDatabase() will result in 404 error
|
|
|
+ if (throwable instanceof DocumentClientException) {
|
|
|
+ DocumentClientException de =
|
|
|
+ (DocumentClientException) throwable;
|
|
|
+ if (de.getStatusCode() == 404) {
|
|
|
+ // if the database doesn't exist, create it.
|
|
|
+ LOG.info("Creating new Database : {}", databaseName);
|
|
|
+
|
|
|
+ Database dbDefinition = new Database();
|
|
|
+ dbDefinition.setId(databaseName);
|
|
|
+
|
|
|
+ return client.createDatabase(dbDefinition, null);
|
|
|
+ }
|
|
|
+ }
|
|
|
+ // some unexpected failure in reading database happened.
|
|
|
+ // pass the error up.
|
|
|
+ LOG.error("Reading database : {} if it exists failed.",
|
|
|
+ databaseName, throwable);
|
|
|
+ return Observable.error(throwable);
|
|
|
+ });
|
|
|
+ // wait for completion
|
|
|
+ databaseExistenceObs.toCompletable().await();
|
|
|
}
|
|
|
|
|
|
@Override
|
|
|
public void createCollection(final String collectionName) {
|
|
|
LOG.info("Creating Timeline Collection : {} for Database : {}",
|
|
|
collectionName, databaseName);
|
|
|
- try {
|
|
|
- client.readCollection(String.format(COLLECTION_LINK, databaseName,
|
|
|
- collectionName), new RequestOptions());
|
|
|
- LOG.info("Collection {} already exists.", collectionName);
|
|
|
- } catch (DocumentClientException docExceptionOnRead) {
|
|
|
- if (docExceptionOnRead.getStatusCode() == 404) {
|
|
|
- DocumentCollection collection = new DocumentCollection();
|
|
|
- collection.setId(collectionName);
|
|
|
- LOG.info("Creating collection {} under Database {}",
|
|
|
- collectionName, databaseName);
|
|
|
- try {
|
|
|
- client.createCollection(
|
|
|
- String.format(DATABASE_LINK, databaseName),
|
|
|
- collection, new RequestOptions());
|
|
|
- } catch (DocumentClientException docExceptionOnCreate) {
|
|
|
- LOG.error("Unable to create Collection : {} under Database : {}",
|
|
|
- collectionName, databaseName, docExceptionOnCreate);
|
|
|
- }
|
|
|
- } else {
|
|
|
- LOG.error("Error while reading Collection : {} under Database : {}",
|
|
|
- collectionName, databaseName, docExceptionOnRead);
|
|
|
- }
|
|
|
- }
|
|
|
+ client.queryCollections(String.format(DATABASE_LINK, databaseName),
|
|
|
+ new SqlQuerySpec(QUERY_COLLECTION_IF_EXISTS,
|
|
|
+ new SqlParameterCollection(
|
|
|
+ new SqlParameter(ID, collectionName))), null)
|
|
|
+ .single() // there should be single page of result
|
|
|
+ .flatMap((Func1<FeedResponse<DocumentCollection>, Observable<?>>)
|
|
|
+ page -> {
|
|
|
+ if (page.getResults().isEmpty()) {
|
|
|
+ // if there is no matching collection create one.
|
|
|
+ DocumentCollection collection = new DocumentCollection();
|
|
|
+ collection.setId(collectionName);
|
|
|
+ LOG.info("Creating collection {}", collectionName);
|
|
|
+ return client.createCollection(
|
|
|
+ String.format(DATABASE_LINK, databaseName),
|
|
|
+ collection, null);
|
|
|
+ } else {
|
|
|
+ // collection already exists, nothing else to be done.
|
|
|
+ LOG.info("Collection {} already exists.", collectionName);
|
|
|
+ return Observable.empty();
|
|
|
+ }
|
|
|
+ })
|
|
|
+ .doOnError(throwable -> LOG.error("Unable to create collection : {}",
|
|
|
+ collectionName, throwable))
|
|
|
+ .toCompletable().await();
|
|
|
}
|
|
|
|
|
|
@Override
|
|
|
@@ -156,32 +192,40 @@ public class CosmosDBDocumentStoreWriter<TimelineDoc extends TimelineDocument>
|
|
|
AccessCondition accessCondition = new AccessCondition();
|
|
|
StringBuilder eTagStrBuilder = new StringBuilder();
|
|
|
|
|
|
- TimelineDoc updatedTimelineDoc = applyUpdatesOnPrevDoc(collectionType,
|
|
|
+ final TimelineDoc updatedTimelineDoc = applyUpdatesOnPrevDoc(collectionType,
|
|
|
timelineDoc, eTagStrBuilder);
|
|
|
|
|
|
accessCondition.setCondition(eTagStrBuilder.toString());
|
|
|
accessCondition.setType(AccessConditionType.IfMatch);
|
|
|
requestOptions.setAccessCondition(accessCondition);
|
|
|
|
|
|
- try {
|
|
|
- client.upsertDocument(collectionLink, updatedTimelineDoc,
|
|
|
- requestOptions, true);
|
|
|
+ ResourceResponse<Document> resourceResponse =
|
|
|
+ client.upsertDocument(collectionLink, updatedTimelineDoc,
|
|
|
+ requestOptions, true)
|
|
|
+ .subscribeOn(schedulerForBlockingWork)
|
|
|
+ .doOnError(throwable ->
|
|
|
+ LOG.error("Error while upserting Collection : {} " +
|
|
|
+ "with Doc Id : {} under Database : {}",
|
|
|
+ collectionType.getCollectionName(),
|
|
|
+ updatedTimelineDoc.getId(), databaseName, throwable))
|
|
|
+ .toBlocking()
|
|
|
+ .single();
|
|
|
+
|
|
|
+ if (resourceResponse.getStatusCode() == 409) {
|
|
|
+ LOG.warn("There was a conflict while upserting, hence retrying...",
|
|
|
+ resourceResponse);
|
|
|
+ upsertDocument(collectionType, updatedTimelineDoc);
|
|
|
+ } else if (resourceResponse.getStatusCode() >= 200 && resourceResponse
|
|
|
+ .getStatusCode() < 300) {
|
|
|
LOG.debug("Successfully wrote doc with id : {} and type : {} under " +
|
|
|
"Database : {}", timelineDoc.getId(), timelineDoc.getType(),
|
|
|
databaseName);
|
|
|
- } catch (DocumentClientException e) {
|
|
|
- if (e.getStatusCode() == 409) {
|
|
|
- LOG.warn("There was a conflict while upserting, hence retrying...", e);
|
|
|
- upsertDocument(collectionType, updatedTimelineDoc);
|
|
|
- }
|
|
|
- LOG.error("Error while upserting Collection : {} with Doc Id : {} under" +
|
|
|
- " Database : {}", collectionType.getCollectionName(),
|
|
|
- updatedTimelineDoc.getId(), databaseName, e);
|
|
|
}
|
|
|
}
|
|
|
|
|
|
+ @VisibleForTesting
|
|
|
@SuppressWarnings("unchecked")
|
|
|
- private TimelineDoc applyUpdatesOnPrevDoc(CollectionType collectionType,
|
|
|
+ TimelineDoc applyUpdatesOnPrevDoc(CollectionType collectionType,
|
|
|
TimelineDoc timelineDoc, StringBuilder eTagStrBuilder) {
|
|
|
TimelineDoc prevDocument = fetchLatestDoc(collectionType,
|
|
|
timelineDoc.getId(), eTagStrBuilder);
|
|
|
@@ -192,14 +236,15 @@ public class CosmosDBDocumentStoreWriter<TimelineDoc extends TimelineDocument>
|
|
|
return timelineDoc;
|
|
|
}
|
|
|
|
|
|
+ @VisibleForTesting
|
|
|
@SuppressWarnings("unchecked")
|
|
|
- private TimelineDoc fetchLatestDoc(final CollectionType collectionType,
|
|
|
+ TimelineDoc fetchLatestDoc(final CollectionType collectionType,
|
|
|
final String documentId, StringBuilder eTagStrBuilder) {
|
|
|
final String documentLink = String.format(DOCUMENT_LINK, databaseName,
|
|
|
collectionType.getCollectionName(), documentId);
|
|
|
try {
|
|
|
Document latestDocument = client.readDocument(documentLink, new
|
|
|
- RequestOptions()).getResource();
|
|
|
+ RequestOptions()).toBlocking().single().getResource();
|
|
|
TimelineDoc timelineDoc;
|
|
|
switch (collectionType) {
|
|
|
case FLOW_RUN:
|
|
|
@@ -227,9 +272,17 @@ public class CosmosDBDocumentStoreWriter<TimelineDoc extends TimelineDocument>
|
|
|
@Override
|
|
|
public synchronized void close() {
|
|
|
if (client != null) {
|
|
|
- LOG.info("Closing Cosmos DB Client...");
|
|
|
+ LOG.info("Closing Cosmos DB Writer Async Client...");
|
|
|
client.close();
|
|
|
client = null;
|
|
|
}
|
|
|
}
|
|
|
+
|
|
|
+ private void addShutdownHook() {
|
|
|
+ Runtime.getRuntime().addShutdownHook(new Thread(() -> {
|
|
|
+ if (executorService != null) {
|
|
|
+ executorService.shutdown();
|
|
|
+ }
|
|
|
+ }));
|
|
|
+ }
|
|
|
}
|
Akira Ajisaka