class Resque::DataStore::Workers

def acquire_pruning_dead_worker_lock(worker, expiry)
  @redis.set(redis_key_for_worker_pruning, worker.to_s, :ex => expiry, :nx => true)
end

def all_heartbeats
  @redis.hgetall(HEARTBEAT_KEY)
end

def get_worker_payload(worker_id)
  @redis.get("worker:#{worker_id}")
end

def heartbeat(worker)
  heartbeat = @redis.hget(HEARTBEAT_KEY, worker.to_s)
  heartbeat && Time.parse(heartbeat)
end

def heartbeat!(worker, time)
  @redis.hset(HEARTBEAT_KEY, worker.to_s, time.iso8601)
end

def initialize(redis)
  @redis = redis
end

def redis_key_for_worker(worker)
  "worker:#{worker}"
end

def redis_key_for_worker_pruning
  "pruning_dead_workers_in_progress"
end

def redis_key_for_worker_start_time(worker)
  "#{redis_key_for_worker(worker)}:started"
end

def register_worker(worker)
  @redis.pipelined do |piped|
    piped.sadd(:workers, [worker.id])
    worker_started(worker, redis: piped)
  end
end

def remove_heartbeat(worker)
  @redis.hdel(HEARTBEAT_KEY, worker.to_s)
end

def set_worker_payload(worker, data)
  @redis.set(redis_key_for_worker(worker), data)
end

def unregister_worker(worker, &block)
  @redis.pipelined do |piped|
    piped.srem(:workers, [worker.id])
    piped.del(redis_key_for_worker(worker))
    piped.del(redis_key_for_worker_start_time(worker))
    piped.hdel(HEARTBEAT_KEY, worker.to_s)
    block.call redis: piped
  end
end

def worker_done_working(worker, &block)
  @redis.pipelined do |piped|
    piped.del(redis_key_for_worker(worker))
    block.call redis: piped
  end
end

def worker_exists?(worker_id)
  @redis.sismember(:workers, worker_id)
end

def worker_ids
  Array(@redis.smembers(:workers))
end

def worker_start_time(worker)
  @redis.get(redis_key_for_worker_start_time(worker))
end

def worker_started(worker, redis: @redis)
  redis.set(redis_key_for_worker_start_time(worker), Time.now.to_s)
end

def workers_map(worker_ids)
  redis_keys = worker_ids.map { |id| "worker:#{id}" }
  @redis.mapped_mget(*redis_keys)
end