summaryrefslogtreecommitdiff
path: root/lib/custodian/queue.rb
blob: 9c33825ffe824346b4fac7df53eb0133ca636b96 (plain)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
#
# Attempt to load the Redis-library.
#
# Without this we cannot connect to our queue.
#
%w( redis ).each do |library|
  begin
    require library
  rescue LoadError
    puts("Failed to load the #{library} library - queue access will fail!")
  end
end



module Custodian


  #
  # An abstraction layer for our queue.
  #
  class QueueType

    #
    # Retrieve a job from the queue.
    #
    def fetch(_timeout)
      raise 'Subclasses must implement this method!'
    end


    #
    # Add a new job to the queue.
    #
    def add(_job_string)
      raise 'Subclasses must implement this method!'
    end


    #
    # Get the size of the queue
    #
    def size?
      raise 'Subclasses must implement this method!'
    end


    #
    # Empty the queue
    #
    def flush!
      raise 'Subclasses must implement this method!'
    end
  end




  #
  # This is a simple queue which uses Redis for storage.
  #
  class RedisQueueType < QueueType


    #
    # Connect to the server on localhost, unless QUEUE_ADDRESS is set.
    #
    def initialize
      host = ENV['QUEUE_ADDRESS'] || '127.0.0.1'
      @redis = Redis.new(:host => host)
    end


    #
    #  Fetch a job from the queue.
    #
    #  The timeout is used to specify the period we wait for a new job, and
    # we pause that same period between fetches.
    #
    def fetch(timeout = 1)
      job = nil

      loop do
        job = @redis.zrange('zset', '0', '0')

        if !job.empty?
          # We only have one entry in our array
          job = job[0]

          # Remove from the queue
          @redis.zrem('zset', job)

          return job
        else
          sleep(timeout)
        end
      end
    end


    #
    #  Add a new job to the queue - this can stall for the case where the
    # job is already pending.
    #
    def add(test)

        @redis.watch('zset')

        #
        # Count the number of times we attempt to add the test
        #
        attempts = 0
        added = false


        #
        # This is run in a loop, as we have to wait until both
        #
        # (a) the score is missing
        # (b) the zadd function succeeds
        #
        while (attempts < 40) do

          #
          # Only update if no score is set
          #
          if !@redis.zscore('zset', test)

            #
            # If MULTI returns nil, the transaction failed, so we need to try
            # again.
            #
            break unless @redis.multi do |r|
              @redis.zadd('zset', Time.now.to_f, test)
              added = true
            end.nil?
          end

          #
          # This could be tighter..
          #
          sleep 0.1

          #
          # Bump the count of attempts.
          #
          attempts += 1
        end

        #
        # Do we need to unwatch here?
        #
        @redis.unwatch

        #
        # Return the success/fail
        #
        return added
    end


    #
    #  How many jobs in the queue?
    #
    def size?
      @redis.scard('custodian_queue')
    end


    #
    #  Empty the queue, discarding all pending jobs.
    #
    def flush!
      @redis.del('custodian_queue')
    end

  end

end