JDBI, Race Conditions and You

Here at Twilio, we've built a number of RESTful services using Dropwizard. Dropwizard assembles "stable, mature libraries from the Java ecosystem into a simple, light-weight package".

During a recent deploy, we noticed some odd behavior immediately after service startup, in which queries meant for one database were being executed on another. Some investigation into our data revealed that these occurrences were correlated with the first few moments after deploying the service. After several days of debugging, we tracked down the source of the problem to our usage of JDBI for executing SQL queries.

JDBI APIs

Included with Dropwizard is JDBI, a library that provides a much improved interface over JDBC for working with SQL databases in Java. JDBI provides two different style APIs, a fluent style and an object style. The fluent style is similar to the JDBC interface:

DataSource ds = JdbcConnectionPool.create("jdbc:h2:mem:test",
                                          "username",
                                          "password");
DBI dbi = new DBI(ds);
Handle h = dbi.open();

h.execute("insert into foo (bar, baz) values (?, ?)", 10, "beep");
h.close();

The object style API provides an idiomatic Java interface for executing queries. Here's a simple example of a data access object (DAO) interface which defines the insert query seen above:

public interface FooDAO
{
  @SqlUpdate("insert into foo (bar, baz) values (:bar, :baz)")
  void insert(@Bind("bar") int bar, @Bind("baz") String baz);
}

Once we have defined the interface, we use JDBI to create instances of our DAO:

JdbcConnectionPool ds = JdbcConnectionPool.create("jdbc:h2:mem:test2",
                                                  "username",
                                                  "password");
DBI dbi = new DBI(ds);

MyDAO dao = dbi.onDemand(MyDAO.class);
dao.insert(10, "beep");

Thats it! The object style API does most of the heavy lifting with a nice clean interface. Both styles make use of the DBI object which encapsulates a connection to a specific database. DBI.onDemand takes the class object corresponding to our interface and instantiates a DAO tied to our DBI. One thing to note is that with the JDBI object API, queries are expressed via annotations, the values of which must be compile-time constants. This provides some certainty about what queries can be run, but can also create headaches when you have many similar queries as each one must have its own annotated method signature.

Under the Hood

But how does JDBI actually create implementations of our interface? The heart of this functionality lies in JDBI's SqlObject class. Of particular interest is the buildSqlObject method.

Looking at this method, we see the use of the Enhancer class and the Factory interface. These are both part of cglib, a code generation library used to create dynamic proxy objects and intercept field access. (That sounds like exactly the kind of thing that would allow JDBI to create instances of our DAO interfaces.)

Enhancers are the primary class for creating dynamic proxy objects in cglib. For cases where we want to make many instances of our proxy object, the cglib Javadoc recommends using the Factory to create instances: "Using this [Factory] interface for new instances is faster than going through the Enhancer interface".

But where do we get Factory instances if Enhancer creates instances of our proxy objects? Well, "all enhanced instances returned by the Enhancer class implement this [Factory] interface." So to create a new Factory for our proxy object, we just create a single instance and cast it to a Factory.

Concurrent Execution

Let's consider the creation of our generated DAO objects in a concurrent context. Assume we've just deployed our code, and two separate threads are getting instances of the same DAO concurrently, so there is no available Factory instance.

T t = (T) e.create();
T actual = (T) factories.putIfAbsent(sqlObjectType, (Factory) t);
return actual != null ? actual : t;

Both threads instantiate instances via an Enhancer. Then, each thread attempts to insert its newly created instance into factories, with the putIfAbsent method. This method inserts the given value into the map if it isn't already set, returning null if there was no existing value, or the existing value.

Since ConcurrentHashMap's putIfAbsent method is atomic, one of our threads will receive null from this call, while the other will receive the instance created by the other thread. At this point, both threads will return the same instance of our DAO, handle and all!

Fixing the Race

Now that we've identified the problem, this fix for this race condition is fairly trivial:

- return actual != null ? actual : t;
+ if (actual == null) {
+   return t;
+ }
+ f = (Factory) actual;

Instead of returning an existing instance from our factory map, we simply assign our local Factory variable, f, with the returned value, and let execution continue.

Down the Rabbit Hole

But let's consider how we got here. Here's the diff which introduced our race condition:

- factories.putIfAbsent(sqlObjectType, (Factory) t);
- return t;
+ T actual = (T) factories.putIfAbsent(sqlObjectType, (Factory) t);
+ return actual != null ? actual : t;

The description of the commit gives us some context around this change: "Fix findbugs, pmd and javadoc problems." Sure enough, reverting this change, we see that findbugs complains with the original code when we run 'mvn install':

[INFO] --- findbugs-maven-plugin:2.5.2:check (default) @ jdbi ---
[INFO] BugInstance size is 1
[INFO] Error size is 0
[INFO] Total bugs: 1
[INFO] Return value of putIfAbsent is ignored, but t is reused in org.skife.jdbi.v2.sqlobject.SqlObject.buildSqlObject(Class, HandleDing) ["org.skife.jdbi.v2.sqlobject.SqlObject"] At SqlObject.java:[lines 38-221]
[INFO] ------------------------------------------------------------------------
[INFO] BUILD FAILURE

Since findbugs doesn't like the original implementation, another way to prevent this type of issue from occurring would be to only intercept methods on the instances created from a Factory. So let's see what happens when we run the tests after removing the initial callback altogether, and only add callbacks when we call Factory.newInstance:

testSqlObjectApi(org.skife.jdbi.v2.docs.TestFoldToObjectGraph)  Time elapsed: 0.522 sec  <<< ERROR!
java.lang.IllegalStateException: Callbacks are required
    at net.sf.cglib.proxy.Enhancer.validate(Enhancer.java:333)
    at net.sf.cglib.proxy.Enhancer.createHelper(Enhancer.java:371)
    at net.sf.cglib.proxy.Enhancer.create(Enhancer.java:285)
    at org.skife.jdbi.v2.sqlobject.SqlObject.buildSqlObject(SqlObject.java:77)
    at org.skife.jdbi.v2.sqlobject.SqlObjectBuilder.attach(SqlObjectBuilder.java:39)
    at org.skife.jdbi.v2.BasicHandle.attach(BasicHandle.java:380)
    at org.skife.jdbi.v2.docs.TestFoldToObjectGraph.testSqlObjectApi(TestFoldToObjectGraph.java:130)

Hmm, not good, looks like we need to add callbacks. Ok, since we're going to override our callbacks later anyway, let's just set an initial callback that throws an exception, like this:

public Object intercept(Object o, Method method, Object[] objects, MethodProxy methodProxy) throws Throwable {
    throw new UnsupportedOperationException("Cannnot intercept on factory objects.");
}

With this change, now all our JDBI tests pass, so let's include it in our project and run our tests. Uh oh:

org.skife.jdbi.cglib.core.CodeGenerationException: java.lang.UnsupportedOperationException-->Cannnot intercept on factory objects.
    at org.skife.jdbi.cglib.core.ReflectUtils.newInstance(ReflectUtils.java:235)
    at org.skife.jdbi.cglib.core.ReflectUtils.newInstance(ReflectUtils.java:220)
    at org.skife.jdbi.cglib.core.ReflectUtils.newInstance(ReflectUtils.java:216)
    at org.skife.jdbi.cglib.proxy.Enhancer.createUsingReflection(Enhancer.java:643)
    at org.skife.jdbi.cglib.proxy.Enhancer.firstInstance(Enhancer.java:538)
    at org.skife.jdbi.cglib.core.AbstractClassGenerator.create(AbstractClassGenerator.java:231)
    at org.skife.jdbi.cglib.proxy.Enhancer.createHelper(Enhancer.java:377)
    at org.skife.jdbi.cglib.proxy.Enhancer.create(Enhancer.java:285)
    at org.skife.jdbi.v2.sqlobject.SqlObject.buildSqlObject(SqlObject.java:77)
    at org.skife.jdbi.v2.sqlobject.SqlObjectBuilder.onDemand(SqlObjectBuilder.java:67)
    at org.skife.jdbi.v2.DBI.onDemand(DBI.java:338)

We're now intercepting every method call on our Factory, including those made by cglib itself during initial creation! At this point the rabbit-hole was getting too deep to safely continue with this approach, especially considering the final words of warning in this post, cglib: The missing manual.

What now?

The fix for this case was submitted upstream on 4/22/2014, and merged into master on 4/28/2014. So if you're pulling down JDBI master on every deploy and building from source, you're all set. But what if you're using JDBI as a transitive dependency of your framework? Even on the latest release of DropWizard, 0.7.0, you're pulling in JDBI 2.53, which still has the broken behavior.

At minimum, I recommend auditing your application's use of DBI.open, DBI.onDemand, Handle.attach, and the static SqlObjectBuilder equivalents. As long as these methods are called only from a single thread, the race condition can be avoided, as the DAOs are generally safe to use across threads.

Another option is to avoid using the generated DAOs entirely. In this case, I recommend taking a look at jOOQ. jOOQ provides a builder interface for safely constructing SQL queries. Here's our previous example, building the query with jOOQ, and executing it with JDBI:

String sql = create.insertInto(table("foo"), field("bar"), field("baz", classOf[Integer]).values(param("barVal"), param("bazVal")).getSQL(ParamType.NAMED)

JdbcConnectionPool ds = JdbcConnectionPool.create("jdbc:h2:mem:test2",
                                                  "username",
                                                  "password");
DBI dbi = new DBI(ds);

Statement statement = dbi.createQuery(sql)
statement.bind("barVal", 10).bind("bazVal", "beep").execute()

While maybe not as clean as the JDBI object API, jOOQ is still a major improvement over writing our own SQL directly. As a bonus, it provides a much greater degree of flexibility with crafting our queries, as our SQL can be assembled at runtime, while avoiding direct manipulation of raw SQL.