[SPARK-47910][CORE] close stream when DiskBlockObjectWriter closeResources to avoid memory leak #46131

JacobZheng0927 · 2024-04-19T06:46:14Z

What changes were proposed in this pull request?

close stream when DiskBlockObjectWriter closeResources to avoid memory leak

Why are the changes needed?

SPARK-34647 replaced the ZstdInputStream with ZstdInputStreamNoFinalizer. This meant that all usages of CompressionCodec.compressedOutputStream would need to manually close the stream as this would no longer be handled by the finalizer mechanism.
When using zstd for shuffle write compression, if for some reason the execution of this process is interrupted(eg. enable spark.sql.execution.interruptOnCancel and cancel Job). The memory used by ZstdInputStreamNoFinalizer may not be freed, causing a memory leak.

Does this PR introduce any user-facing change?

No

How was this patch tested?

Spark Shell Configuration

$> export SPARK_SUBMIT_OPTS="-XX:+AlwaysPreTouch -Xms1g"
$> $SPARK_HOME/bin/spark-shell --conf spark.io.compression.codec=zstd

Test Script

import java.util.concurrent.TimeUnit
import scala.concurrent.Future
import scala.concurrent.ExecutionContext.Implicits.global
import scala.util.Random

sc.setJobGroup("jobA", "this is a job to be cancelled", interruptOnCancel = true)
(1 to 50).foreach { batch => {
  val jobA = Future {
    val df1 = spark.range(2000000).map { _ =>(Random.nextString(20),Random.nextInt(1000),Random.nextInt(1000),Random.nextInt(10))}.toDF("a","b","c","d")
    val df2 = spark.range(2000000).map { _ =>(Random.nextString(20),Random.nextInt(1000),Random.nextInt(1000),Random.nextInt(10))}.toDF("a","b","c","d")
    df1.join(df2,"b").show()  }
  Thread.sleep(5000)
  sc.cancelJobGroup("jobA")
}}

Memory Monitor

$> while true; do echo \"$(date +%Y-%m-%d' '%H:%M:%S)\",$(pmap -x <PID> | grep "total kB" | awk '{print $4}'); sleep 10; done;

Results

Before

"2024-05-13 16:54:23",1332384
"2024-05-13 16:54:33",1417112
"2024-05-13 16:54:43",2211684
"2024-05-13 16:54:53",3060820
"2024-05-13 16:55:03",3850444
"2024-05-13 16:55:14",4631744
"2024-05-13 16:55:24",5317200
"2024-05-13 16:55:34",6019464
"2024-05-13 16:55:44",6489180
"2024-05-13 16:55:54",7255548
"2024-05-13 16:56:05",7718728
"2024-05-13 16:56:15",8388392
"2024-05-13 16:56:25",8927636
"2024-05-13 16:56:36",9473412
"2024-05-13 16:56:46",10000380
"2024-05-13 16:56:56",10344024
"2024-05-13 16:57:07",10734204
"2024-05-13 16:57:17",11211900
"2024-05-13 16:57:27",11665524
"2024-05-13 16:57:38",12268976
"2024-05-13 16:57:48",12896264
"2024-05-13 16:57:58",13572244
"2024-05-13 16:58:09",14252416
"2024-05-13 16:58:19",14915560
"2024-05-13 16:58:30",15484196
"2024-05-13 16:58:40",16170324

After

"2024-05-13 16:35:44",1355428
"2024-05-13 16:35:54",1391028
"2024-05-13 16:36:04",1673720
"2024-05-13 16:36:14",2103716
"2024-05-13 16:36:24",2129876
"2024-05-13 16:36:35",2166412
"2024-05-13 16:36:45",2177672
"2024-05-13 16:36:55",2188340
"2024-05-13 16:37:05",2190688
"2024-05-13 16:37:15",2195168
"2024-05-13 16:37:26",2199296
"2024-05-13 16:37:36",2228052
"2024-05-13 16:37:46",2238104
"2024-05-13 16:37:56",2260624
"2024-05-13 16:38:06",2307184
"2024-05-13 16:38:16",2331140
"2024-05-13 16:38:27",2323388
"2024-05-13 16:38:37",2357552
"2024-05-13 16:38:47",2352948
"2024-05-13 16:38:57",2364744
"2024-05-13 16:39:07",2368528
"2024-05-13 16:39:18",2385492
"2024-05-13 16:39:28",2389184
"2024-05-13 16:39:38",2388060
"2024-05-13 16:39:48",2388336
"2024-05-13 16:39:58",2386916

Was this patch authored or co-authored using generative AI tooling?

No

JacobZheng0927 · 2024-04-22T02:08:16Z

cc @dongjoon-hyun
This is a similar to #35613. Please take a look, thanks!

dongjoon-hyun

Thank you for making a PR, @JacobZheng0927 .

However, your PR fails to compile. Please make GitHub Action CI green.

[error] (core / Compile / compileIncremental) Compilation failed

JacobZheng0927 · 2024-04-22T05:46:33Z

Thank you for making a PR, @JacobZheng0927 .

However, your PR fails to compile. Please make GitHub Action CI green.

[error] (core / Compile / compileIncremental) Compilation failed

Done.

dongjoon-hyun · 2024-04-22T22:12:38Z

core/src/main/scala/org/apache/spark/storage/DiskBlockObjectWriter.scala

nit. Redundant empty line. Please remove this.

dongjoon-hyun

Like #35613, do you think you can provide a way to validate your PR, @JacobZheng0927 ?

mridulm · 2024-04-23T06:20:58Z

core/src/main/scala/org/apache/spark/storage/DiskBlockObjectWriter.scala

+    } catch {
+      case e: IOException =>
+        logError(log"Exception occurred while closing the output stream" +
+          log"${MDC(ERROR, e.getMessage)}")


error -> info
and the exception stack trace would be useful

mridulm · 2024-04-23T06:24:14Z

core/src/main/scala/org/apache/spark/storage/DiskBlockObjectWriter.scala

+          objOut = closeIfNonNull(objOut)
+          bs = null
+        } {
+          bs = closeIfNonNull(bs)


Suggested change

objOut = closeIfNonNull(objOut)

bs = null

} {

bs = closeIfNonNull(bs)

if (null != objOut) objOut.close()

bs = null

} {

objOut = null

if (null != bs) bs.close()

bs = null

And remove closeIfNonNull

ok. I'll fix it.

JacobZheng0927 · 2024-04-23T11:12:06Z

Like #35613, do you think you can provide a way to validate your PR, @JacobZheng0927 ?

Ok, I'll try to reproduce the problem with a simple script.

JacobZheng0927 · 2024-05-13T09:07:22Z

I apologize for the long delay in updating. I've just added the steps for reproduction, please take a look. @dongjoon-hyun @mridulm

mridulm · 2024-05-13T18:24:19Z

@JacobZheng0927, please add it as a unit test.

JacobZheng0927 · 2024-05-15T07:22:40Z

@JacobZheng0927, please add it as a unit test.

I'm not sure how to test for native memory leak cases in unit tests. is there a relevant example I can refer to?

mridulm · 2024-05-23T06:47:57Z

One way I can quickly think of is to check if objOut.close() or bs.close() is being called or not.
For example, adapt the "calling closeAndDelete() on a partial write file" test and use either a custom serializer or a custom compression codec to check for close being invoked ?

(Sorry for the delay in getting back to you - this PR dropped off my todo list unfortunately)

github-actions bot added the CORE label Apr 19, 2024

JacobZheng0927 force-pushed the zstdMemoryLeak branch from 44a5604 to 6bb8917 Compare April 19, 2024 07:20

HyukjinKwon changed the title ~~SPARK-47910; close stream when DiskBlockObjectWriter closeResources to avoid memory leak~~ [SPARK-47910][CORE] close stream when DiskBlockObjectWriter closeResources to avoid memory leak Apr 19, 2024

dongjoon-hyun requested changes Apr 22, 2024

View reviewed changes

JacobZheng0927 force-pushed the zstdMemoryLeak branch from 6bb8917 to 8088e75 Compare April 22, 2024 03:18

JacobZheng0927 requested a review from dongjoon-hyun April 22, 2024 05:44

dongjoon-hyun reviewed Apr 22, 2024

View reviewed changes

mridulm reviewed Apr 23, 2024

View reviewed changes

JacobZheng0927 added 2 commits May 11, 2024 11:01

SPARK-47910; close stream when DiskBlockObjectWriter closeResources

877acaf

SPARK-47910; Avoiding IOException thrown by closeResources

c90e9f4

github-actions bot added SQL ML MLLIB STRUCTURED STREAMING KUBERNETES WEB UI GRAPHX BUILD YARN EXAMPLES DOCS INFRA PYTHON R DSTREAM AVRO labels May 11, 2024

github-actions bot added CONNECT PROTOBUF labels May 11, 2024

SPARK-47910; remove closeIfNonNull

2db509c

JacobZheng0927 force-pushed the zstdMemoryLeak branch from e7160f8 to 2db509c Compare May 11, 2024 03:50

github-actions bot added SQL and removed SQL ML MLLIB STRUCTURED STREAMING KUBERNETES WEB UI GRAPHX BUILD YARN EXAMPLES DOCS INFRA PYTHON R DSTREAM AVRO PANDAS API ON SPARK CONNECT PROTOBUF labels May 11, 2024

JacobZheng0927 force-pushed the zstdMemoryLeak branch from 859a7f3 to 2db509c Compare May 13, 2024 08:01

github-actions bot removed the SQL label May 13, 2024

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

[SPARK-47910][CORE] close stream when DiskBlockObjectWriter closeResources to avoid memory leak #46131

[SPARK-47910][CORE] close stream when DiskBlockObjectWriter closeResources to avoid memory leak #46131

JacobZheng0927 commented Apr 19, 2024 •

edited

JacobZheng0927 commented Apr 22, 2024

dongjoon-hyun left a comment

JacobZheng0927 commented Apr 22, 2024

dongjoon-hyun Apr 22, 2024

dongjoon-hyun left a comment

mridulm Apr 23, 2024

mridulm Apr 23, 2024 •

edited

JacobZheng0927 Apr 23, 2024

JacobZheng0927 commented Apr 23, 2024

JacobZheng0927 commented May 13, 2024

mridulm commented May 13, 2024

JacobZheng0927 commented May 15, 2024

mridulm commented May 23, 2024

[SPARK-47910][CORE] close stream when DiskBlockObjectWriter closeResources to avoid memory leak #46131

Are you sure you want to change the base?

[SPARK-47910][CORE] close stream when DiskBlockObjectWriter closeResources to avoid memory leak #46131

Conversation

JacobZheng0927 commented Apr 19, 2024 • edited

What changes were proposed in this pull request?

Why are the changes needed?

Does this PR introduce any user-facing change?

How was this patch tested?

Spark Shell Configuration

Test Script

Memory Monitor

Results

Before

After

Was this patch authored or co-authored using generative AI tooling?

JacobZheng0927 commented Apr 22, 2024

dongjoon-hyun left a comment

Choose a reason for hiding this comment

JacobZheng0927 commented Apr 22, 2024

dongjoon-hyun Apr 22, 2024

Choose a reason for hiding this comment

dongjoon-hyun left a comment

Choose a reason for hiding this comment

mridulm Apr 23, 2024

Choose a reason for hiding this comment

mridulm Apr 23, 2024 • edited

Choose a reason for hiding this comment

JacobZheng0927 Apr 23, 2024

Choose a reason for hiding this comment

JacobZheng0927 commented Apr 23, 2024

JacobZheng0927 commented May 13, 2024

mridulm commented May 13, 2024

JacobZheng0927 commented May 15, 2024

mridulm commented May 23, 2024

JacobZheng0927 commented Apr 19, 2024 •

edited

mridulm Apr 23, 2024 •

edited