Reject truncated Zstandard event logs

yarn-resource-cost/test_calculate_yarn_job_cost.py

@@ -1273,6 +1273,21 @@ def test_portable_eventlog_reader_handles_zstd(self):
             ),
         )
 
+    def test_portable_eventlog_reader_handles_concatenated_zstandard(self):
+        compressor = zstandard.ZstdCompressor()
+        compressed = compressor.compress(b"first\n") + compressor.compress(
+            b"second\n"
+        )
+
+        self.assertEqual(
+            ["first", "second"],
+            list(
+                EVENTLOG.iter_text_lines(
+                    io.BytesIO(compressed), EVENTLOG.ZSTD_CODEC
+                )
+            ),
+        )
+
     def test_portable_eventlog_reader_preserves_split_utf8(self):
         payload = '{"name":"caf\u00e9"}\n'.encode()
         split = payload.index("\u00e9".encode()) + 1
@@ -1310,6 +1325,37 @@ def test_portable_eventlog_reader_classifies_corrupt_zstandard_as_retryable(self
             )
         self.assertTrue(raised.exception.retryable)
 
+    def test_portable_eventlog_reader_classifies_truncated_zstandard_as_retryable(
+        self,
+    ):
+        compressed = zstandard.ZstdCompressor().compress(b"one\ntwo\n")
+
+        with self.assertRaisesRegex(ValueError, "Truncated Zstandard") as raised:
+            list(
+                EVENTLOG.iter_text_lines(
+                    io.BytesIO(compressed[:-1]), EVENTLOG.ZSTD_CODEC
+                )
+            )
+
+        self.assertTrue(raised.exception.retryable)
+
+    def test_portable_eventlog_reader_rejects_truncated_concatenated_zstandard(
+        self,
+    ):
+        compressor = zstandard.ZstdCompressor()
+        compressed = compressor.compress(b"one\n") + compressor.compress(b"two\n")[
+            :-1
+        ]
+        lines = EVENTLOG.iter_text_lines(
+            io.BytesIO(compressed), EVENTLOG.ZSTD_CODEC
+        )
+
+        self.assertEqual("one", next(lines))
+        with self.assertRaisesRegex(ValueError, "Truncated Zstandard") as raised:
+            list(lines)
+
+        self.assertTrue(raised.exception.retryable)
+
     def test_portable_eventlog_reader_handles_tar_bundle(self):
         with tempfile.TemporaryDirectory() as directory:
             root = Path(directory)

yarn-resource-cost/test_yarn_job_cost_api.py

@@ -240,6 +240,40 @@ def test_corrupt_zstd_event_log_is_retryable(self):
         self.assertTrue(result.retryable)
         self.assertIn("Unable to decompress", result.warnings[0])
 
+    def test_truncated_concatenated_zstd_event_log_is_retryable(self):
+        event_prefix = "spark-events/eventlog_v2_application_1_0001"
+        events = "\n".join(
+            (
+                '{"Event":"SparkListenerApplicationStart","App Name":'
+                '"portable-cost-sample","App ID":"application_1_0001",'
+                '"Timestamp":1000}',
+                '{"Event":"SparkListenerApplicationEnd","Timestamp":11000}',
+            )
+        ).encode()
+        compressor = zstandard.ZstdCompressor()
+        compressed = compressor.compress(events) + compressor.compress(b"later\n")[
+            :-1
+        ]
+        s3 = FakeS3Client(
+            {
+                f"{event_prefix}/events_1_application_1_0001.zstd": compressed,
+                "emr-logs/j-TEST/node/i-1/applications/"
+                "hadoop-yarn-resourcemanager-rm.log": (
+                    self.yarn_log_with_instance_type().encode()
+                ),
+            }
+        )
+
+        result = calculate_emr_application_usage(
+            self.request(f"s3://test-bucket/{event_prefix}"),
+            emr_client=FakeEmrClient(),
+            s3_client=s3,
+        )
+
+        self.assertFalse(result.complete)
+        self.assertTrue(result.retryable)
+        self.assertIn("Truncated Zstandard frame", " | ".join(result.warnings))
+
     def test_permanent_event_log_read_error_is_not_retryable(self):
         event_log_dir = FIXTURE / "eventlog_v2_application_1_0001"
         metadata = SimpleNamespace(

yarn-resource-cost/yarn_job_cost_eventlog.py

@@ -117,11 +117,28 @@ def iter_zstd_chunks(stream: BinaryIO) -> Iterable[bytes]:
             retryable=False,
         ) from error
     try:
-        decompressor = zstandard.ZstdDecompressor()
-        with decompressor.stream_reader(
-            stream, read_across_frames=True, closefd=False
-        ) as reader:
-            yield from iter(lambda: reader.read(READ_SIZE), b"")
+        compressed = b""
+        while True:
+            frame = zstandard.ZstdDecompressor().decompressobj(
+                write_size=READ_SIZE
+            )
+            while not frame.eof:
+                if not compressed:
+                    compressed = stream.read(READ_SIZE)
+                    if not compressed:
+                        raise EventLogReadError(
+                            "Truncated Zstandard frame", retryable=True
+                        )
+                output = frame.decompress(compressed)
+                compressed = frame.unconsumed_tail
+                for offset in range(0, len(output), READ_SIZE):
+                    yield output[offset : offset + READ_SIZE]
+
+            compressed = frame.unused_data
+            if not compressed:
+                compressed = stream.read(READ_SIZE)
+                if not compressed:
+                    return
     except zstandard.ZstdError as error:
         raise EventLogReadError(
             f"Unable to decompress Zstandard event log: {error}",