Merge pull request #28 from JuliaComputing/tan/misc

fix race condition in basic_consume

Author: tanmaykm

src/protocol.jl

@@ -210,8 +210,11 @@ mutable struct MessageConsumer
     callback::Function
     receiver::Task
 
-    function MessageConsumer(chan_id::TAMQPChannel, consumer_tag::String, callback::Function; buffer_size::Int=typemax(Int))
-        c = new(chan_id, consumer_tag, Channel{Message}(buffer_size), callback)
+    function MessageConsumer(chan_id::TAMQPChannel, consumer_tag::String, callback::Function;
+        buffer_size::Int=typemax(Int),
+        buffer::Channel{Message}=Channel{Message}(buffer_size))
+
+        c = new(chan_id, consumer_tag, buffer, callback)
         c.receiver = @async connection_processor(c, "Consumer $consumer_tag", channel_message_consumer)
         c
     end
@@ -232,14 +235,16 @@ mutable struct MessageChannel <: AbstractChannel
     partial_msgs::Vector{Message} # holds partial messages while they are getting read (message bodies arrive in sequence)
     chan_get::Channel{Union{Message, Nothing}}  # channel used for received messages, in sync get call (TODO: maybe type more strongly?)
     consumers::Dict{String,MessageConsumer}
+    pending_msgs::Dict{String,Channel{Message}} # holds messages received that do not have a consumer registered
+    lck::ReentrantLock
 
     closereason::Union{CloseReason, Nothing}
 
     function MessageChannel(id, conn)
         new(id, conn, CONN_STATE_CLOSED, true,
             Channel{TAMQPGenericFrame}(CONN_MAX_QUEUED), nothing, Dict{Tuple,Tuple{Function,Any}}(),
             Message[], Channel{Union{Message, Nothing}}(1), Dict{String,MessageConsumer}(),
-            nothing)
+            Dict{String,Channel{Message}}(), ReentrantLock(), nothing)
     end
 end
 
@@ -422,7 +427,7 @@ end
 clear_handlers(c::MessageChannel) = (empty!(c.callbacks); nothing)
 function handle(c::MessageChannel, classname::Symbol, methodname::Symbol, cb=nothing, ctx=nothing)
     cbkey = method_key(classname, methodname)
-    if cb == nothing
+    if cb === nothing
         delete!(c.callbacks, cbkey)
     else
         c.callbacks[cbkey] = (cb, ctx)
@@ -431,7 +436,7 @@ function handle(c::MessageChannel, classname::Symbol, methodname::Symbol, cb=not
 end
 function handle(c::MessageChannel, frame_type::Integer, cb=nothing, ctx=nothing)
     cbkey = frame_key(frame_type)
-    if cb == nothing
+    if cb === nothing
         delete!(c.callbacks, cbkey)
     else
         c.callbacks[cbkey] = (cb, ctx)
@@ -779,15 +784,28 @@ nowait: do not send a reply method
 """
 function basic_consume(chan::MessageChannel, queue::String, consumer_fn::Function; consumer_tag::String="", no_local::Bool=false, no_ack::Bool=false,
     exclusive::Bool=false, nowait::Bool=false, arguments::Dict{String,Any}=Dict{String,Any}(), timeout::Int=DEFAULT_TIMEOUT, buffer_sz::Int=typemax(Int))
+
+    # register the consumer and get the consumer_tag
     result = _wait_resp(chan, (true, ""), nowait, on_basic_consume_ok, :Basic, :ConsumeOk, (false, ""), timeout) do
         send_basic_consume(chan, queue, consumer_tag, no_local, no_ack, exclusive, nowait, arguments)
     end
 
-    # setup a message consumer
+    # start the message consumer
     if result[1]
         consumer_tag = result[2]
-        chan.consumers[consumer_tag] = MessageConsumer(chan.id, consumer_tag, consumer_fn; buffer_size=buffer_sz)
+
+        # set up message buffer beforehand to store messages that the consumer may receive while we are still setting things up,
+        # or get the buffer that was set up already because we received messages
+        lock(chan.lck) do
+            consumer_buffer = get!(chan.pending_msgs, consumer_tag) do
+                Channel{Message}(buffer_sz)
+            end
+            consumer_buffer.sz_max = buffer_sz
+            chan.consumers[consumer_tag] = MessageConsumer(chan.id, consumer_tag, consumer_fn; buffer=consumer_buffer)
+            delete!(chan.pending_msgs, consumer_tag)
+        end
     end
+
     result
 end
 
@@ -1177,15 +1195,15 @@ function on_basic_get_empty_or_ok(chan::MessageChannel, m::TAMQPMethodFrame, ctx
 end
 
 function on_channel_message_in(chan::MessageChannel, m::TAMQPContentHeaderFrame, ctx)
-    msg = chan.partial_msgs[1]
+    msg = last(chan.partial_msgs)
     msg.properties = m.hdrpayload.proplist
     msg.data = Vector{UInt8}(undef, m.hdrpayload.bodysize)
     msg.filled = 0
     nothing
 end
 
 function on_channel_message_in(chan::MessageChannel, m::TAMQPContentBodyFrame, ctx)
-    msg = chan.partial_msgs[1]
+    msg = last(chan.partial_msgs)
     data = m.payload.data
     startpos = msg.filled + 1
     endpos = min(length(msg.data), msg.filled + length(data))
@@ -1195,11 +1213,16 @@ function on_channel_message_in(chan::MessageChannel, m::TAMQPContentBodyFrame, c
     if msg.filled >= length(msg.data)
         # got all data for msg
         if isempty(msg.consumer_tag)
-            put!(chan.chan_get, popfirst!(chan.partial_msgs))
-        elseif msg.consumer_tag in keys(chan.consumers)
-            put!(chan.consumers[msg.consumer_tag].recvq, popfirst!(chan.partial_msgs))
+            put!(chan.chan_get, pop!(chan.partial_msgs))
         else
-            @debug("discarding message, no consumer with tag", tag=msg.consumer_tag)
+            lock(chan.lck) do
+                if msg.consumer_tag in keys(chan.consumers)
+                    put!(chan.consumers[msg.consumer_tag].recvq, pop!(chan.partial_msgs))
+                else
+                    put!(get!(()->Channel{Message}(typemax(Int)), chan.pending_msgs, msg.consumer_tag), msg)
+                    @debug("holding message, no consumer yet with tag", tag=msg.consumer_tag)
+                end
+            end
         end
     end
 
@@ -1210,6 +1233,4 @@ on_confirm_select_ok(chan::MessageChannel, m::TAMQPMethodFrame, ctx) = _on_ack(c
 
 # ----------------------------------------
 # send and recv for methods end
-# ----------------------------------------
-
-
+# ----------------------------------------

test/test_rpc.jl

@@ -3,18 +3,17 @@ module AMQPTestRPC
 using AMQPClient, Test, Random
 
 const JULIA_HOME = Sys.BINDIR
-
 const QUEUE_RPC = "queue_rpc"
-
-testlog(msg) = println(msg)
-reply_queue_id = 0
-server_id = 0
-
 const NRPC_MSGS = 100
 const NRPC_CLNTS = 4
 const NRPC_SRVRS = 4
+const server_lck = Ref(ReentrantLock())
+const queue_declared = Ref(false)
+const servers_done = Channel{Int}(NRPC_SRVRS)
+
+testlog(msg) = println(msg)
 
-function test_rpc_client(;virtualhost="/", host="localhost", port=AMQPClient.AMQP_DEFAULT_PORT, auth_params=AMQPClient.DEFAULT_AUTH_PARAMS)
+function test_rpc_client(reply_queue_id; virtualhost="/", host="localhost", port=AMQPClient.AMQP_DEFAULT_PORT, auth_params=AMQPClient.DEFAULT_AUTH_PARAMS)
     # open a connection
     testlog("client opening connection...")
     conn = connection(;virtualhost=virtualhost, host=host, port=port, auth_params=auth_params)
@@ -24,8 +23,6 @@ function test_rpc_client(;virtualhost="/", host="localhost", port=AMQPClient.AMQ
     chan1 = channel(conn, AMQPClient.UNUSED_CHANNEL, true)
 
     # create a reply queue for a client
-    global reply_queue_id
-    reply_queue_id += 1
     queue_name = QUEUE_RPC * "_" * string(reply_queue_id) * "_" * string(getpid())
     testlog("client creating queue " * queue_name * "...")
     success, queue_name, message_count, consumer_count = queue_declare(chan1, queue_name; exclusive=true)
@@ -83,11 +80,7 @@ function test_rpc_client(;virtualhost="/", host="localhost", port=AMQPClient.AMQ
     testlog("client done.")
 end
 
-function test_rpc_server(;virtualhost="/", host="localhost", port=AMQPClient.AMQP_DEFAULT_PORT, auth_params=AMQPClient.DEFAULT_AUTH_PARAMS)
-    global server_id
-    server_id += 1
-    my_server_id = server_id
-
+function test_rpc_server(my_server_id; virtualhost="/", host="localhost", port=AMQPClient.AMQP_DEFAULT_PORT, auth_params=AMQPClient.DEFAULT_AUTH_PARAMS)
     # open a connection
     testlog("server $my_server_id opening connection...")
     conn = connection(;virtualhost=virtualhost, host=host, port=port, auth_params=auth_params)
@@ -97,10 +90,15 @@ function test_rpc_server(;virtualhost="/", host="localhost", port=AMQPClient.AMQ
     chan1 = channel(conn, AMQPClient.UNUSED_CHANNEL, true)
 
     # create queues (no need to bind if we are using the default exchange)
-    testlog("server $my_server_id creating queues...")
-    # this is the callback queue
-    success, message_count, consumer_count = queue_declare(chan1, QUEUE_RPC)
-    @test success
+    lock(server_lck[]) do
+        if !(queue_declared[])
+            testlog("server $my_server_id creating queues...")
+            # this is the callback queue
+            success, message_count, consumer_count = queue_declare(chan1, QUEUE_RPC)
+            @test success
+            queue_declared[] = true
+        end
+    end
 
     # test RPC
     testlog("server $my_server_id testing rpc...")
@@ -131,16 +129,24 @@ function test_rpc_server(;virtualhost="/", host="localhost", port=AMQPClient.AMQ
     end
 
     testlog("server $my_server_id closing down...")
-    success, message_count = queue_purge(chan1, QUEUE_RPC)
-    @test success
-    @test message_count == 0
-
     @test basic_cancel(chan1, consumer_tag)
-
-    success, message_count = queue_delete(chan1, QUEUE_RPC)
-    @test success
-    @test message_count == 0
-    testlog("server $my_server_id deleted rpc queue")
+    testlog("server $my_server_id cancelled consumer...")
+
+    lock(server_lck[]) do
+        take!(servers_done)
+        # the last server to finish will purge and delete the queue
+        if length(servers_done.data) == 0
+            success, message_count = queue_purge(chan1, QUEUE_RPC)
+            @test success
+            @test message_count == 0
+            testlog("server $my_server_id purged queue...")
+
+            success, message_count = queue_delete(chan1, QUEUE_RPC)
+            @test success
+            @test message_count == 0
+            testlog("server $my_server_id deleted rpc queue")
+        end
+    end
 
     # close channels and connection
     close(chan1)
@@ -157,29 +163,23 @@ end
 
 function runtests()
     testlog("testing multiple client server rpc")
-    clients = Vector{Task}()
-    servers = Vector{Task}()
 
     for idx in 1:NRPC_SRVRS
-        push!(servers, @async test_rpc_server())
+        put!(servers_done, idx)
     end
 
-    for idx in 1:NRPC_CLNTS
-        push!(clients, @async test_rpc_client())
-    end
-
-    tasks_active = NRPC_CLNTS + NRPC_SRVRS
-    while tasks_active > 0
-        tasks_active = 0
+    @sync begin
         for idx in 1:NRPC_SRVRS
-            istaskdone(servers[idx]) || (tasks_active += 1)
+            @async test_rpc_server(idx)
         end
+
         for idx in 1:NRPC_CLNTS
-            istaskdone(clients[idx]) || (tasks_active += 1)
+            @async test_rpc_client(idx)
         end
-        sleep(5)
     end
+
     testlog("done")
 end
+
 end # module AMQPTestRPC