zhouyejoe commented on code in PR #35906:
URL: https://github.com/apache/spark/pull/35906#discussion_r923751792
##########
common/network-shuffle/src/main/java/org/apache/spark/network/shuffle/RemoteBlockPushResolver.java:
##########
@@ -316,22 +355,39 @@ public String[] getMergedBlockDirs(String appId) {
@Override
public void applicationRemoved(String appId, boolean cleanupLocalDirs) {
logger.info("Application {} removed, cleanupLocalDirs = {}", appId,
cleanupLocalDirs);
- AppShuffleInfo appShuffleInfo = appsShuffleInfo.remove(appId);
+ // Cleanup the DB within critical section to gain the consistency between
+ // DB and in-memory hashmap.
+ AtomicReference<AppShuffleInfo> ref = new AtomicReference<>(null);
+ appsShuffleInfo.compute(appId, (id, info) -> {
+ if (null != info) {
+ // Try cleaning up this application attempt local paths information
+ // and also the local paths information from former attempts in DB.
+ removeAppAttemptPathInfoFromDB(info.appId, info.attemptId);
+ if (info.attemptId != UNDEFINED_ATTEMPT_ID) {
+ for (int formerAttemptId = info.attemptId - 1; formerAttemptId >= 0;
formerAttemptId--) {
+ removeAppAttemptPathInfoFromDB(info.appId, formerAttemptId);
+ }
+ }
Review Comment:
Updated
##########
common/network-shuffle/src/main/java/org/apache/spark/network/shuffle/RemoteBlockPushResolver.java:
##########
@@ -632,6 +746,15 @@ public void registerExecutor(String appId,
ExecutorShuffleInfo executorInfo) {
appsShuffleInfo.compute(appId, (id, appShuffleInfo) -> {
if (appShuffleInfo == null || attemptId >
appShuffleInfo.attemptId) {
originalAppShuffleInfo.set(appShuffleInfo);
+ AppPathsInfo appPathsInfo = new AppPathsInfo(appId,
executorInfo.localDirs,
+ mergeDir, executorInfo.subDirsPerLocalDir);
+ // Clean up the outdated App Attempt local path info in the DB
and
+ // put the newly registered local path info from newer attempt
into the DB.
+ // Deletion or insertion may fail as of various reasons.
Review Comment:
Removed.
##########
common/network-shuffle/src/main/java/org/apache/spark/network/shuffle/RemoteBlockPushResolver.java:
##########
@@ -656,6 +781,236 @@ public void registerExecutor(String appId,
ExecutorShuffleInfo executorInfo) {
}
}
+ /**
+ * Close the DB during shutdown
+ */
+ @Override
+ public void close() {
+ if (db != null) {
+ try {
+ db.close();
+ } catch (IOException e) {
+ logger.error("Exception closing leveldb with registered app paths info
and "
+ + "shuffle partition info", e);
+ }
+ }
+ }
+
+ /**
+ * Write the application attempt's local path information to the DB
+ */
+ private void writeAppPathsInfoToDb(String appId, int attemptId, AppPathsInfo
appPathsInfo) {
+ if (db != null) {
+ AppAttemptId appAttemptId = new AppAttemptId(appId, attemptId);
+ try {
+ byte[] key = getDbAppAttemptPathsKey(appAttemptId);
+ String valueStr = mapper.writeValueAsString(appPathsInfo);
+ byte[] value = valueStr.getBytes(StandardCharsets.UTF_8);
+ db.put(key, value);
+ } catch (Exception e) {
+ logger.error("Error saving registered app paths info for {}",
appAttemptId, e);
+ }
+ }
+ }
+
+ /**
+ * Write the finalized shuffle merge partition information into the DB
+ */
+ private void writeAppAttemptShuffleMergeInfoToDB(
+ AppAttemptShuffleMergeId appAttemptShuffleMergeId) {
+ if (db != null) {
+ // Write AppAttemptShuffleMergeId into LevelDB for finalized shuffles
+ try{
+ byte[] dbKey =
getDbAppAttemptShufflePartitionKey(appAttemptShuffleMergeId);
+ db.put(dbKey, new byte[0]);
+ } catch (Exception e) {
+ logger.error("Error saving active app shuffle partition {}",
appAttemptShuffleMergeId, e);
+ }
+ }
+ }
+
+ /**
+ * Parse the DB key with the prefix and the expected return value type
+ */
+ private <T> T parseDbKey(String key, String prefix, Class<T> valueType)
throws IOException {
+ String json = key.substring(prefix.length() + 1);
+ return mapper.readValue(json, valueType);
+ }
+
+ /**
+ * Generate AppAttemptId from the DB key
+ */
+ private AppAttemptId parseDbAppAttemptPathsKey(String key) throws
IOException {
+ return parseDbKey(key, APP_ATTEMPT_PATH_KEY_PREFIX, AppAttemptId.class);
+ }
+
+ /**
+ * Generate AppAttemptShuffleMergeId from the DB key
+ */
+ private AppAttemptShuffleMergeId parseDbAppAttemptShufflePartitionKey(
+ String key) throws IOException {
+ return parseDbKey(
+ key, APP_ATTEMPT_SHUFFLE_FINALIZE_STATUS_KEY_PREFIX,
AppAttemptShuffleMergeId.class);
+ }
+
+ /**
+ * Generate the DB key with the key object and the specified string prefix
+ */
+ private byte[] getDbKey(Object key, String prefix) throws IOException {
+ // We add a common prefix on all the keys so we can find them in the DB
+ String keyJsonString = prefix + DB_KEY_DELIMITER +
mapper.writeValueAsString(key);
+ return keyJsonString.getBytes(StandardCharsets.UTF_8);
+ }
+
+ /**
+ * Generate the DB key from AppAttemptShuffleMergeId object
+ */
+ private byte[] getDbAppAttemptShufflePartitionKey(
+ AppAttemptShuffleMergeId appAttemptShuffleMergeId) throws IOException {
+ return getDbKey(appAttemptShuffleMergeId,
APP_ATTEMPT_SHUFFLE_FINALIZE_STATUS_KEY_PREFIX);
+ }
+
+ /**
+ * Generate the DB key from AppAttemptId object
+ */
+ private byte[] getDbAppAttemptPathsKey(AppAttemptId appAttemptId) throws
IOException {
+ return getDbKey(appAttemptId, APP_ATTEMPT_PATH_KEY_PREFIX);
+ }
+
+ /**
+ * Reload the DB to recover the meta data stored in the hashmap for merged
shuffles.
+ * The application attempts local paths information will be firstly
reloaded, and then
+ * the finalized shuffle merges will be updated.
+ * This method will also try deleting dangling key/values in DB, which
includes:
+ * 1) Outdated application attempt local paths information as of some DB
deletion failures
+ * 2) The deletion of finalized shuffle merges are triggered asynchronously,
there can be cases
+ * that deletions miss the execution during restart. These finalized shuffle
merges should have
+ * no relevant application attempts local paths information registered in
the DB and the hashmap.
+ */
+ @VisibleForTesting
+ void reloadAndCleanUpAppShuffleInfo(DB db) throws IOException {
+ logger.info("Reload applications merged shuffle information from DB");
+ List<byte[]> dbKeysToBeRemoved = new ArrayList<>();
+ dbKeysToBeRemoved.addAll(reloadActiveAppAttemptsPathInfo(db));
+ dbKeysToBeRemoved.addAll(reloadFinalizedAppAttemptsShuffleMergeInfo(db));
+ // Clean up invalid data stored in DB
+ submitCleanupTask(() ->
+ dbKeysToBeRemoved.forEach(
+ (key) -> {
+ try {
+ if (logger.isDebugEnabled()) {
+ logger.debug("Removing dangling key {} in DB",
+ parseDbAppAttemptShufflePartitionKey(
Review Comment:
Removed this debugging logging since it needs to deal with 2 types of DB
keys.
This is done async, it is calling the submitCleanupTask
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
