Handle buffered dwrf write exception to avoid server crash

Catch exception during BufferedWriter processing to avoid server crash as ~BufferedWriter expect either the object is aborted or all the data mutations have been flushed.
facebookincubator · May 13, 2024 · 8c2ca34 · 8c2ca34
1 parent 49c3ebb
commit 8c2ca34
Show file tree

Hide file tree

Showing 2 changed files with 5 additions and 0 deletions.
diff --git a/velox/dwio/dwrf/writer/DictionaryEncodingUtils.h b/velox/dwio/dwrf/writer/DictionaryEncodingUtils.h
@@ -80,6 +80,8 @@ class DictionaryEncodingUtils {
     uint32_t newIndex = 0;
     auto dictLengthWriter =
         createBufferedWriter<uint32_t>(pool, 64 * 1024, lengthFn);
+    auto errorGuard =
+        folly::makeGuard([&dictLengthWriter]() { dictLengthWriter.abort(); });
     for (uint32_t i = 0; i != numKeys; ++i) {
       auto origIndex = (sort ? sortedIndex[i] : i);
       if (!dropInfrequentKeys || shouldWriteKey(dictEncoder, origIndex)) {
@@ -94,6 +96,7 @@ class DictionaryEncodingUtils {
         inDict[origIndex] = false;
       }
     }
+    errorGuard.dismiss();:
     dictLengthWriter.close();
 
     return newIndex;

diff --git a/velox/dwio/dwrf/writer/IntegerDictionaryEncoder.h b/velox/dwio/dwrf/writer/IntegerDictionaryEncoder.h
@@ -207,6 +207,7 @@ class IntegerDictionaryEncoder : public AbstractIntegerDictionaryEncoder {
     uint32_t newIndex = 0;
 
     auto dictWriter = createBufferedWriter<Integer>(writeBuffer, fn);
+    auto errorGuard = folly::makeGuard([&dictWriter]() { dictWriter.abort(); });
     for (uint32_t i = 0; i != numKeys; ++i) {
       auto origIndex = (sort ? sortedIndex[i] : i);
       if (!dropInfrequentKeys || shouldWriteKey(dictEncoder, origIndex)) {
@@ -218,6 +219,7 @@ class IntegerDictionaryEncoder : public AbstractIntegerDictionaryEncoder {
         inDict[origIndex] = false;
       }
     }
+    errorGuard.dismiss();
     dictWriter.close();
     return newIndex;
   }