Message deduplication for exchanges

shard.lock

@@ -16,6 +16,10 @@ shards:
     git: https://github.com/84codes/lz4.cr.git
     version: 1.0.0+git.commit.96d714f7593c66ca7425872fd26c7b1286806d3d
 
+  mqtt-protocol:
+    git: https://github.com/84codes/mqtt-protocol.cr.git
+    version: 0.2.0+git.commit.3f82ee85d029e6d0505cbe261b108e156df4e598
+
   systemd:
     git: https://github.com/84codes/systemd.cr.git
     version: 2.0.0

shard.yml

@@ -32,6 +32,8 @@ dependencies:
     github: 84codes/systemd.cr
   lz4:
     github: 84codes/lz4.cr
+  mqtt-protocol:
+    github: 84codes/mqtt-protocol.cr
 
 development_dependencies:
   ameba:

spec/clustering_spec.cr

@@ -2,6 +2,8 @@ require "./spec_helper"
 require "../src/lavinmq/clustering/client"
 require "../src/lavinmq/clustering/controller"
 
+alias IndexTree = LavinMQ::MQTT::TopicTree(String)
+
 describe LavinMQ::Clustering::Client do
   follower_data_dir = "/tmp/lavinmq-follower"
 
@@ -72,6 +74,46 @@ describe LavinMQ::Clustering::Client do
     end
   end
 
+  it "replicates and streams retained messages to followers" do
+    replicator = LavinMQ::Clustering::Server.new(LavinMQ::Config.instance, LavinMQ::Etcd.new, 0)
+    tcp_server = TCPServer.new("localhost", 0)
+
+    spawn(replicator.listen(tcp_server), name: "repli server spec")
+    config = LavinMQ::Config.new.tap &.data_dir = follower_data_dir
+    repli = LavinMQ::Clustering::Client.new(config, 1, replicator.password, proxy: false)
+    done = Channel(Nil).new
+    spawn(name: "follow spec") do
+      repli.follow("localhost", tcp_server.local_address.port)
+      done.send nil
+    end
+    wait_for { replicator.followers.size == 1 }
+
+    retain_store = LavinMQ::MQTT::RetainStore.new("#{LavinMQ::Config.instance.data_dir}/retain_store", replicator)
+    props = LavinMQ::AMQP::Properties.new
+    msg1 = LavinMQ::Message.new(100, "test", "rk", props, 10, IO::Memory.new("body1"))
+    msg2 = LavinMQ::Message.new(100, "test", "rk", props, 10, IO::Memory.new("body2"))
+    retain_store.retain("topic1", msg1.body_io, msg1.bodysize)
+    retain_store.retain("topic2", msg2.body_io, msg2.bodysize)
+
+    wait_for(10.seconds) { replicator.followers.first?.try &.lag_in_bytes == 0 }
+    repli.close
+    done.receive
+
+    follower_retain_store = LavinMQ::MQTT::RetainStore.new("#{follower_data_dir}/retain_store", LavinMQ::Clustering::NoopServer.new)
+    a = Array(String).new(2)
+    b = Array(String).new(2)
+    follower_retain_store.each("#") do |topic, bytes|
+      a << topic
+      b << String.new(bytes)
+    end
+
+    a.sort!.should eq(["topic1", "topic2"])
+    b.sort!.should eq(["body1", "body2"])
+    follower_retain_store.retained_messages.should eq(2)
+  ensure
+    replicator.try &.close
+  end
+
   it "can stream full file" do
     replicator = LavinMQ::Clustering::Server.new(LavinMQ::Config.instance, LavinMQ::Etcd.new, 0)
     tcp_server = TCPServer.new("localhost", 0)

spec/deduper_spec.cr

@@ -0,0 +1,138 @@
+require "./spec_helper"
+require "../src/lavinmq/exchange/dedup_ext.cr"
+
+describe LavinMQ::Deduplication do
+  describe LavinMQ::Deduplication::MemoryCache do
+    it "should have a max size" do
+      cache = LavinMQ::Deduplication::MemoryCache(String).new(2)
+      cache.insert("item1")
+      cache.insert("item2")
+      cache.insert("item3")
+      cache.contains?("item1").should be_false
+      cache.contains?("item2").should be_true
+      cache.contains?("item3").should be_true
+    end
+
+    it "should store item without ttl" do
+      cache = LavinMQ::Deduplication::MemoryCache(String).new(10)
+      cache.insert("item1")
+      cache.contains?("item1").should be_true
+      cache.contains?("item2").should be_false
+    end
+
+    it "should respect ttl" do
+      cache = LavinMQ::Deduplication::MemoryCache(String).new(3)
+      cache.insert("item1", 1)
+      cache.insert("item2", 300)
+      cache.insert("item3")
+      sleep 0.2.seconds
+      cache.contains?("item1").should be_false
+      cache.contains?("item2").should be_true
+      cache.contains?("item3").should be_true
+    end
+  end
+end
+
+class MockCache(T) < LavinMQ::Deduplication::Cache(T)
+  @counter = Hash(String, Array({T, UInt32?})).new do |h, k|
+    h[k] = Array({T, UInt32?}).new
+  end
+
+  def contains?(key : T) : Bool
+    @counter["contains?"] << {key, nil}
+    false
+  end
+
+  def insert(key : T, ttl : UInt32? = nil)
+    @counter["insert"] << {key, ttl}
+  end
+
+  def calls(key : String) : Array({T, UInt32?})
+    @counter[key]
+  end
+end
+
+describe LavinMQ::Deduplication::Deduper do
+  describe "duplicate?" do
+    it "should return false if \"x-deduplication-header\" is missing (no identifier, always unique)" do
+      mock = MockCache(AMQ::Protocol::Field).new
+      deduper = LavinMQ::Deduplication::Deduper.new(mock)
+      props = LavinMQ::AMQP::Properties.new
+      msg = LavinMQ::Message.new("ex", "rk", "body", props)
+      res = deduper.duplicate?(msg)
+      res.should be_false
+    end
+
+    it "should check cache if entry exists" do
+      mock = MockCache(AMQ::Protocol::Field).new
+      deduper = LavinMQ::Deduplication::Deduper.new(mock)
+      props = LavinMQ::AMQP::Properties.new(headers: LavinMQ::AMQP::Table.new({
+        "x-deduplication-header" => "msg1",
+      }))
+      msg = LavinMQ::Message.new("ex", "rk", "body", props)
+      deduper.duplicate?(msg)
+      mock.calls("contains?").size.should eq 1
+    end
+
+    it "should only insert into cache if header has a value" do
+      mock = MockCache(AMQ::Protocol::Field).new
+      deduper = LavinMQ::Deduplication::Deduper.new(mock)
+      props = LavinMQ::AMQP::Properties.new(headers: LavinMQ::AMQP::Table.new)
+      msg = LavinMQ::Message.new("ex", "rk", "body", props)
+      deduper.add(msg)
+      mock.calls("insert").size.should eq 0
+    end
+
+    it "should only insert into cache if header has a value" do
+      mock = MockCache(AMQ::Protocol::Field).new
+      deduper = LavinMQ::Deduplication::Deduper.new(mock)
+      props = LavinMQ::AMQP::Properties.new(headers: LavinMQ::AMQP::Table.new({
+        "x-deduplication-header" => "msg1",
+      }))
+      msg = LavinMQ::Message.new("ex", "rk", "body", props)
+      deduper.add(msg)
+      mock.calls("insert").size.should eq 1
+    end
+
+    it "should respect x-cache-ttl on message" do
+      mock = MockCache(AMQ::Protocol::Field).new
+      deduper = LavinMQ::Deduplication::Deduper.new(mock)
+      props = LavinMQ::AMQP::Properties.new(headers: LavinMQ::AMQP::Table.new({
+        "x-deduplication-header" => "msg1",
+        "x-cache-ttl"            => 10,
+      }))
+      msg = LavinMQ::Message.new("ex", "rk", "body", props)
+      deduper.add(msg)
+      calls = mock.calls("insert")
+      calls.first[0].should eq "msg1"
+      calls.first[1].should eq 10
+    end
+
+    it "should fallback to default ttl" do
+      mock = MockCache(AMQ::Protocol::Field).new
+      deduper = LavinMQ::Deduplication::Deduper.new(mock, 12)
+      props = LavinMQ::AMQP::Properties.new(headers: LavinMQ::AMQP::Table.new({
+        "x-deduplication-header" => "msg1",
+      }))
+      msg = LavinMQ::Message.new("ex", "rk", "body", props)
+      deduper.add(msg)
+      calls = mock.calls("insert")
+      calls.first[0].should eq "msg1"
+      calls.first[1].should eq 12
+    end
+
+    it "should prio message ttl over default ttl" do
+      mock = MockCache(AMQ::Protocol::Field).new
+      deduper = LavinMQ::Deduplication::Deduper.new(mock, 12)
+      props = LavinMQ::AMQP::Properties.new(headers: LavinMQ::AMQP::Table.new({
+        "x-deduplication-header" => "msg1",
+        "x-cache-ttl"            => 10,
+      }))
+      msg = LavinMQ::Message.new("ex", "rk", "body", props)
+      deduper.add(msg)
+      calls = mock.calls("insert")
+      calls.first[0].should eq "msg1"
+      calls.first[1].should eq 10
+    end
+  end
+end

spec/exchange_spec.cr

@@ -170,4 +170,35 @@ describe LavinMQ::Exchange do
       end
     end
   end
+  describe "message deduplication" do
+    args = AMQP::Client::Arguments.new({
+      "x-message-deduplication" => true,
+      "x-cache-size"            => 10,
+    })
+
+    it "should handle message deduplication" do
+      with_amqp_server do |s|
+        with_channel(s) do |ch|
+          ch.exchange("test", "topic", args: args)
+          ch.queue.bind("test", "#")
+          ex = s.vhosts["/"].exchanges["test"]
+          q = s.vhosts["/"].queues.first_value
+          props = LavinMQ::AMQP::Properties.new(headers: LavinMQ::AMQP::Table.new({
+            "x-deduplication-header" => "msg1",
+          }))
+          msg = LavinMQ::Message.new("ex", "rk", "body", props)
+          ex.publish(msg, false).should eq 1
+          ex.dedup_count.should eq 0
+          props = LavinMQ::AMQP::Properties.new(headers: LavinMQ::AMQP::Table.new({
+            "x-deduplication-header" => "msg1",
+          }))
+          msg = LavinMQ::Message.new("ex", "rk", "body", props)
+          ex.publish(msg, false).should eq 0
+          ex.dedup_count.should eq 1
+
+          q.message_count.should eq 1
+        end
+      end
+    end
+  end
 end

spec/message_routing_spec.cr

@@ -421,3 +421,42 @@ describe LavinMQ::Exchange do
     end
   end
 end
+
+describe LavinMQ::MQTT::Exchange do
+  it "should only allow Session to bind" do
+    with_amqp_server do |s|
+      vhost = s.vhosts.create("x")
+      q1 = LavinMQ::AMQP::Queue.new(vhost, "q1")
+      s1 = LavinMQ::MQTT::Session.new(vhost, "q1")
+      index = LavinMQ::MQTT::TopicTree(String).new
+      store = LavinMQ::MQTT::RetainStore.new("tmp/retain_store", LavinMQ::Clustering::NoopServer.new, index)
+      x = LavinMQ::MQTT::Exchange.new(vhost, "", store)
+      x.bind(s1, "s1", LavinMQ::AMQP::Table.new)
+      expect_raises(LavinMQ::Exchange::AccessRefused) do
+        x.bind(q1, "q1", LavinMQ::AMQP::Table.new)
+      end
+    end
+  end
+
+  it "publish messages to queues with it's own publish method" do
+    with_amqp_server do |s|
+      vhost = s.vhosts.create("x")
+      s1 = LavinMQ::MQTT::Session.new(vhost, "session 1")
+      index = LavinMQ::MQTT::TopicTree(String).new
+      store = LavinMQ::MQTT::RetainStore.new("tmp/retain_store", LavinMQ::Clustering::NoopServer.new, index)
+      x = LavinMQ::MQTT::Exchange.new(vhost, "mqtt.default", store)
+      x.bind(s1, "s1", LavinMQ::AMQP::Table.new)
+      pub_args = {
+        packet_id: 1u16,
+        payload:   Bytes.new(0),
+        dup:       false,
+        qos:       0u8,
+        retain:    false,
+        topic:     "s1",
+      }
+      msg = MQTT::Protocol::Publish.new(**pub_args)
+      x.publish(msg)
+      s1.message_count.should eq 1
+    end
+  end
+end