forked from benmanns/goworker
/
goworker.go
123 lines (105 loc) · 3 KB
/
goworker.go
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
package goworker
import (
"github.com/youtube/vitess/go/pools"
"github.com/cihub/seelog"
"golang.org/x/net/context"
"os"
"strconv"
"sync"
"time"
)
var (
logger seelog.LoggerInterface
pool *pools.ResourcePool
ctx context.Context
)
// Init initializes the goworker process. This will be
// called by the Work function, but may be used by programs
// that wish to access goworker functions and configuration
// without actually processing jobs.
func Init() error {
var err error
logger, err = seelog.LoggerFromWriterWithMinLevel(os.Stdout, seelog.InfoLvl)
if err != nil {
return err
}
if err := flags(); err != nil {
return err
}
ctx = context.Background()
pool = newRedisPool(uri, connections, connections, time.Minute)
return nil
}
// GetConn returns a connection from the goworker Redis
// connection pool. When using the pool, check in
// connections as quickly as possible, because holding a
// connection will cause concurrent worker functions to lock
// while they wait for an available connection. Expect this
// API to change drastically.
func GetConn() (*RedisConn, error) {
resource, err := pool.Get(ctx)
if err != nil {
return nil, err
}
// Applied suggested fix from https://github.com/benmanns/goworker/issues/15
// to force a redis reconnect on disconnect.
// Performs simple connection check. Redial if needed
_, err = resource.(*RedisConn).Do("PING")
if err != nil {
resource.(*RedisConn).Close()
resource, err = redisConnFromUri(uri)
if err != nil {
return nil, err
}
}
return resource.(*RedisConn), nil
}
// PutConn puts a connection back into the connection pool.
// Run this as soon as you finish using a connection that
// you got from GetConn. Expect this API to change
// drastically.
func PutConn(conn *RedisConn) {
pool.Put(conn)
}
// Close cleans up resources initialized by goworker. This
// will be called by Work when cleaning up. However, if you
// are using the Init function to access goworker functions
// and configuration without processing jobs by calling
// Work, you should run this function when cleaning up. For
// example,
//
// if err := goworker.Init(); err != nil {
// fmt.Println("Error:", err)
// }
// defer goworker.Close()
func Close() {
pool.Close()
}
// Work starts the goworker process. Check for errors in
// the return value. Work will take over the Go executable
// and will run until a QUIT, INT, or TERM signal is
// received, or until the queues are empty if the
// -exit-on-complete flag is set.
func Work() error {
err := Init()
if err != nil {
return err
}
defer Close()
quit := signals()
poller, err := newPoller(queues, isStrict)
if err != nil {
return err
}
jobs := poller.poll(time.Duration(interval), quit)
var monitor sync.WaitGroup
for id := 0; id < concurrency; id++ {
worker, err := newWorker(strconv.Itoa(id), queues)
if err != nil {
return err
}
worker.work(jobs, &monitor)
}
monitor.Wait()
return nil
}