Prevent lost updates in concurrent writes

Author: preciz

lib/abit.ex

@@ -100,9 +100,7 @@ defmodule Abit do
   defp do_union(ref_a, _, 0), do: ref_a
 
   defp do_union(ref_a, ref_b, index) do
-    unioned_value = :atomics.get(ref_a, index) ||| :atomics.get(ref_b, index)
-
-    :atomics.put(ref_a, index, unioned_value)
+    update_atomic(ref_a, ref_b, index, :union)
 
     do_union(ref_a, ref_b, index - 1)
   end
@@ -124,9 +122,7 @@ defmodule Abit do
   defp do_intersect(ref_a, _, 0), do: ref_a
 
   defp do_intersect(ref_a, ref_b, index) do
-    intersected_value = :atomics.get(ref_a, index) &&& :atomics.get(ref_b, index)
-
-    :atomics.put(ref_a, index, intersected_value)
+    update_atomic(ref_a, ref_b, index, :intersect)
 
     do_intersect(ref_a, ref_b, index - 1)
   end
@@ -150,9 +146,7 @@ defmodule Abit do
   defp do_difference(ref_a, _, 0), do: ref_a
 
   defp do_difference(ref_a, ref_b, index) do
-    diff_value = :atomics.get(ref_a, index) &&& bnot(:atomics.get(ref_b, index))
-
-    :atomics.put(ref_a, index, diff_value)
+    update_atomic(ref_a, ref_b, index, :difference)
 
     do_difference(ref_a, ref_b, index - 1)
   end
@@ -175,13 +169,35 @@ defmodule Abit do
   defp do_symmetric_difference(ref_a, _, 0), do: ref_a
 
   defp do_symmetric_difference(ref_a, ref_b, index) do
-    xor_value = :atomics.get(ref_a, index) |> bxor(:atomics.get(ref_b, index))
-
-    :atomics.put(ref_a, index, xor_value)
+    update_atomic(ref_a, ref_b, index, :symmetric_difference)
 
     do_symmetric_difference(ref_a, ref_b, index - 1)
   end
 
+  defp update_atomic(ref_a, ref_b, index, operation) do
+    current_value = :atomics.get(ref_a, index)
+
+    do_update_atomic(ref_a, ref_b, index, operation, current_value)
+  end
+
+  defp do_update_atomic(ref_a, ref_b, index, operation, current_value) do
+    other_value = :atomics.get(ref_b, index)
+    next_value = apply_bitwise_operation(operation, current_value, other_value)
+
+    case :atomics.compare_exchange(ref_a, index, current_value, next_value) do
+      :ok ->
+        :ok
+
+      new_current_value ->
+        do_update_atomic(ref_a, ref_b, index, operation, new_current_value)
+    end
+  end
+
+  defp apply_bitwise_operation(:union, left, right), do: left ||| right
+  defp apply_bitwise_operation(:intersect, left, right), do: left &&& right
+  defp apply_bitwise_operation(:difference, left, right), do: left &&& bnot(right)
+  defp apply_bitwise_operation(:symmetric_difference, left, right), do: bxor(left, right)
+
   @doc """
   Inverts all bits in the signed atomics reference `ref` using bitwise NOT.
 

lib/abit/counter.ex

@@ -155,7 +155,7 @@ defmodule Abit.Counter do
   end
 
   def put(
-        %Counter{atomics_ref: atomics_ref, signed: signed, counters_bit_size: counters_bit_size},
+        %Counter{atomics_ref: atomics_ref, counters_bit_size: counters_bit_size} = counter,
         index,
         value
       )
@@ -164,12 +164,36 @@ defmodule Abit.Counter do
 
     atomics_value = :atomics.get(atomics_ref, atomics_index)
 
+    do_put(counter, index, value, atomics_index, bit_index, atomics_value)
+  end
+
+  defp do_put(
+         %Counter{
+           atomics_ref: atomics_ref,
+           signed: signed,
+           counters_bit_size: counters_bit_size
+         } = counter,
+         index,
+         value,
+         atomics_index,
+         bit_index,
+         atomics_value
+       ) do
     {final_counter_value, <<next_atomics_value::64>>} =
       put_value(signed, counters_bit_size, bit_index, <<atomics_value::64>>, value)
 
-    :atomics.put(atomics_ref, atomics_index, next_atomics_value)
+    case :atomics.compare_exchange(
+           atomics_ref,
+           atomics_index,
+           atomics_value,
+           next_atomics_value
+         ) do
+      :ok ->
+        {:ok, {index, final_counter_value}}
 
-    {:ok, {index, final_counter_value}}
+      new_atomics_value ->
+        do_put(counter, index, value, atomics_index, bit_index, new_atomics_value)
+    end
   end
 
   @doc """

test/abit/counter_test.exs

@@ -144,6 +144,21 @@ defmodule Abit.CounterTest do
     for i <- 1..8, do: assert Counter.get(counter, i) == 1
   end
 
+  test "put/3 preserves concurrent writes to counters in the same atomics integer" do
+    for _trial <- 1..100 do
+      counter = Counter.new(8, 8, signed: false)
+
+      operations =
+        for index <- 0..7 do
+          fn -> Counter.put(counter, index, index + 1) end
+        end
+
+      run_simultaneously(operations)
+
+      assert Enum.to_list(counter) == Enum.to_list(1..8)
+    end
+  end
+
   test "Enumerable protocol implementation" do
     counter = Counter.new(10, 8)
     Counter.put(counter, 1, 42)
@@ -192,4 +207,30 @@ defmodule Abit.CounterTest do
       assert Counter.get(c, 0) == 10
     end
   end
+
+  defp run_simultaneously(operations) do
+    parent = self()
+    gate = make_ref()
+
+    tasks =
+      Enum.map(operations, fn operation ->
+        Task.async(fn ->
+          send(parent, {gate, :ready, self()})
+
+          receive do
+            {^gate, :go} -> operation.()
+          end
+        end)
+      end)
+
+    pids =
+      for _task <- tasks do
+        receive do
+          {^gate, :ready, pid} -> pid
+        end
+      end
+
+    Enum.each(pids, &send(&1, {gate, :go}))
+    Enum.each(tasks, &Task.await/1)
+  end
 end

test/abit_test.exs

@@ -170,6 +170,41 @@ defmodule AbitTest do
     end
   end
 
+  test "binary operations preserve concurrent updates to the same atomic element" do
+    operation_cases = [
+      {:union, 0, 0xFF},
+      {:intersect, 0xFF, 0},
+      {:difference, 0xFF, 0},
+      {:symmetric_difference, 0, 0xFF}
+    ]
+
+    for {operation, initial_value, expected_value} <- operation_cases,
+        _trial <- 1..100 do
+      ref_l = :atomics.new(1, signed: false)
+      :atomics.put(ref_l, 1, initial_value)
+
+      operations =
+        for bit <- 0..7 do
+          ref_r = :atomics.new(1, signed: false)
+          bit_value = Bitwise.bsl(1, bit)
+
+          right_value =
+            case operation do
+              :intersect -> Bitwise.bxor(0xFF, bit_value)
+              _other -> bit_value
+            end
+
+          :atomics.put(ref_r, 1, right_value)
+
+          fn -> apply(Abit, operation, [ref_l, ref_r]) end
+        end
+
+      run_simultaneously(operations)
+
+      assert :atomics.get(ref_l, 1) == expected_value
+    end
+  end
+
   test "invert atomics bit arrays returns reference" do
     ref = :atomics.new(2, signed: true)
 
@@ -350,4 +385,30 @@ defmodule AbitTest do
       assert Abit.set_bits_count(ref) == 96
     end
   end
+
+  defp run_simultaneously(operations) do
+    parent = self()
+    gate = make_ref()
+
+    tasks =
+      Enum.map(operations, fn operation ->
+        Task.async(fn ->
+          send(parent, {gate, :ready, self()})
+
+          receive do
+            {^gate, :go} -> operation.()
+          end
+        end)
+      end)
+
+    pids =
+      for _task <- tasks do
+        receive do
+          {^gate, :ready, pid} -> pid
+        end
+      end
+
+    Enum.each(pids, &send(&1, {gate, :go}))
+    Enum.each(tasks, &Task.await/1)
+  end
 end