Streams and Temp File Cleanup: Fixing a Real Production Issue

Mismanaging resources is one of the easiest ways to bring down a production system. In a development environment, it’s easy to ignore open streams and temporary files, but at production scale, they can have a hugely negative impact. Unfortunately, it’s far too easy to overlook these issues while developing, even when you’re looking out for them.

Dealing with multiple streams

A few months ago, I dealt with a bug in one of our production services that centered around Scala code that looked something like this:

def getImageBytes(image: Image): Array[Byte] = {
  val byteArrayOutputStream = new ByteArrayOutputStream()
  val writer: ImageWriter = openWriter()

  val output: ImageOutputStream = createImageOutputStream(byteArrayOutputStream)
  writer.setOutput(output)
  writer.write(image)

  byteArrayOutputStream.flush()
  val bytes = byteArrayOutputStream.toByteArray
  byteArrayOutputStream.close()
  writer.dispose()
  bytes
}

At first glance, this doesn’t look so bad. A stream is being closed, and the writer is disposed of. However, there are two big issues which need to be addressed to provide better production stability. The most obvious problem is that this code won’t close the stream or dispose of the writer if any exceptions are thrown. However, even in the normal case, a temp image file created in this code snippet isn't properly deleted when the code is done running. As this repeatedly occurs, the disk could easily fill up and no new images would be processed.

After a bit of investigation, I discovered that the call to createImageOutputStream was creating a temp file, but output was never getting closed, so the file was never cleaned up. byteArrayOutputStream’s close and writer’s dispose calls made it seem like all the streams and writers were being properly handled, but in code that uses multiple resources, it’s easy to miss closing one. In this case, closing byteArrayOutputStream and not output is a particularly unfortunate event: closing a ByteArrayOutputStream is a no-op, but closing output actually cleans up temp files. Simply adding the line output.close(), fixes the main bug, but that’s not really enough to provide production stability.

Error cases

In a production system, it’s inevitable that code will throw exceptions. While programmers can do many things to mitigate issues, from using Option instead of null to using Either instead of throwing an exception, error states resulting from calls to external code are essentially unavoidable. In the above example, if the writer.write call throws an exception for some reason, neither stream would get closed, and the writer wouldn’t be disposed. If errors frequently occur, this small issue can quickly cascade into more serious problems like running out of file handles or disk space. To make sure that everything is properly cleaned up, it’s important to always add a try-finally block. The cleanup code needs to be in the finally block, not try or catch, because the resources need to be properly disposed of regardless of whether the code succeeds or fails. So in a simple case, the code might look like this:

def operationWithStream(): Unit = {
  val streamVal = openStream()
  var secondStreamVar = Option.empty[Stream]
  try {
    secondStreamVar = Some(openStream(streamVar))
    secondStreamVar.read()
  } finally {
    streamVal.close()
    secondStreamVar.foreach(_.close())
  }
}

As a Scala developer, however, the use of var isn’t idiomatic, and it could be a source of errors. You could avoid this issue by nesting the try blocks, but that quickly becomes hard to read as more resources are added.

Ensuring streams and temp files are cleaned up

Even when a good-faith attempt is made to be careful with resources, human error often leaves places for file handles and temp files to leak. The smaller a burden you can place on the developer’s good practice, the better. When I originally fixed this issue, I used our own closeable function which is basically syntactic sugar around the nested try approach.

//Our closeable function
def closeable[A, B](create: => A)(run: A => B)(close: A => Unit): B = {
  val closeable = create
  try {
    run(closeable)
  } finally {
    close(closeable)
  }
}

//Refactored to use closeable
def getImageBytes(image: Image): Array[Byte] = {
  closeable(new ByteArrayOutputStream()) { byteArrayOutputStream =>
    closeable(openWriter()) { writer =>
      closeable(createImageOutputStream(byteArrayOutputStream)) { output =>
        writer.setOutput(output)
        writer.write(image)

        byteArrayOutputStream.flush()
        byteArrayOutputStream.toByteArray
      }(_.close())
    }(_.dispose())
  }(_.close())
}

For a more robust solution with less boilerplate, the Scala ARM library is a better choice. It allows you to wrap a resource with managed, creating a ManagedResource which will close or dispose of the resource as soon as you’re finished with it.

Here is a simple example fixed using Scala ARM:

def operationWithStream(): Either[Seq[Throwable], Array[Bytes] = {
  val managedResult = for {
    firstStream <- managed(openStream())
    secondStream <- managed(openStream(firstStream))
  } {
    //For this simple example, read returns an array of bytes
    secondStream.read()
  }
  managedResult.map(identity).either
}

Using Scala ARM to properly handle the image byte streams would produce something like this:

def getImageBytes(image: Image): Either[Seq[Throwable], Array[Byte]] = {
  val managedBytes = for {
    byteArrayOutputStream <- managed(new ByteArrayOutputStream())
    writer <- managed(openWriter())
    output <- managed(createImageOutputStream(byteArrayOutputStream))
  } yield {
    writer.setOutput(output)
    writer.write(image)

    byteArrayOutputStream.flush()
    byteArrayOutputStream.toByteArray
  }
  managedBytes.map(identity).either
}

As shown above, ARM allows you to use streams and other managed resources in a more functional style while making certain that every resource is cleaned up for you. Soon, Scala will offer a very similar approach with Using, coming in Scala 12.3.

With both approaches, you’re able to treat the open stream as part of a scope—a structure that guarantees the stream will be closed at the correct time. Other languages offer similar features, such as Python’s with statement. Servers have finite resources, so it’s important not to give them up because of a small mistake. Because these issues aren’t immediately apparent on a dev machine due to the difference in request volume, it’s important to have compile-time checks in place to make resource management as easy as possible. Handling resources needs to be an integral part of your workflow to prevent costly mistakes. If you clean up streams correctly as you go, it’s much easier to maintain a stable production environment.