Add support for stream_id

Author: drolando

lib/cassandra/cql.lua

@@ -945,10 +945,15 @@ do
 
     header:write_byte(flags)
 
+    local stream_id = 0
+    if self.opts and self.opts.stream_id then
+      stream_id = self.opts.stream_id
+    end
+
     if version < 3 then
-      header:write_byte(0)       -- stream_id
+      header:write_byte(stream_id)
     else
-      header:write_short(0)      -- stream_id
+      header:write_short(stream_id)
     end
 
     header:write_byte(self.op_code)

lib/cassandra/init.lua

@@ -150,28 +150,34 @@ function _Host:send(request)
   local sent, err = self.sock:send(frame)
   if not sent then return nil, err end
 
-  -- receive frame version byte
-  local v_byte, err = self.sock:receive(1)
-  if not v_byte then return nil, err end
+  while true do
+    -- receive frame version byte
+    local v_byte, err = self.sock:receive(1)
+    if not v_byte then return nil, err end
 
-  -- -1 because of the v_byte we just read
-  local version, n_bytes = cql.frame_reader.version(v_byte)
+    -- -1 because of the v_byte we just read
+    local version, n_bytes = cql.frame_reader.version(v_byte)
 
-  -- receive frame header
-  local header_bytes, err = self.sock:receive(n_bytes)
-  if not header_bytes then return nil, err end
+    -- receive frame header
+    local header_bytes, err = self.sock:receive(n_bytes)
+    if not header_bytes then return nil, err end
 
-  local header = cql.frame_reader.read_header(version, header_bytes)
+    local header = cql.frame_reader.read_header(version, header_bytes)
 
-  -- receive frame body
-  local body_bytes
-  if header.body_length > 0 then
-    body_bytes, err = self.sock:receive(header.body_length)
-    if not body_bytes then return nil, err end
-  end
+    -- receive frame body
+    local body_bytes
+    if header.body_length > 0 then
+      body_bytes, err = self.sock:receive(header.body_length)
+      if not body_bytes then return nil, err end
+    end
 
-  -- res, err, cql_err_code
-  return cql.frame_reader.read_body(header, body_bytes)
+    -- If stream_id was set in request.opts, only return a response
+    -- with a matching stream_id and drop everything else
+    if not request.opts or not request.opts.stream_id or request.opts.stream_id == header.stream_id then
+      -- res, err, cql_err_code
+      return cql.frame_reader.read_body(header, body_bytes)
+    end
+  end
 end
 
 local function send_startup(self)

lib/resty/cassandra/cluster.lua

@@ -387,6 +387,7 @@ function _Cluster.new(opts)
                 or require('resty.cassandra.policies.reconnection.exp').new(1000, 60000),
     retry_policy = opts.retry_policy
                 or require('resty.cassandra.policies.retry.simple').new(3),
+    stream_ids = nil,
   }, _Cluster)
 end
 
@@ -550,6 +551,15 @@ function _Cluster:refresh()
   -- initiate the load balancing policy
   self.lb_policy:init(peers)
 
+  if self.stream_ids == nil then
+    -- Initialize the list of available seed ids (only once)
+    self.stream_ids = {}
+    local max_id = protocol_version < 3 and 2^7-1 or 2^15-1
+    for i=1,max_id do
+      self.stream_ids[i] = i
+    end
+  end
+
   -- cluster is ready to be queried
   self.init = true
 
@@ -748,8 +758,45 @@ local function handle_error(self, err, cql_code, coordinator, request)
   return nil, err, cql_code
 end
 
+local function get_stream_id(self)
+  local stream_id = 0
+
+  local lock = resty_lock:new(self.dict_name, self.lock_opts)
+  local elapsed, err = lock:lock('stream_id')
+  if not elapsed then return stream_id, 'failed to acquire stream_id lock: '..err end
+
+  if table.getn(self.stream_ids) > 0 then
+    stream_id = table.remove(self.stream_ids)
+  else
+    err = 'Too many inflight requests. No new stream ids available.'
+  end
+
+  local ok, err = lock:unlock()
+  if not ok then return stream_id, 'failed to unlock: '..err end
+
+  return stream_id, nil
+end
+
+local function release_stream_id(self, stream_id)
+  local lock = resty_lock:new(self.dict_name, self.lock_opts)
+  local elapsed, err = lock:lock('stream_id')
+  if not elapsed then return stream_id, 'failed to acquire stream_id lock: '..err end
+
+  table.insert(self.stream_ids, 1, stream_id)
+
+  local ok, err = lock:unlock()
+  if not ok then return stream_id, 'failed to unlock: '..err end
+end
+
 send_request = function(self, coordinator, request)
+  local err
+  request.opts.stream_id, err = get_stream_id(self)
+  if err ~= nil and self.logging then
+    log(WARN, _log_prefix, err)
+  end
+
   local res, err, cql_code = coordinator:send(request)
+  release_stream_id(self, request.opts.stream_id)
   if not res then
     return handle_error(self, err, cql_code, coordinator, request)
   elseif res.warnings and self.logging then

spec/01-unit/03-cql_spec.lua

@@ -166,6 +166,7 @@ for protocol_version = 2, 3 do
 
   describe("CQL requests", function()
     local requests = cql.requests
+    local consistencies = cql.consistencies
 
     it("sanity", function()
       local r = requests.query.new("SELECT * FROM peers")
@@ -190,6 +191,14 @@ for protocol_version = 2, 3 do
         local frame3 = r:build_frame(protocol_version)
         assert.matches("SELECT key FROM local", frame3, nil, true)
       end)
+      it("sets the stream_id if provided", function()
+        local r = requests.query.new("SELECT * FROM local")
+        r.opts = {stream_id = 255, consistency = consistencies.one}
+        local frame = r:build_frame(protocol_version)
+
+        local header = cql.frame_reader.read_header(protocol_version, string.sub(frame, 2, -1))
+        assert.equal(255, header.stream_id)
+      end)
     end)
 
     describe("execute_prepared", function()