redstruct 0.1.7 → 0.2.0

data/lib/redstruct/list.rb

@@ -0,0 +1,315 @@
+# frozen_string_literal: true
+
+require 'redstruct/struct'
+require 'redstruct/utils/scriptable'
+require 'redstruct/utils/iterable'
+
+module Redstruct
+  # Class to manipulate redis lists, modeled after Ruby's Array class.
+  # TODO: Add maximum instance variable and modify all methods (where applicable)
+  # to take it into account.
+  class List < Redstruct::Struct
+    include Redstruct::Utils::Scriptable
+    include Redstruct::Utils::Iterable
+
+    # Clears the set by simply removing the key from the DB
+    # @see Redstruct::Struct#clear
+    def clear
+      delete
+    end
+
+    # Checks if the set is empty by checking if the key actually exists on the underlying redis db
+    # @see Redstruct::Struct#exists?
+    # @return [Boolean] true if it is empty, false otherwise
+    def empty?
+      return !exists?
+    end
+
+    # Returns the item located at index
+    # @param [Integer] index the item located at index
+    # @return [String, nil] nil if no item at index, otherwise the value
+    def [](index)
+      return self.connection.lindex(@key, index.to_i)
+    end
+
+    # Sets or updates the value for item at index
+    # @param [Integer] index the index
+    # @param [#to_s] value the new value
+    # @raise Redis::BaseError when index is out of range
+    # @return [Boolean] true if set, false otherwise
+    def []=(index, value)
+      return coerce_bool(set_script(keys: @key, argv: [index.to_i, value]))
+    end
+
+    # Inserts the given value at the given zero-based index.
+    # TODO: Support multiple insertion like Array#insert? The biggest issue
+    # here is that concatenating lists in Lua is O(n), so on very large lists,
+    # this operation would become slow. There are Redis Modules which implement
+    # splice operations (so a O(1) list split/merge), but there's no way to
+    # guarantee if the module will be present. Perhaps provide optional support
+    # if the module is detected?
+    # @param [#to_s] value the value to insert
+    # @param [#to_i] index the index at which to insert the value
+    def insert(value, index)
+      result = case index
+      when 0 then prepend(value)
+      when -1 then append(value)
+      else
+        index += 1 if index.negative?
+        insert_script(keys: @key, argv: [value, index])
+      end
+
+      return coerce_bool(result)
+    end
+
+    # Appends the given items (from the right) to the list
+    # @param [Array<#to_s>] items the items to append
+    # @param [Integer] max optional; if given, appends the items and trims down the list to max afterwards
+    # @param [Boolean] exists optional; if true, only appends iff the list already exists (i.e. is not empty)
+    # @return [Integer] the number of items appended to the list
+    def append(*items, max: 0, exists: false)
+      max = max.to_i
+      results = if max.positive? || exists
+        push_and_trim_script(keys: @key, argv: [max - 1, false, exists] + items)
+      else
+        self.connection.rpush(@key, items)
+      end
+
+      return results
+    end
+    alias push append
+
+    # Pushes the given element onto the list. As << is a binary operator, it can
+    # only take one argument in. It's more of a convenience method.
+    # @param [#to_s] item the item to append to the list
+    # @return [Integer] 1 if appended, 0 otherwise
+    def <<(item)
+      return append(item)
+    end
+
+    # Prepends the given items (from the right) to the list
+    # @param [Array<#to_s>] items the items to prepend
+    # @param [Integer] max optional; if given, prepends the items and trims down the list to max afterwards
+    # @param [Boolean] exists optional; if true, only prepends iff the list already exists (i.e. is not empty)
+    # @return [Integer] the number of items prepended to the list
+    def prepend(*items, max: nil, exists: false)
+      max = max.to_i
+
+      # redis literally prepends each element one at a time, so 1 2 will end up 2 1
+      # to keep behaviour in sync with Array#unshift we preemptively reverse the list
+      items = items.reverse
+
+      results = if max.positive? || exists
+        push_and_trim_script(keys: @key, argv: [max - 1, true, exists] + items)
+      else
+        self.connection.lpush(@key, items)
+      end
+
+      return results
+    end
+    alias unshift prepend
+
+    # Pops an item from the list, optionally blocking to wait until the list is non-empty
+    # @param [Integer] timeout the amount of time to wait in seconds; if 0, waits indefinitely
+    # @return [nil, String] nil if the list was empty, otherwise the item
+    def pop(size = 1, timeout: nil)
+      raise ArgumentError, 'size must be positive' unless size.positive?
+
+      if timeout.nil?
+        return self.connection.rpop(@key) if size == 1
+        return shift_pop_script(keys: @key, argv: [-size, -1, 1])
+      else
+        raise ArgumentError, 'timeout is only supported if size == 1' unless size == 1
+        return self.connection.brpop(@key, timeout: timeout)&.last
+      end
+    end
+
+    # Shifts an item from the list, optionally blocking to wait until the list is non-empty
+    # @param [Integer] timeout the amount of time to wait in seconds; if 0, waits indefinitely
+    # @return [nil, String] nil if the list was empty, otherwise the item
+    def shift(size = 1, timeout: nil)
+      raise ArgumentError, 'size must be positive' unless size.positive?
+
+      if timeout.nil?
+        return self.connection.lpop(@key) if size == 1
+        return shift_pop_script(keys: @key, argv: [0, size - 1, 0])
+      else
+        raise ArgumentError, 'timeout is only supported if size == 1' unless size == 1
+        return self.connection.blpop(@key, timeout: timeout)&.last
+      end
+    end
+
+    # Pops an element from this list and shifts it onto the given list.
+    # @param [Redstruct::List] list the list to shift the element onto
+    # @param [#to_i] timeout optional timeout to wait for in seconds
+    # @return [String] the element that was popped from the list and pushed onto the other
+    def popshift(list, timeout: nil)
+      raise ArgumentError, 'list must respond to #key' unless list.respond_to?(:key)
+
+      if timeout.nil?
+        return self.connection.rpoplpush(@key, list.key)
+      else
+        return self.connection.brpoplpush(@key, list.key, timeout: timeout)
+      end
+    end
+
+    # Removes the given item from the list.
+    # @param [Integer] count count > 0: Remove items equal to value moving from head to tail.
+    #                        count < 0: Remove items equal to value moving from tail to head.
+    #                        count = 0: Remove all items equal to value.
+    # @return [Integer] the number of items removed
+    def remove(value, count: 1)
+      self.connection.lrem(@key, count.to_i, value)
+    end
+
+    # Checks how many items are in the list.
+    # @return [Integer] the number of items in the list
+    def size
+      return self.connection.llen(@key)
+    end
+
+    # Returns a slice of the list starting at start (inclusively), up to length (inclusively)
+    # @example
+    #   pry> list.slice(start: 1, length: 10) #=> Array<...> # array with 11 items
+    # @param [Integer] start the starting index for the slice; if start is larger than the end of the list, an empty list is returned
+    # @param [Integer] length the length of the slice (inclusively); if -1, returns everything
+    # @return [Array<String>] the requested slice, or an empty list
+    def slice(start: 0, length: -1)
+      length = length.to_i
+      end_index = length.positive? ? start + length - 1 : length
+
+      return self.connection.lrange(@key, start.to_i, end_index)
+    end
+
+    # Loads all items into memory and returns an array.
+    # NOTE: if the list is expected to be large, use to_enum
+    # @return [Array<String>] the items in the list
+    def to_a
+      return slice(start: 0, length: -1)
+    end
+
+    # Since the list can be modified in between loops, this does not guarantee
+    # completion of the operation, nor that every single element of the list
+    # will be visited once; rather, it guarantees that it loops until no more
+    # elements are returned, using an incrementing offset.
+    # This means that should elements be removed in the meantime, they will
+    # not be seen, and others might be skipped as a result of this.
+    # If elements are added, it is however not an issue (although keep in mind
+    # that if elements are added faster than consumed, this can loop forever)
+    # @return [Enumerator] base enumerator to iterate of the list elements
+    def to_enum(match: '*', count: 10)
+      pattern = Regexp.compile("^#{Regexp.escape(match).gsub('\*', '.*')}$")
+
+      return Enumerator.new do |yielder|
+        offset = 0
+        loop do
+          items = slice(start: offset, length: offset + count)
+
+          offset += items.size
+          matched = items.select { |item| item =~ pattern }
+          yielder << matched unless matched.empty?
+
+          raise StopIteration if items.size < count
+        end
+      end
+    end
+
+    # @!group Lua Scripts
+
+    # Appends or prepends (argv[1]) a number of items (argv[2]) to a list (keys[1]),
+    # then trims it out to size (argv[3])
+    # @param [Array<#to_s>] keys first key should be the key to the list to prepend to and resize
+    # @param [Array<Integer, Integer, Integer, Array<#to_s>>] argv the maximum size of the list; if 1, will lpush, otherwise rpush;
+    #                                                              if 1, will push only if the list exists; the list of items to prepend
+    # @return [Integer] the length of the list after the operation
+    defscript :push_and_trim_script, <<~LUA
+      local max = tonumber(table.remove(ARGV, 1))
+      local prepend = tonumber(table.remove(ARGV, 1)) == 1
+      local exists = tonumber(table.remove(ARGV, 1)) == 1
+      local push = prepend and 'lpush' or 'rpush'
+      local size = 0
+
+      if exists then
+        if redis.call('exists', KEYS[1]) == 0 then
+          return nil
+        end
+      end
+
+      size = redis.call(push, KEYS[1], unpack(ARGV))
+      if max > 0 and size > max then
+        redis.call('ltrim', KEYS[1], 0, max)
+        size = max + 1
+      end
+
+      return size
+    LUA
+    protected :push_and_trim_script
+
+    # Removes N elements from the list (either from the head or the tail) and trims
+    # the list down to size (either from the head or the tail).
+    # @param [Array<#to_s>] keys first key should be the key to the list to prepend to and resize
+    # @param [Array<Integer, Integer>] argv the start of the slice; the end index of the slice
+    # @return [Integer] the sliced list
+    defscript :shift_pop_script, <<~LUA
+      local range_start = tonumber(ARGV[1])
+      local range_end = tonumber(ARGV[2])
+      local direction = tonumber(ARGV[3])
+      local list = redis.call('lrange', KEYS[1], range_start, range_end)
+
+      if #list > 0 then
+        if direction == 0 then
+          redis.call('ltrim', KEYS[1], range_end + 1, -1)
+        else
+          redis.call('ltrim', KEYS[1], 0, range_start - 1)
+        end
+      end
+
+      return list
+    LUA
+    protected :shift_pop_script
+
+    # Inserts the given element at the given index. Can raise out of bound error
+    # @param [Array<#to_s>] keys first key is the list key
+    # @param [Array<#to_s, #to_i>] argv the value to insert; the index at which to insert it
+    # @return [Boolean] true if inserted, false otherwise
+    defscript :insert_script, <<~LUA
+      local value = ARGV[1]
+      local index = tonumber(ARGV[2])
+      local pivot = redis.call('lindex', KEYS[1], index - 1)
+
+      if pivot ~= nil then
+        return redis.call('linsert', KEYS[1], 'AFTER', pivot, value)
+      end
+
+      return false
+    LUA
+    protected :insert_script
+
+    # Sets the element in much the same way a ruby array would, by padding
+    # with empty strings (redis equivalent of nil) when the index is out of range
+    # @param [Array<#to_s>] keys first key is the list key
+    # @param [Array<#to_s, #to_i>] argv the index to set; the value to set
+    # @return [Boolean] true if inserted, false otherwise
+    defscript :set_script, <<~LUA
+      local index = tonumber(ARGV[1])
+      local value = ARGV[2]
+      local max = redis.call('llen', KEYS[1]) - 1
+
+      if max < index then
+        local upto = index - max - 1
+        local items = {}
+        for i = index - max - 1, 1, -1 do
+          items[#items+1] = ''
+        end
+
+        items[#items+1] = value
+        return redis.call('rpush', KEYS[1], unpack(items))
+      end
+
+      return redis.call('lset', KEYS[1], index, value)
+    LUA
+    protected :set_script
+
+    # @!endgroup
+  end
+end

data/lib/redstruct/lock.rb

@@ -0,0 +1,183 @@
+# frozen_string_literal: true
+
+require 'securerandom'
+require 'redstruct/factory/object'
+require 'redstruct/utils/scriptable'
+require 'redstruct/utils/coercion'
+
+module Redstruct
+  # Implementation of a simple binary lock (locked/not locked), with option to block and wait for the lock.
+  # Uses two redis structures: a string for the lease, and a list for blocking operations.
+  class Lock < Redstruct::Factory::Object
+    include Redstruct::Utils::Scriptable
+    include Redstruct::Utils::Coercion
+
+    # The default expiry on the underlying redis keys, in seconds; can be between 0 and 1 as a float for milliseconds
+    DEFAULT_EXPIRY = 1
+
+    # The default timeout when blocking, in seconds
+    DEFAULT_TIMEOUT = nil
+
+    # @return [String] the resource name (or ID of the lock)
+    attr_reader :resource
+
+    # @return [String] the current token
+    attr_reader :token
+
+    # @return [Float, Integer] the expiry of the underlying redis structure in seconds
+    attr_reader :expiry
+
+    # @return [Integer] if greater than 0, will block until timeout is reached or the lock is acquired
+    attr_reader :timeout
+
+    # @param [String] resource the name of the resource to be locked (or ID)
+    # @param [Integer] expiry in seconds; to prevent infinite locking, you should pass a minimum expiry; you can pass 0 if you want to control it yourself
+    # @param [Integer] timeout in seconds; if > 0, will block when trying to obtain the lock; if 0, blocks indefinitely; if nil, does not block
+    def initialize(resource, expiry: DEFAULT_EXPIRY, timeout: DEFAULT_TIMEOUT, **options)
+      super(**options)
+
+      @resource = resource
+      @token = nil
+      @expiry = expiry
+      @timeout = case timeout
+      when nil then nil
+      when Float::INFINITY then 0
+      else
+        timeout.to_i
+      end
+
+      factory = @factory.factory(@resource)
+      @lease = factory.string('lease')
+      @tokens = factory.list('tokens')
+    end
+
+    # Executes the given block if the lock can be acquired
+    # @yield Block to be executed if the lock is acquired
+    def locked
+      Thread.handle_interrupt(Exception => :never) do
+        begin
+          if acquire
+            Thread.handle_interrupt(Exception => :immediate) do
+              yield
+            end
+          end
+        ensure
+          release
+        end
+      end
+    end
+
+    # Whether or not the lock will block when attempting to acquire it
+    # @return [Boolean]
+    def blocking?
+      return !@timeout.nil?
+    end
+
+    # Attempts to acquire the lock. First attempts to grab the lease (a redis string).
+    # If the current token is already the lease token, the lock is considered acquired.
+    # If there is no current lease, then sets it to the current token.
+    # If there is a current lease that is not the current token, then:
+    #   1) If this not a blocking lock (see Lock#blocking?), return false
+    #   2) If this is a blocking lock, block and wait for the next token to be pushed on the tokens list
+    #   3) If a token was pushed, set it as our token and refresh the expiry
+    # @return [Boolean] True if acquired, false otherwise
+    def acquire
+      acquired = false
+
+      token = non_blocking_acquire
+      token = blocking_acquire if token.nil? && blocking?
+
+      unless token.nil?
+        @lease.expire(@expiry)
+        @token = token
+        acquired = true
+      end
+
+      return acquired
+    end
+
+    # Releases the lock only if the current token is the value of the lease.
+    # If the lock is a blocking lock (see Lock#blocking?), push the next token on the tokens list.
+    # @return [Boolean] True if released, false otherwise
+    def release
+      return false if @token.nil?
+
+      keys = [@lease.key, @tokens.key]
+      argv = [@token, generate_token, (@expiry.to_f * 1000).floor]
+
+      released = release_script(keys: keys, argv: argv)
+      @token = nil
+
+      return coerce_bool(released)
+    end
+
+    private
+
+    def non_blocking_acquire
+      keys = [@lease.key, @tokens.key]
+      argv = [generate_token]
+
+      return acquire_script(keys: keys, argv: argv)
+    end
+
+    def blocking_acquire
+      return @tokens.pop(timeout: @timeout)
+    end
+
+    # The acquire script attempts to set the lease (keys[1]) to the given token (argv[1]), only
+    # if it wasn't already set. It then compares to check if the value of the lease is that of the token,
+    # and if so refreshes the expiry (argv[2]) time of the lease.
+    # @param [Array<(::String)>] keys The lease key specifying who owns the mutex at the moment
+    # @param [Array<(::String, Fixnum)>] argv the current token
+    # @return [::String] Returns the token if acquired, nil otherwise.
+    defscript :acquire_script, <<~LUA
+      local token = ARGV[1]
+      local lease = redis.call('get', KEYS[1])
+
+      if not lease then
+        redis.call('set', KEYS[1], token)
+      elseif token ~= lease then
+        token = redis.call('lpop', KEYS[2])
+        if not token or token ~= lease then
+          return false
+        end
+      end
+
+      return token
+    LUA
+
+    # The release script compares the given token (argv[1]) with the lease value (keys[1]); if they are the same,
+    # then a new token (argv[2]) is set as the lease, and pushed on the tokens (keys[2]) list
+    # for the next acquire request.
+    # @param [Array<(::String, ::String)>] keys the lease key; the tokens list key
+    # @param [Array<(::String, ::String, Fixnum)>] argv the current token; the next token to push; the expiry time of both keys
+    # @return [Fixnum] 1 if released, 0 otherwise
+    defscript :release_script, <<~LUA
+      local currentToken = ARGV[1]
+      local nextToken = ARGV[2]
+      local expiry = tonumber(ARGV[3])
+
+      if redis.call('get', KEYS[1]) == currentToken then
+        redis.call('set', KEYS[1], nextToken)
+        redis.call('lpush', KEYS[2], nextToken)
+
+        if expiry > 0 then
+          redis.call('pexpire', KEYS[1], expiry)
+          redis.call('pexpire', KEYS[2], expiry)
+        end
+
+        return true
+      end
+
+      return false
+    LUA
+
+    def generate_token
+      return SecureRandom.uuid
+    end
+
+    def inspectable_attributes
+      super.merge(expiry: @expiry, blocking: blocking?)
+    end
+  end
+end

data/lib/redstruct/script.rb

@@ -0,0 +1,104 @@
+# frozen_string_literal: true
+
+require 'digest'
+require 'redis'
+require 'redstruct/connection_proxy'
+require 'redstruct/utils/inspectable'
+
+module Redstruct
+  # Utility class to interact with Lua scripts on Redis.
+  # It is recommended you flush your script cache on the redis server every once in a while to remove scripts that
+  # are not used anymore.
+  class Script
+    # Redis returns an error starting with NOSCRIPT when we try to evaluate am unknown script using its sha1.
+    ERROR_MESSAGE_PREFIX = 'NOSCRIPT'
+
+    # @return [String] the Lua script to evaluate
+    attr_reader :script
+
+    # @return [Redstruct::ConnectionProxy] the connection used for script commands
+    attr_reader :connection
+
+    # @param [String] script the lua source code for the script
+    # @param [Redstruct::ConnectionProxy] connection connection to use for script commands
+    # @param [String] sha1 the sha1 representation of the script; optional
+    def initialize(script:, connection:, sha1: nil)
+      self.script = script
+      @connection = connection
+      @sha1 = sha1 unless sha1.nil?
+
+      raise ArgumentError, 'requires a connection proxy' unless @connection.is_a?(Redstruct::ConnectionProxy)
+    end
+
+    # Duplicates and freezes the given script, and reinitializes the sha1 (which later gets lazily computed)
+    # @param [String] script the lua source code
+    def script=(script)
+      script = script&.strip
+      raise ArgumentError, 'No source script given' if script.empty?
+
+      @sha1 = nil
+      @script = script.dup.freeze
+    end
+
+    # Returns the sha1 representation of the source code at `script`
+    # When running a lua script, redis will compile it once and cache the bytecode, using the sha1 of the source code
+    # as the cache key.
+    # @return [String] sha1 representation of `script`
+    def sha1
+      return @sha1 ||= begin
+        Digest::SHA1.hexdigest(@script)
+      end
+    end
+
+    # Checks if the script was already loaded for the given redis db using #sha1
+    # @return [Boolean] true if the script was already loaded, false otherwise
+    def exists?
+      return @connection.script(:exists, self.sha1)
+    end
+
+    # Loads the given script to redis (i.e. sends the source, which gets compiled and saved) and saves the returned sha1
+    # @return [String] the new sha1
+    def load
+      @sha1 = @connection.script(:load, @script)
+      return @sha1
+    end
+
+    # Evaluates the script using the given keys and argv arrays, and returns the unparsed result. Caller is in charge
+    # of interpreting the result.
+    # NOTE: To minimize the number of redis commands, this always first assumes that the script was already loaded using
+    # its sha1 representation, and tells redis to execute the script cached by `sha1`. If it receives as error that the
+    # script does not exist, only then will it send the source to be executed. So in the worst case you get 2 redis
+    # commands, but in the average case you get 1, and it's much faster as redis does not have to reparse the script,
+    # and we don't need to send the lua source every time.
+    # @param [Array<String>] keys the KEYS array as described in the Redis doc for eval
+    # @param [Array<String>] argv the ARGV array as described in the Redis doc for eval
+    # @return [nil, Boolean, String, Numeric] returns whatever redis returns
+    def eval(keys: [], argv: [])
+      keys = [keys] unless keys.is_a?(Array)
+      argv = [argv] unless argv.is_a?(Array)
+      argv = normalize(argv)
+
+      @connection.evalsha(self.sha1, keys, argv)
+    rescue Redis::CommandError => err
+      raise unless err.message.start_with?(ERROR_MESSAGE_PREFIX)
+      @connection.eval(@script, keys, argv)
+    end
+
+    private
+
+    def inspectable_attributes # :nodoc:
+      return super.merge(sha1: self.sha1, script: @script.slice(0, 20))
+    end
+
+    def normalize(values)
+      values.map do |value|
+        case value
+        when true then 1
+        when false then 0
+        else
+          value.to_s
+        end
+      end
+    end
+  end
+end