diff --git a/src/riak_pipe_vnode_worker.erl b/src/riak_pipe_vnode_worker.erl
index 449210d..87de791 100644
--- a/src/riak_pipe_vnode_worker.erl
+++ b/src/riak_pipe_vnode_worker.erl
@@ -139,7 +139,6 @@
          send_output/3,
          send_output/4,
          send_output/5]).
--export([behaviour_info/1]).
 
 %% gen_fsm callbacks
 -export([
@@ -163,14 +162,16 @@
 -opaque state() :: #state{}.
 -export_type([state/0]).
 
-%% @doc Get information about this behavior.
--spec behaviour_info(atom()) -> 'undefined' | [{atom(), arity()}].
-behaviour_info(callbacks) ->
-    [{init,2},
-     {process,3},
-     {done,1}];
-behaviour_info(_Other) ->
-    undefined.
+-type callback_state() :: term().
+
+-callback init(riak_pipe_vnode:partition(),
+               riak_pipe_fitting:details()) ->
+    {ok, callback_state()}.
+-callback process(term(), boolean(), callback_state()) ->
+    {ok, callback_state()} |
+    {forward_preflist, callback_state()} |
+    {{error, term()}, callback_state()}.
+-callback done(callback_state()) -> ok.
 
 %%%===================================================================
 %%% API