Automatic Resource Management in Scala: Revisiting Test Context Classes

I have written previously about how I like to use what I call call context classes for functional tests in order to encapsulate test setup, reliably cleanup resources at the end of a test, and to promote code reuse across tests. I started by writing about using disposable objects for writing functional tests in C# and I followed this by detailing how I use the RAII pattern in C++ for the same purpose. Now that I'm writing more tests in Scala, I would like to revisit the concept of resource context classes one more time and use it as an opportunity to explore automatic resource management in Scala.

For the examples in this blog post, I will consider writing tests for a WebSocket server where I want to create a unique WebSocket connection at the beginning of each test, then reliably close the WebSocket connection at the end of the test, independent of the test result, so that the server will release the resources associated with the connection. The WebSocket server has the endpoint http://<hostname>:8080/<UUID>, where the UUID uniquely identifies a specific client. I will use the Java-WebSocket client implemented in the following Scala class.

import java.net.URI
import org.java_websocket.drafts.Draft_17
import org.java_websocket.handshake.ServerHandshake

class WebSocketClient(url: String, callback: String => Unit)
  extends org.java_websocket.client.WebSocketClient(new URI(url), new Draft_17()) {

  override def onMessage(message: String): Unit = callback(message)

  override def onError(ex: Exception): Unit = throw ex

  override def onClose(code: Int, reason: String, remote: Boolean): Unit =
    println(s"WebSocket closed for url : $url")

  override def onOpen(handshakedata: ServerHandshake): Unit =
    println(s"WebSocket opened for url : $url")
}

object WebSocketClient {
  def apply(url: String, callback: String => Unit): WebSocketClient = {
    new WebSocketClient(url, callback)
  }
}

Automatic Resource Management in Scala

In C#, the scope of a resource managed by an object that implements the IDisposable interface can be controlled with a using statement. The using statement is equivalent to creating the resource, performing a set of operations in a try block, then calling the Dispose method in a finally block in order to cleanup the resource. In C++, with deterministic destruction and RAII, the class constructor creates the resource and, symmetrically, the class destructor cleans up the resource. The object encapsulating the resource can simply be created on the stack and the resource will automatically be cleaned up when it goes out of scope.

Classes in Scala do not have destructors and there is no equivalent of the C# using statement native to the language. It is not hard, however, to use Scala's parameter groups and function literals to create an object that encapsulates resource acquisition and deterministically performs the cleanup in a finally statement. This is known as the Loan Pattern. Consider the following UsingWebSocket object that creates a WebSocket connection to the server, calls the function literal defined by the caller — passing to it the WebSocket client as an argument — and, finally, closes the WebSocket connection.

object UsingWebSocket {
  def apply[T](url: String, callback: String => Unit)(f: WebSocketClient => T): T = {
    val client = WebSocketClient(url, callback)
    assert(client.connectBlocking(), s"Failed to connect to $url")
    try {
      f(client)
    } finally {
      client.close()
    }
  }
}

The UsingWebSocket object can then be used in a test by providing the URL and callback method, defining the test code that makes use of the WebSocket connection in a function literal. For example, in the following test, which I will use throughout this article, the test passes if the client sends a message to the server. If it fails to send a message to the server, an exception will be thrown and the test fails. In this simple test, I do not make use of the callback method and simply define a empty function.

class WebSocketTests extends FlatSpec with ShouldMatchers {
  val url = s"http://localhost:8080/${uuid()}"
  def callback(response: String) {}

  "Client" should "send message to server" in {
    UsingWebSocket(url, callback) { client =>
      client.send("Bueller?")
    }
  }
}

If the UsingWebSocket object is going to be used across a large number of tests, the test code can be streamlined further by using implicit parameters to define the URL and callback function. Note that implicit parameters must be defined as the last parameter group to a method.

object UsingWebSocket {
  def apply[T](f: WebSocketClient => T)
              (implicit url: String, callback: String => Unit): T = {
    val client = WebSocketClient(url, callback)
    assert(client.connectBlocking(), s"Failed to connect to $url")
    try {
      f(client)
    } finally {
      client.close()
    }
  }
}

With this redefined UsingWebSocket object, the URL and callback can be defined once for a group of tests and implicitly passed to the object each time its used. I also adopt the convention of adding an explicit return type for the callback method now that it is an implicit method.

class WebSocketTests extends FlatSpec with ShouldMatchers {
  implicit val url = s"http://localhost:8080/${uuid()}"
  implicit def callback(response: String): Unit = {}

  "Client" should "send message to server" in {
    UsingWebSocket { client =>
      client.send("Bueller?")
    }
  }
}

If you want to define consistent resource management patterns for objects across your code base, it is possible to implement a using pattern like in C# by defining a Disposable trait and a using object that operates on objects that mix-in this trait.

trait Disposable {
  def dispose()
}

object using {
  def apply[D, T](resource: D)(f: D => T)(implicit d: D => Disposable): T = {
    try {
      f(resource)
    }
    finally {
      resource.dispose()
    }
  }
}

I can define a new WebSocket class that extends WebSocketClient, mixes in the Disposable trait, and overrides the dispose method to call the close method of the WebSocket client.

class WebSocket(url: String, callback: String => Unit) extends WebSocketClient(url, callback) with Disposable {
  assert(super.connectBlocking(), s"Failed to connect to $url")
  override def dispose() = super.close()
}

It can be used in the test with the using object to control the scope of the WebSocket connection.

class WebSocketTests extends FlatSpec with ShouldMatchers {
  val url = s"http://localhost:8080/${uuid()}"
  def callback(response: String) {}

  "Client" should "send message to server" in {
    using(new WebSocket(url, callback)) { client =>
      client.send("Bueller?")
    }
  }
}

A final option for automatic resource management is to adopt a library like Scala ARM. This library is quite sophisticated and it can be used in many different ways, the most straightforward of which is the imperative style using a for comprehension. A resource that supports either a dispose or close method is wrapped in the managed method and will be automatically cleaned up at the end of the for comprehension.

class WebSocketTests extends FlatSpec with ShouldMatchers {
  val url = s"http://localhost:8080/${uuid()}"
  def callback(response: String) {}

  "Client" should "send message to server" in {
    for (client <- managed(WebSocketClient(url, callback))) {
      assert(client.connectBlocking(), s"Failed to connect to $url")
      client.send("Bueller?")
    }
  }
}

These examples show the diversity of options available for automatic resource management in Scala and highlight the flexibility of the language. The Scala ARM library even supports additional patterns which I did not explore here. This flexibility is powerful and interesting, however, I do think it makes Scala somewhat intimidating. Like a number of things in Scala, because there are so many ways to do the same thing, it makes it hard to know how to write idiomatic Scala, if such a thing even exists for automatic resource management. The Scala ARM project aims to be an incubator project for automatic resource management in Scala and I think standardization, through Scala ARM or another library, is welcome.

I do find it strange, for such a modern language, that there is no native resource management in Scala, either directly, through destructors like C++, or by convention, like the using pattern in C#. I understand the reduced focus on class destructors in managed languages like C#, Java, and Scala, where the lifetime of an object is ultimately determined by the garbage collector, but I think the simplicity of the class destructor for automatic resource management is often under appreciated. Certainly any language that supports a finally block can implement reliable resource management, but the code can become complex when managing multiple resources. This is something the using statement in C# helps address, and the Scala ARM library attempts to address it as well, but I agree with Bjarne Stroustrup's opinion regarding the finally construct:

In a system, we need a "resource handle" class for each resource. However, we don't have to have an "finally" clause for each acquisition of a resource. In realistic systems, there are far more resource acquisitions than kinds of resources, so the "resource acquisition is initialization" technique leads to less code than use of a "finally" construct.

As I write functional tests that will benefit from automatic resource management, I will most likely use the Scala ARM library or I will create classes dedicated to acquiring and releasing specific resources, like the UsingWebSocket class I provided in the first example.

In my next article, I will explore extending the test context classes that I developed in this article to make writing functional tests even easier and more reliable.