Commits

dizzyd  committed 09d189c

Add support for terminate/2 call on drivers; enables driver-specific cleanup prior to exit

  • Participants
  • Parent commits 168ae92

Comments (0)

Files changed (1)

File src/basho_bench_worker.erl

 
 
 worker_init(State) ->
+    %% Trap exits from linked parent process; use this to ensure the driver
+    %% gets a chance to cleanup
+    process_flag(trap_exit, true),
     random:seed(State#state.rng_seed),
     worker_idle_loop(State).
 
             %% Driver crashed, generate a crash error and terminate. This will take down
             %% the corresponding worker which will get restarted by the appropriate supervisor.
             basho_bench_stats:op_complete(Next, {error, crash}, ElapsedUs),
+
+            %% Give the driver a chance to cleanup
+            (catch (State#state.driver):terminate({'EXIT', Reason}, State#state.driver_state)),
+
             ?DEBUG("Driver ~p crashed: ~p\n", [State#state.driver, Reason]),
             crash;
 
             %% Driver (or something within it) has requested that this worker
             %% terminate cleanly.
             ?INFO("Driver ~p (~p) has requested stop: ~p\n", [State#state.driver, self(), Reason]),
+
+            %% Give the driver a chance to cleanup
+            (catch (State#state.driver):terminate(normal, State#state.driver_state)),
+
             normal
     end.
 
+needs_shutdown(State) ->
+    Parent = State#state.parent_pid,
+    receive
+        {'EXIT', Parent, _Reason} ->
+            %% Give the driver a chance to cleanup
+            (catch (State#state.driver):terminate(normal, State#state.driver_state)),
+            true
+    after 0 ->
+            false
+    end.
+
+
 max_worker_run_loop(State) ->
     case worker_next_op(State) of
         {ok, State2} ->
-            max_worker_run_loop(State2);
+            case needs_shutdown(State2) of
+                true ->
+                    ok;
+                false ->
+                    max_worker_run_loop(State2)
+            end;
         ExitReason ->
             exit(ExitReason)
     end.
     timer:sleep(trunc(stats_rv:exponential(Lambda))),
     case worker_next_op(State) of
         {ok, State2} ->
-            rate_worker_run_loop(State2, Lambda);
+            case needs_shutdown(State2) of
+                true ->
+                    ok;
+                false ->
+                    rate_worker_run_loop(State2, Lambda)
+            end;
         ExitReason ->
             exit(ExitReason)
     end.