// GroupTaskStates - returns states of all tasks in the group
func (redisBackend *RedisBackend) GroupTaskStates(groupUUID string, groupTaskCount int) ([]*TaskState, error) {
taskStates := make([]*TaskState, groupTaskCount)
groupMeta := GroupMeta{}
conn, err := redisBackend.open()
if err != nil {
return taskStates, err
}
defer conn.Close()
item, err := redis.Bytes(conn.Do("GET", groupUUID))
if err != nil {
return taskStates, err
}
if err := json.Unmarshal(item, &groupMeta); err != nil {
return taskStates, err
}
for i, taskUUID := range groupMeta.TaskUUIDs {
taskState, err := redisBackend.GetState(taskUUID)
if err != nil {
return taskStates, err
}
taskStates[i] = taskState
}
return taskStates, nil
}
// GroupCompleted - returns true if all tasks in a group finished
func (redisBackend *RedisBackend) GroupCompleted(groupUUID string, groupTaskCount int) (bool, error) {
groupMeta := GroupMeta{}
conn, err := redisBackend.open()
if err != nil {
return false, err
}
defer conn.Close()
item, err := redis.Bytes(conn.Do("GET", groupUUID))
if err != nil {
return false, err
}
if err := json.Unmarshal(item, &groupMeta); err != nil {
return false, err
}
for _, taskUUID := range groupMeta.TaskUUIDs {
taskState, err := redisBackend.GetState(taskUUID)
if err != nil {
return false, err
}
if !taskState.IsCompleted() {
return false, nil
}
}
return true, nil
}
// Consumes messages
func (redisBroker *RedisBroker) consume(errorsChan chan error, taskProcessor TaskProcessor) {
defer redisBroker.wg.Done()
for {
conn, err := redisBroker.open()
if err != nil {
redisBroker.retryFunc()
errorsChan <- err
return
}
redisBroker.retryFunc = utils.RetryClosure()
defer conn.Close()
log.Print("[*] Waiting for messages. To exit press CTRL+C")
select {
case <-redisBroker.quitChan:
return
default:
// Return value of BLPOP is an array. For example:
// redis> RPUSH list1 a b c
// (integer) 3
// redis> BLPOP list1 list2 0
// 1) "list1"
// 2) "a"
multiBulk, err := redis.MultiBulk(conn.Do("BLPOP", redisBroker.config.DefaultQueue, "0"))
if err != nil {
errorsChan <- err
return
}
item, err := redis.Bytes(multiBulk[1], nil)
if err != nil {
errorsChan <- err
return
}
log.Printf("Received new message: %s", item)
signature := signatures.TaskSignature{}
if err := json.Unmarshal(item, &signature); err != nil {
errorsChan <- err
return
}
if err := taskProcessor.Process(&signature); err != nil {
errorsChan <- err
return
}
}
}
}
// GetState - returns the latest task state
func (redisBackend *RedisBackend) GetState(taskUUID string) (*TaskState, error) {
taskState := TaskState{}
conn, err := redisBackend.open()
if err != nil {
return nil, err
}
defer conn.Close()
item, err := redis.Bytes(conn.Do("GET", taskUUID))
if err != nil {
return nil, err
}
if err := json.Unmarshal(item, &taskState); err != nil {
return nil, err
}
return &taskState, nil
}
// Updates a task state group
func (redisBackend *RedisBackend) updateStateGroup(groupUUID string, groupTaskCount int, taskState *TaskState) (*TaskStateGroup, error) {
var taskStateGroup *TaskStateGroup
conn, err := redisBackend.open()
if err != nil {
return nil, err
}
defer conn.Close()
item, err := redis.Bytes(conn.Do("GET", groupUUID))
if err != nil {
taskStateGroup = &TaskStateGroup{
GroupUUID: groupUUID,
GroupTaskCount: groupTaskCount,
States: make(map[string]TaskState),
}
} else {
if err := json.Unmarshal(item, &taskStateGroup); err != nil {
return nil, err
}
}
taskStateGroup.States[taskState.TaskUUID] = *taskState
encoded, err := json.Marshal(taskStateGroup)
if err != nil {
return nil, fmt.Errorf("JSON Encode Message: %v", err)
}
_, err = conn.Do("SET", groupUUID, encoded)
if err != nil {
return nil, err
}
return taskStateGroup, redisBackend.setExpirationTime(groupUUID)
}
// StartConsuming enters a loop and waits for incoming messages
func (redisBroker *RedisBroker) StartConsuming(consumerTag string, taskProcessor TaskProcessor) (bool, error) {
if redisBroker.retryFunc == nil {
redisBroker.retryFunc = utils.RetryClosure()
}
redisBroker.pool = redisBroker.newPool()
defer redisBroker.pool.Close()
_, err := redisBroker.pool.Get().Do("PING")
if err != nil {
redisBroker.retryFunc()
return true, err // retry true
}
redisBroker.retryFunc = utils.RetryClosure()
redisBroker.stopChan = make(chan int)
redisBroker.stopReceivingChan = make(chan int)
redisBroker.errorsChan = make(chan error)
deliveries := make(chan []byte)
redisBroker.wg.Add(1)
go func() {
defer redisBroker.wg.Done()
log.Print("[*] Waiting for messages. To exit press CTRL+C")
conn := redisBroker.pool.Get()
for {
select {
// A way to stop this goroutine from redisBroker.StopConsuming
case <-redisBroker.stopReceivingChan:
return
default:
itemBytes, err := conn.Do("LPOP", redisBroker.config.DefaultQueue)
if err != nil {
redisBroker.errorsChan <- err
return
}
// Unline BLPOP, LPOP is non blocking so nil means we can keep iterating
if itemBytes == nil {
continue
}
item, err := redis.Bytes(itemBytes, nil)
if err != nil {
redisBroker.errorsChan <- err
return
}
deliveries <- item
}
}
}()
if err := redisBroker.consume(deliveries, taskProcessor); err != nil {
return true, err // retry true
}
return false, nil
}