GRPC Cleanup Extension for JUnit 5

The grpc-java project provides GrpcCleanupRule which is a helpful way to clean up Grpc Service resources in JUnit 4 tests, but it doesn’t work with JUnit 5. JUnit 5 doesn’t support@Rule classes, but rather has moved that functionality into Extensions.

Additionally, the GRPC folks have said they aren’t moving to JUnit 5 for a while, and they suggest sticking with JUnit 4.Here at StackHawk, we want to use JUnit 5 exclusively, and we want the test resources to get automatically cleaned up to avoid test pollution. So we wrote a JUnit 5 Extension called Grpc Cleanup Extension that handles the channel/server resource cleanup after each test invocation.

The Problem

If you run a JUnit 5 test with theGrpc cleanup rule, you’ll see something this in the logs:

2021-07-08 09:33:21,962 ERROR [Test worker] io.grpc.internal.ManagedChannelOrphanWrapper$ManagedChannelReference: *~*~*~ Channel ManagedChannelImpl{logId=348, target=directaddress:///6f89ffb5-fd55-46c3-94fa-1700e758525b} was not shutdown properly!!! ~*~*~*
    Make sure to call shutdown()/shutdownNow() and wait until awaitTermination() returns true.

Which means that the GrpcCleanupRule did not fire (because JUnit 5 doesn’t support @Rule), and the previous server instance was replaced without shutting down correctly.

Grpc Cleanup Extension Use

Here’s a look at how our Grpc Services are set up and how we use the cleanup Extension in tests.

Let’s say we have a Grpc Service implementation called integration service that’s set up like this:

@GRpcService
@Component
class IntegrationService() : IntegrationServiceGrpc.IntegrationServiceImplBase() {
    // do integration stuff
}

When we go to write a test forIntegrationServer, we set up a GrpcCleanupExtension
instance. The syntax is a little non-Kotlinesque because extensions need to be final static member variables (not in a companion object).

@ExtendWith(SpringExtension::class)
class IntegrationServiceTests {
    @JvmField
    @RegisterExtension
    final val grpcCleanupExtension = GrpcCleanupExtension() // yes, the final is needed in kotlin

    lateinit var blockingStub: IntegrationServiceGrpc.IntegrationServiceBlockingStub

    @BeforeEach
    fun setup() {
        val integrationService = IntegrationService()
        blockingStub = IntegrationServiceGrpc.newBlockingStub(grpcCleanupExtension.addService(integrationService))
    }
    
    @Test
    fun testThingOne() {
        assertNotNull(blockingStub.thingOne())
    }
}

And that’s it!

How The Cleanup Extension Works

Not surprisingly, the GrpcCleanupExtension looks very similar to the Grpc cleanup Rule. It keeps track of a list of servers/channels, and shuts them down after each test execution in the after each function.

class GrpcCleanupExtension : AfterEachCallback {

    private var cleanupTargets: MutableList<CleanupTarget> = mutableListOf()

    companion object {
        private val logger = LoggerFactory.getLogger(GrpcCleanupExtension::class.java)

        const val TERMINATION_TIMEOUT_MS = 250L
        const val MAX_NUM_TERMINATIONS = 100
    }

    fun addService(service: BindableService): ManagedChannel {
        val serverName: String = InProcessServerBuilder.generateName()

        cleanupTargets.add(
            ServerCleanupTarget(
                InProcessServerBuilder
                    .forName(serverName)
                    .directExecutor()
                    .intercept(GlobalGrpcExceptionHandler())
                    .addService(service)
                    .build()
                    .start()
            )
        )

        val channel = InProcessChannelBuilder.forName(serverName)
            .directExecutor()
            .build()

        cleanupTargets.add(ManagedChannelCleanupTarget(channel))

        return channel
    }

    override fun afterEach(context: ExtensionContext?) {
        cleanupTargets.forEach { cleanupTarget ->
            try {
                var count = 0
                cleanupTarget.shutdown()
                do {
                    cleanupTarget.awaitTermination(TERMINATION_TIMEOUT_MS, TimeUnit.MILLISECONDS)
                    count++
                    if (count > MAX_NUM_TERMINATIONS) {
                        logger.error("Hit max count $count trying to shut down down cleanupTarget $cleanupTarget")
                        break
                    }
                } while (!cleanupTarget.isTerminated())
            } catch (e: Exception) {
                logger.error("Problem shutting down cleanupTarget $cleanupTarget", e)
            }
        }

        if (isAllTerminated()) {
            cleanupTargets.clear()
        } else {
            logger.error("Not all cleanupTargets are terminated")
        }
    }

    fun isAllTerminated(): Boolean = cleanupTargets.all { it.isTerminated() }
}

interface CleanupTarget {
    fun shutdown()
    fun awaitTermination(timeout: Long, timeUnit: TimeUnit): Boolean
    fun isTerminated(): Boolean
}

class ServerCleanupTarget(private val server: Server) : CleanupTarget {
    override fun shutdown() {
        server.shutdown()
    }

    override fun awaitTermination(timeout: Long, timeUnit: TimeUnit): Boolean =
        server.awaitTermination(timeout, timeUnit)

    override fun isTerminated(): Boolean = server.isTerminated
}

class ManagedChannelCleanupTarget(private val managedChannel: ManagedChannel) : CleanupTarget {
    override fun shutdown() {
        managedChannel.shutdown()
    }

    override fun awaitTermination(timeout: Long, timeUnit: TimeUnit): Boolean =
        managedChannel.awaitTermination(timeout, timeUnit)

    override fun isTerminated(): Boolean = managedChannel.isTerminated
}

C’est La Fin

In order to avoid JUnit 4 when testing Grpc service classes, we wrote a GrpcCleanupExtension. It hooks into the JUnit 5 test lifecycle to clean up server/channel resources. Enjoy!