summaryrefslogtreecommitdiffstats
path: root/modules/task/queue_redis.go
blob: 127de0cdbf1d3865589949298459b61bde081e2c (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
// Copyright 2019 The Gitea Authors. All rights reserved.
// Use of this source code is governed by a MIT-style
// license that can be found in the LICENSE file.

package task

import (
	"encoding/json"
	"errors"
	"strconv"
	"strings"
	"time"

	"code.gitea.io/gitea/models"
	"code.gitea.io/gitea/modules/log"

	"github.com/go-redis/redis"
)

var (
	_ Queue = &RedisQueue{}
)

type redisClient interface {
	RPush(key string, args ...interface{}) *redis.IntCmd
	LPop(key string) *redis.StringCmd
	Ping() *redis.StatusCmd
}

// RedisQueue redis queue
type RedisQueue struct {
	client    redisClient
	queueName string
	closeChan chan bool
}

func parseConnStr(connStr string) (addrs, password string, dbIdx int, err error) {
	fields := strings.Fields(connStr)
	for _, f := range fields {
		items := strings.SplitN(f, "=", 2)
		if len(items) < 2 {
			continue
		}
		switch strings.ToLower(items[0]) {
		case "addrs":
			addrs = items[1]
		case "password":
			password = items[1]
		case "db":
			dbIdx, err = strconv.Atoi(items[1])
			if err != nil {
				return
			}
		}
	}
	return
}

// NewRedisQueue creates single redis or cluster redis queue
func NewRedisQueue(addrs string, password string, dbIdx int) (*RedisQueue, error) {
	dbs := strings.Split(addrs, ",")
	var queue = RedisQueue{
		queueName: "task_queue",
		closeChan: make(chan bool),
	}
	if len(dbs) == 0 {
		return nil, errors.New("no redis host found")
	} else if len(dbs) == 1 {
		queue.client = redis.NewClient(&redis.Options{
			Addr:     strings.TrimSpace(dbs[0]), // use default Addr
			Password: password,                  // no password set
			DB:       dbIdx,                     // use default DB
		})
	} else {
		// cluster will ignore db
		queue.client = redis.NewClusterClient(&redis.ClusterOptions{
			Addrs:    dbs,
			Password: password,
		})
	}
	if err := queue.client.Ping().Err(); err != nil {
		return nil, err
	}
	return &queue, nil
}

// Run starts to run the queue
func (r *RedisQueue) Run() error {
	for {
		select {
		case <-r.closeChan:
			return nil
		case <-time.After(time.Millisecond * 100):
		}

		bs, err := r.client.LPop(r.queueName).Bytes()
		if err != nil {
			if err != redis.Nil {
				log.Error("LPop failed: %v", err)
			}
			time.Sleep(time.Millisecond * 100)
			continue
		}

		var task models.Task
		err = json.Unmarshal(bs, &task)
		if err != nil {
			log.Error("Unmarshal task failed: %s", err.Error())
		} else {
			err = Run(&task)
			if err != nil {
				log.Error("Run task failed: %s", err.Error())
			}
		}
	}
}

// Push implements Queue
func (r *RedisQueue) Push(task *models.Task) error {
	bs, err := json.Marshal(task)
	if err != nil {
		return err
	}
	return r.client.RPush(r.queueName, bs).Err()
}

// Stop stop the queue
func (r *RedisQueue) Stop() {
	r.closeChan <- true
}