// UpdateQueues should be called whenever a client changes what queues it's
// registered for. The new full list of registered queues should be passed in,
// this method will do a diff and figure it out what was removed
func UpdateQueues(client *clients.Client, queues []string) error {
respCh := make(chan error)
callCh <- func(s *state) {
oldQueues := s.clientQueues(client.ID)
removed := stringSliceSub(oldQueues, queues)
s.addClientQueues(client, queues)
s.removeClientQueues(client, removed)
select {
case updateNotifyCh <- struct{}{}:
default:
}
ts := time.Now().Unix()
pipe := make([]*db.PipePart, 0, len(removed)+len(queues))
for _, queueName := range removed {
consumersKey := db.ConsumersKey(queueName)
pipe = append(pipe, db.PP("ZREM", consumersKey, client.ID))
}
for _, queueName := range queues {
consumersKey := db.ConsumersKey(queueName)
pipe = append(pipe, db.PP("ZADD", consumersKey, ts, client.ID))
}
_, err := db.Inst.Pipe(pipe...)
respCh <- err
}
return <-respCh
}
func updateActiveConsumers() error {
log.L.Debug("updating active consumers")
ts := time.Now().Unix()
// A list of args to pass into ZADD for each consumer
consumersArgs := map[string][]interface{}{}
// Populate consumersArgs arguments to the ZADD commands we're going to
// need to perform
for _, queue := range registeredQueues() {
for _, client := range queueClients(queue) {
args, ok := consumersArgs[queue]
if !ok {
args = make([]interface{}, 0, 3)
args = append(args, db.ConsumersKey(queue))
consumersArgs[queue] = args
}
consumersArgs[queue] = append(args, ts, client.ID)
}
}
for _, args := range consumersArgs {
if err := db.Inst.Cmd("ZADD", args...).Err; err != nil {
return err
}
}
return nil
}
func TestStaleCleanup(t *T) {
queue := clients.RandQueueName()
client := clients.NewClient(clients.NewFakeClientConn())
err := UpdateQueues(client, []string{queue})
require.Nil(t, err)
// Make sure the queue has this clientId as a consumer
key := db.ConsumersKey(queue)
res := db.Inst.Cmd("ZRANK", key, client.ID)
assert.Equal(t, true, res.IsType(redis.Int), "res: %s", res)
// Remove all knowledge about this client from the consumer state
callCh <- func(s *state) {
s.removeClientQueues(client, []string{queue})
}
// Wait a little bit and try to remove stale consumers manually
time.Sleep(2 * time.Second)
err = removeStaleConsumers(1 * time.Second)
require.Nil(t, err)
// Make sure this client is no longer a consumer
res = db.Inst.Cmd("ZRANK", key, client.ID)
assert.Equal(t, true, res.IsType(redis.Nil), "key: %s clientId: %s res: %s", key, client.ID, res)
}
func TestUpdateQueues(t *T) {
queues := []string{
clients.RandQueueName(),
clients.RandQueueName(),
clients.RandQueueName(),
}
client := clients.NewClient(clients.NewFakeClientConn())
err := UpdateQueues(client, queues)
require.Nil(t, err)
// Make sure the client.Id appears in the consumers set for those queues
for i := range queues {
key := db.ConsumersKey(queues[i])
res := db.Inst.Cmd("ZRANK", key, client.ID)
assert.Equal(t, true, res.IsType(redis.Int), "res: %s", res)
}
err = UpdateQueues(client, queues[1:])
require.Nil(t, err)
// Make sure the first queue had this clientId removed from it
key := db.ConsumersKey(queues[0])
res := db.Inst.Cmd("ZRANK", key, client.ID)
assert.Equal(t, true, res.IsType(redis.Nil), "res: %s", res)
// Make sure the rest of the queues still have it
for i := range queues[1:] {
key := db.ConsumersKey(queues[1:][i])
res := db.Inst.Cmd("ZRANK", key, client.ID)
assert.Equal(t, true, res.IsType(redis.Int), "res: %s", res)
}
err = UpdateQueues(client, []string{})
require.Nil(t, err)
// Make sure the clientId appears nowhere
for i := range queues {
key := db.ConsumersKey(queues[i])
res := db.Inst.Cmd("ZRANK", key, client.ID)
assert.Equal(t, true, res.IsType(redis.Nil), "res: %s", res)
}
}
func removeStaleConsumers(timeout time.Duration) error {
log.L.Debug("removing stale consumers")
wildcardKey := db.ConsumersKey("*")
staleTS := time.Now().Add(timeout * -1).Unix()
for key := range db.Inst.Scan(wildcardKey) {
r := db.Inst.Cmd("ZREMRANGEBYSCORE", key, "-inf", staleTS)
if err := r.Err; err != nil {
return err
}
}
return nil
}
// QueueConsumerCount returns the total number of consumers registered for the
// given queue, either on this okq instance or others
func QueueConsumerCount(queue string) (int64, error) {
consumersKey := db.ConsumersKey(queue)
i, err := db.Inst.Cmd("ZCARD", consumersKey).Int64()
return i, err
}
