├── example ├── jobs │ ├── say_hello.go │ └── write_message.go ├── main.go └── worker.go ├── .gitignore ├── LICENSE ├── broker.go ├── goku.go ├── util.go ├── README.md ├── goku_test.go └── worker.go /example/jobs/say_hello.go: -------------------------------------------------------------------------------- 1 | package jobs 2 | 3 | import "fmt" 4 | 5 | type SayHelloJob struct { 6 | Recipient string 7 | } 8 | 9 | func (j SayHelloJob) Execute() error { 10 | fmt.Printf("Hello, %s\n", j.Name) 11 | return nil 12 | } 13 | 14 | func (j SayHelloJob) Name() string { 15 | return "say_hello_v0" 16 | } 17 | -------------------------------------------------------------------------------- /.gitignore: -------------------------------------------------------------------------------- 1 | # Compiled Object files, Static and Dynamic libs (Shared Objects) 2 | *.o 3 | *.a 4 | *.so 5 | 6 | # Folders 7 | _obj 8 | _test 9 | 10 | # Architecture specific extensions/prefixes 11 | *.[568vq] 12 | [568vq].out 13 | 14 | *.cgo1.go 15 | *.cgo2.c 16 | _cgo_defun.c 17 | _cgo_gotypes.go 18 | _cgo_export.* 19 | 20 | _testmain.go 21 | 22 | *.exe 23 | *.test 24 | *.prof 25 | 26 | coverage.html 27 | cover.out 28 | -------------------------------------------------------------------------------- /example/main.go: -------------------------------------------------------------------------------- 1 | package main 2 | 3 | import ( 4 | "log" 5 | "time" 6 | 7 | "github.com/i/goku" 8 | "github.com/i/goku/example/jobs" 9 | ) 10 | 11 | func main() { 12 | goku.Configure( 13 | goku.BrokerConfig{ 14 | Hostport: "127.0.0.1:6379", 15 | Timeout: time.Second, 16 | }, 17 | ) 18 | 19 | j := jobs.WriteMessageJob{ 20 | To: "Xzibit", 21 | Message: "Hey man", 22 | } 23 | 24 | // schedule the job 25 | if err := goku.Run(j, "lo_priority"); err != nil { 26 | log.Fatal(err) 27 | } 28 | } 29 | -------------------------------------------------------------------------------- /example/jobs/write_message.go: -------------------------------------------------------------------------------- 1 | package jobs 2 | 3 | import ( 4 | "errors" 5 | "fmt" 6 | "os" 7 | ) 8 | 9 | var ErrInvalidArgs = errors.New("invalid args") 10 | 11 | type WriteMessageJob struct { 12 | To string 13 | Message string 14 | } 15 | 16 | func (j WriteMessageJob) Execute() error { 17 | f, err := os.Create(fmt.Sprintf("./message_to_%s.txt", j.To)) 18 | if err != nil { 19 | return err 20 | } 21 | defer f.Close() 22 | 23 | if _, err := f.Write([]byte(j.Message)); err != nil { 24 | return err 25 | } 26 | 27 | return nil 28 | } 29 | 30 | func (j WriteMessageJob) Name() string { 31 | return "write_messag_job_v0" 32 | } 33 | -------------------------------------------------------------------------------- /LICENSE: -------------------------------------------------------------------------------- 1 | The MIT License (MIT) 2 | 3 | Copyright (c) 2016 Ian Lozinski 4 | 5 | Permission is hereby granted, free of charge, to any person obtaining a copy 6 | of this software and associated documentation files (the "Software"), to deal 7 | in the Software without restriction, including without limitation the rights 8 | to use, copy, modify, merge, publish, distribute, sublicense, and/or sell 9 | copies of the Software, and to permit persons to whom the Software is 10 | furnished to do so, subject to the following conditions: 11 | 12 | The above copyright notice and this permission notice shall be included in all 13 | copies or substantial portions of the Software. 14 | 15 | THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR 16 | IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, 17 | FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE 18 | AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER 19 | LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, 20 | OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE 21 | SOFTWARE. 22 | -------------------------------------------------------------------------------- /example/worker.go: -------------------------------------------------------------------------------- 1 | package main 2 | 3 | import ( 4 | "fmt" 5 | "log" 6 | "os" 7 | "os/signal" 8 | "runtime" 9 | "time" 10 | 11 | "github.com/garyburd/redigo/redis" 12 | "github.com/i/goku" 13 | "github.com/i/goku/example/jobs" 14 | ) 15 | 16 | var rc redis.Conn 17 | 18 | func main() { 19 | numWorkers := runtime.NumCPU() 20 | runtime.GOMAXPROCS(numWorkers) 21 | 22 | config := goku.WorkerConfig{ 23 | NumWorkers: numWorkers, 24 | Queues: []string{"lo_priority"}, 25 | Hostport: "127.0.0.1:6379", 26 | Timeout: time.Second, 27 | } 28 | 29 | opts := goku.WorkerPoolOptions{ 30 | Failure: func(worker int, job goku.Job, r interface{}) { 31 | log.Printf("Worker %d failed while executing: %s\n%v\n", worker, job.Name(), r) 32 | }, 33 | Jobs: []goku.Job{ 34 | jobs.WriteMessageJob{}, 35 | }, 36 | } 37 | 38 | wp, err := goku.NewWorkerPool(config, opts) 39 | if err != nil { 40 | log.Fatalf("Error creating worker pool: %v", err) 41 | } 42 | 43 | wp.Start() 44 | fmt.Printf("Started %d workers\n", config.NumWorkers) 45 | 46 | c := make(chan os.Signal) 47 | signal.Notify(c, os.Interrupt, os.Kill) 48 | <-c 49 | 50 | fmt.Println("Shutting down...") 51 | wp.Stop() 52 | os.Exit(0) 53 | } 54 | -------------------------------------------------------------------------------- /broker.go: -------------------------------------------------------------------------------- 1 | package goku 2 | 3 | import ( 4 | "time" 5 | 6 | "github.com/garyburd/redigo/redis" 7 | ) 8 | 9 | // Broker objects schedule jobs to be processed 10 | type Broker struct { 11 | registry map[string]Job 12 | dq string 13 | redisPool *redis.Pool 14 | } 15 | 16 | // BrokerConfig is the information needed to set up a new broker 17 | type BrokerConfig struct { 18 | Hostport string 19 | Password string 20 | Timeout time.Duration 21 | DefaultQueue string 22 | } 23 | 24 | // NewBroker returns a new *Broker. 25 | func NewBroker(cfg BrokerConfig) (*Broker, error) { 26 | redisPool, err := newRedisPool(cfg.Hostport, cfg.Password, cfg.Timeout) 27 | if err != nil { 28 | return nil, err 29 | } 30 | 31 | if cfg.DefaultQueue == "" { 32 | return nil, ErrNoDefaultQueue 33 | } 34 | 35 | return &Broker{ 36 | redisPool: redisPool, 37 | registry: make(map[string]Job), 38 | dq: cfg.DefaultQueue, 39 | }, nil 40 | } 41 | 42 | // Run schedules jobs to be run asynchronously. If queue is not specified, the 43 | // job will be schedules on the default queue. 44 | func (b *Broker) Run(job Job, opts ...JobOption) error { 45 | var jo jobOptions 46 | for _, opt := range opts { 47 | opt.f(&jo) 48 | } 49 | 50 | jsn, err := marshalJob(job) 51 | if err != nil { 52 | return err 53 | } 54 | 55 | conn := b.redisPool.Get() 56 | defer conn.Close() 57 | 58 | if _, err := conn.Do("RPUSH", b.queueOrDefault(jo.queue), jsn); err != nil { 59 | return err 60 | } 61 | return nil 62 | } 63 | 64 | func (b *Broker) RunAt(job Job, t time.Time, opts ...JobOption) error { 65 | var jo jobOptions 66 | for _, opt := range opts { 67 | opt.f(&jo) 68 | } 69 | 70 | jsn, err := marshalJob(job) 71 | if err != nil { 72 | return err 73 | } 74 | 75 | conn := b.redisPool.Get() 76 | defer conn.Close() 77 | 78 | queue := scheduledQueue(b.queueOrDefault(jo.queue)) 79 | if _, err := conn.Do("ZADD", queue, t.UTC().Unix(), jsn); err != nil { 80 | return err 81 | } 82 | return nil 83 | } 84 | -------------------------------------------------------------------------------- /goku.go: -------------------------------------------------------------------------------- 1 | package goku 2 | 3 | import ( 4 | "errors" 5 | "time" 6 | ) 7 | 8 | // generic goku errors 9 | var ( 10 | ErrPointer = errors.New("method receiver was a pointer when it shouldn't be") 11 | ErrStdNotInitialized = errors.New("default broker hasn't been initialized") 12 | ErrInvalidQueue = errors.New("invalid queue name") 13 | ErrNoDefaultQueue = errors.New("no default queue name provided") 14 | ErrNoRedis = errors.New("can't establish a connection to redis") 15 | ErrInvalidJob = errors.New("invalid job") 16 | ) 17 | 18 | // std is the default broker 19 | var std *Broker 20 | 21 | // Configure configures the default broker for package level use 22 | func Configure(cfg BrokerConfig) error { 23 | b, err := NewBroker(cfg) 24 | if err != nil { 25 | return err 26 | } 27 | std = b 28 | return nil 29 | } 30 | 31 | // Job is any type that implements Execute and Name. In order for a job to be 32 | // valid, all fields used within its Execute method must be exported. 33 | type Job interface { 34 | Name() string 35 | Execute() error 36 | } 37 | 38 | // JobOption specifies an option for a job. 39 | type JobOption struct { 40 | f func(*jobOptions) 41 | } 42 | 43 | // JobQueue specifies which queue to run a job on. 44 | func JobQueue(queue string) JobOption { 45 | return JobOption{func(o *jobOptions) { 46 | o.queue = queue 47 | }} 48 | } 49 | 50 | type jobOptions struct { 51 | queue string 52 | } 53 | 54 | // Run schedules a job using the default broker. Before calling goku.Run, the 55 | // default client must be configured using goku.Configure. 56 | func Run(j Job, opts ...JobOption) error { 57 | if std == nil { 58 | return ErrStdNotInitialized 59 | } 60 | return std.Run(j, opts...) 61 | } 62 | 63 | // RunAt is the same as Run, except it schedules a job to run no sooner than 64 | // time t. 65 | func RunAt(j Job, t time.Time, opts ...JobOption) error { 66 | if std == nil { 67 | return ErrStdNotInitialized 68 | } 69 | return std.RunAt(j, t, opts...) 70 | } 71 | -------------------------------------------------------------------------------- /util.go: -------------------------------------------------------------------------------- 1 | package goku 2 | 3 | import ( 4 | "encoding/json" 5 | "fmt" 6 | "reflect" 7 | "time" 8 | 9 | "github.com/garyburd/redigo/redis" 10 | ) 11 | 12 | type marshalledJob struct { 13 | N string 14 | A map[string]interface{} 15 | } 16 | 17 | func (b *Broker) queueOrDefault(q string) string { 18 | if q == "" { 19 | return b.dq 20 | } 21 | return q 22 | } 23 | 24 | func convertFloat(kind reflect.Kind, f float64) interface{} { 25 | switch kind { 26 | case reflect.Int: 27 | return int(f) 28 | case reflect.Int8: 29 | return int8(f) 30 | case reflect.Int16: 31 | return int16(f) 32 | case reflect.Int32: 33 | return int32(f) 34 | case reflect.Int64: 35 | return int64(f) 36 | case reflect.Uint: 37 | return uint(f) 38 | case reflect.Uint8: 39 | return uint8(f) 40 | case reflect.Uint16: 41 | return uint16(f) 42 | case reflect.Uint32: 43 | return uint32(f) 44 | case reflect.Uint64: 45 | return uint64(f) 46 | case reflect.Uintptr: 47 | return uintptr(f) 48 | case reflect.Float32: 49 | return float32(f) 50 | case reflect.Float64: 51 | return f 52 | default: 53 | return 0 54 | } 55 | } 56 | 57 | func marshalJob(job Job) ([]byte, error) { 58 | rv := reflect.ValueOf(job) 59 | rt := reflect.TypeOf(job) 60 | for rv.Kind() == reflect.Ptr { 61 | return nil, ErrPointer 62 | } 63 | 64 | args := make(map[string]interface{}) 65 | 66 | for i := 0; i < rv.NumField(); i++ { 67 | field := rt.Field(i) 68 | value := rv.Field(i) 69 | args[field.Name] = value.Interface() 70 | } 71 | 72 | return json.Marshal(marshalledJob{N: job.Name(), A: args}) 73 | } 74 | 75 | func scheduledQueue(queue string) string { 76 | return fmt.Sprintf("z:%s", queue) 77 | } 78 | 79 | func newRedisPool(hostport, password string, timeout time.Duration) (*redis.Pool, error) { 80 | pool := &redis.Pool{ 81 | MaxIdle: 3, 82 | IdleTimeout: timeout, 83 | Dial: func() (redis.Conn, error) { 84 | c, err := redis.Dial("tcp", hostport) 85 | if err != nil { 86 | return nil, err 87 | } 88 | if password != "" { 89 | if _, err := c.Do("AUTH", password); err != nil { 90 | c.Close() 91 | return nil, err 92 | } 93 | } 94 | return c, err 95 | }, 96 | } 97 | 98 | conn := pool.Get() 99 | defer conn.Close() 100 | 101 | // test the connection 102 | _, err := conn.Do("SETEX", "FOO", 3, "BAR") 103 | if err != nil { 104 | return nil, ErrNoRedis 105 | } 106 | return pool, nil 107 | } 108 | -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | goku 2 | ======== 3 | distributed task queue for go 4 | 5 | how to get 6 | ------------ 7 | 8 | go get github.com/i/goku 9 | 10 | 11 | defining jobs 12 | -------- 13 | 14 | ```go 15 | package jobs 16 | 17 | import "email" 18 | 19 | /* 20 | Defining jobs is easy. A job only needs to implement two methods: Name() and Execute(). 21 | Name() simply returns the name of the job. The name should be unique so that workers know how what type of job it is and if they can correctly process it. 22 | */ 23 | 24 | // All fields used in Execute() must be exported! 25 | type SendEmailJob struct { 26 | To string 27 | From string 28 | Subject string 29 | Body string 30 | } 31 | 32 | // Receivers of Execute() must be structs. 33 | func (j SendEmailJob) Execute() error { 34 | return email.Send(j.To, j.From, j.Subject, j.Body) 35 | } 36 | 37 | // The return value from Name() should be unique from other jobs because it 38 | // is used to differentiate between different jobs. 39 | func (j SendEmailJob) Name() string { 40 | return "send_email_job_v0" 41 | } 42 | 43 | ``` 44 | 45 | how to queue up jobs 46 | --------- 47 | 48 | ```go 49 | package main 50 | 51 | import ( 52 | "time" 53 | 54 | "github.com/i/goku" 55 | 56 | "./jobs" 57 | ) 58 | 59 | func main() { 60 | err := goku.Configure(goku.BrokerConfig{ 61 | Hostport: "127.0.0.1:6379", 62 | Timeout: time.Second, 63 | DefaultQueue: "goku_queue", 64 | }) 65 | if err != nil { 66 | log.Fatalf("Couldn't configure goku: %v", err) 67 | } 68 | 69 | job := jobs.SendEmailJob{ 70 | To: "Will Smith", 71 | From: "Ian", 72 | Subject: "re: Men in Black 2", 73 | Body: "I thought it was pretty good", 74 | } 75 | 76 | // schedule the job to run immediately on the broker's default queue 77 | if err := goku.Run(job); err != nil { 78 | panic("will probably won't get this...") 79 | } 80 | 81 | // schedule the job to run in an hour 82 | if err := goku.RunAt(time.Now().Add(time.Hour)); err != nil { 83 | panic("he's never gonna read these messages :(") 84 | } 85 | } 86 | 87 | ``` 88 | 89 | how to execute jobs 90 | --------- 91 | 92 | ```go 93 | package main 94 | 95 | import ( 96 | "time" 97 | 98 | "github.com/i/goku" 99 | 100 | "./jobs" 101 | ) 102 | 103 | func main() { 104 | config := goku.WorkerConfig{ 105 | NumWorkers: 1, 106 | Queues: []string{"hi_priority"], 107 | Timeout: time.Second, 108 | Hostport: "127.0.0.1:6379", 109 | } 110 | 111 | opts := goku.WorkerPoolOptions{ 112 | Jobs: []goku.Job{ 113 | jobs.WriteMessageJob{}, 114 | } 115 | } 116 | 117 | wp, err := goku.NewWorkerPool(config, opts) 118 | if err != nil { 119 | log.Fatalf("Error creating worker pool: %v", err) 120 | } 121 | 122 | // doesn't block 123 | wp.Start(config, jobs) 124 | 125 | // wait for something... 126 | 127 | // waits for all current jobs to finish 128 | wp.Stop() 129 | } 130 | ``` 131 | -------------------------------------------------------------------------------- /goku_test.go: -------------------------------------------------------------------------------- 1 | package goku 2 | 3 | import ( 4 | "encoding/json" 5 | "testing" 6 | "time" 7 | 8 | "github.com/garyburd/redigo/redis" 9 | "github.com/stretchr/testify/assert" 10 | "github.com/stretchr/testify/require" 11 | ) 12 | 13 | type TestJob struct { 14 | Foo int 15 | Bar string 16 | } 17 | 18 | func (tj TestJob) Name() string { 19 | return "test_job" 20 | } 21 | 22 | var tjWasCalled bool 23 | 24 | func (tj TestJob) Execute() error { 25 | tjWasCalled = true 26 | return nil 27 | } 28 | 29 | func TestBroker(t *testing.T) { 30 | assert := assert.New(t) 31 | require := require.New(t) 32 | 33 | hostport := "127.0.0.1:6379" 34 | queueName := "goku_test" 35 | 36 | broker, err := NewBroker(BrokerConfig{ 37 | Hostport: hostport, 38 | Timeout: time.Second, 39 | DefaultQueue: queueName, 40 | }) 41 | 42 | require.NoError(err) 43 | 44 | job := TestJob{ 45 | Foo: 4, 46 | Bar: "sup", 47 | } 48 | 49 | err = broker.Run(job) 50 | assert.NoError(err) 51 | 52 | conn, err := redis.Dial("tcp", hostport) 53 | require.NoError(err) 54 | 55 | jsn, err := redis.Bytes(conn.Do("LPOP", queueName)) 56 | assert.NoError(err) 57 | 58 | var m map[string]interface{} 59 | json.Unmarshal(jsn, &m) 60 | args := m["A"].(map[string]interface{}) 61 | 62 | assert.Equal(m["N"], job.Name()) 63 | assert.Equal(args["Foo"], float64(4)) 64 | assert.Equal(args["Bar"], "sup") 65 | } 66 | 67 | func TestRun(t *testing.T) { 68 | assert := assert.New(t) 69 | require := require.New(t) 70 | 71 | queue := "goku_test" 72 | hostport := "127.0.0.1:6379" 73 | 74 | config := WorkerConfig{ 75 | NumWorkers: 1, 76 | Queues: []string{queue}, 77 | Hostport: hostport, 78 | Timeout: time.Second, 79 | } 80 | 81 | opts := WorkerPoolOptions{ 82 | Failure: nil, 83 | Jobs: []Job{ 84 | TestJob{}, 85 | }, 86 | } 87 | 88 | // start the worker 89 | wp, err := NewWorkerPool(config, opts) 90 | assert.NoError(err) 91 | wp.Start() 92 | 93 | tjWasCalled = false 94 | 95 | // schedule the job from the broker 96 | broker, err := NewBroker(BrokerConfig{ 97 | Hostport: hostport, 98 | Timeout: time.Second, 99 | DefaultQueue: queue, 100 | }) 101 | 102 | require.NoError(err) 103 | 104 | job := TestJob{ 105 | Foo: 4, 106 | Bar: "sup", 107 | } 108 | 109 | err = broker.Run(job) 110 | assert.NoError(err) 111 | time.Sleep(time.Second) // give workers some time to pull the job out of the queue 112 | wp.Stop() 113 | 114 | assert.True(tjWasCalled) 115 | } 116 | 117 | func TestBrokerBadConfig(t *testing.T) { 118 | _, err := NewBroker(BrokerConfig{}) 119 | assert.Error(t, err) 120 | } 121 | 122 | func TestWorkerPoolBadConfig(t *testing.T) { 123 | _, err := NewWorkerPool(WorkerConfig{}, WorkerPoolOptions{}) 124 | assert.Error(t, err) 125 | } 126 | 127 | func TestConfigureBadConfig(t *testing.T) { 128 | err := Configure(BrokerConfig{}) 129 | assert.Error(t, err) 130 | } 131 | 132 | func TestConfigureGoodConfig(t *testing.T) { 133 | hostport := "127.0.0.1:6379" 134 | queueName := "goku_test" 135 | err := Configure(BrokerConfig{ 136 | Hostport: hostport, 137 | Timeout: time.Second, 138 | DefaultQueue: queueName, 139 | }) 140 | assert.NoError(t, err) 141 | } 142 | 143 | func TestRunWithPtr(t *testing.T) { 144 | hostport := "127.0.0.1:6379" 145 | queueName := "goku_test" 146 | err := Configure(BrokerConfig{ 147 | Hostport: hostport, 148 | Timeout: time.Second, 149 | DefaultQueue: queueName, 150 | }) 151 | require.NoError(t, err) 152 | 153 | err = Run(&TestJob{}) 154 | assert.Equal(t, ErrPointer, err) 155 | } 156 | 157 | func TestRunAt(t *testing.T) { 158 | assert := assert.New(t) 159 | require := require.New(t) 160 | 161 | queue := "goku_test" 162 | hostport := "127.0.0.1:6379" 163 | 164 | config := WorkerConfig{ 165 | NumWorkers: 1, 166 | Queues: []string{queue}, 167 | Hostport: hostport, 168 | Timeout: time.Second, 169 | } 170 | 171 | opts := WorkerPoolOptions{ 172 | Failure: nil, 173 | Jobs: []Job{ 174 | TestJob{}, 175 | }, 176 | } 177 | 178 | // start the worker 179 | wp, err := NewWorkerPool(config, opts) 180 | assert.NoError(err) 181 | wp.Start() 182 | 183 | // schedule the job from the broker 184 | broker, err := NewBroker(BrokerConfig{ 185 | Hostport: hostport, 186 | Timeout: time.Second, 187 | DefaultQueue: queue, 188 | }) 189 | 190 | require.NoError(err) 191 | 192 | job := TestJob{ 193 | Foo: 4, 194 | Bar: "sup", 195 | } 196 | 197 | tjWasCalled = false 198 | err = broker.RunAt(job, time.Now().Add(3*time.Second)) 199 | assert.NoError(err) 200 | 201 | // give workers some time to pull the job out of the queue 202 | time.Sleep(2 * time.Second) 203 | assert.False(tjWasCalled) 204 | 205 | time.Sleep(2 * time.Second) 206 | wp.Stop() 207 | assert.True(tjWasCalled) 208 | } 209 | -------------------------------------------------------------------------------- /worker.go: -------------------------------------------------------------------------------- 1 | package goku 2 | 3 | import ( 4 | "encoding/json" 5 | "reflect" 6 | "strings" 7 | "sync" 8 | "time" 9 | 10 | "github.com/garyburd/redigo/redis" 11 | ) 12 | 13 | // redis commands (prevent typos) 14 | const ( 15 | blpop = "BLPOP" 16 | rpush = "RPUSH" 17 | zremrangebyscore = "ZREMRANGEBYSCORE" 18 | zrangebyscore = "ZRANGEBYSCORE" 19 | ) 20 | 21 | // FailureFunc is a function that gets executed when a job fails. It will get 22 | // run when a job returns an error or panics. 23 | type FailureFunc func(worker int, job Job, r interface{}) 24 | 25 | // WorkerConfig describes the configuration needed for setting up a new worker 26 | // pool. 27 | type WorkerConfig struct { 28 | NumWorkers int // number of workers that belong to the pool 29 | Queues []string // what queues to pull jobs from 30 | Hostport string // redis hostport 31 | Password string // redis auth password (optional) 32 | Timeout time.Duration // redis timeout 33 | 34 | // If a worker doesn't know how to handle a job it will be requeued. 35 | // sometimes requeuing can fail. This field is max number of retries before 36 | // losing the job. 37 | RequeRetries int 38 | } 39 | 40 | // WorkerPool is what will pull jobs from redis and distribute them to workers 41 | // within the pool. 42 | type WorkerPool struct { 43 | queues []string 44 | redisPool *redis.Pool 45 | fail FailureFunc 46 | workCh chan qj 47 | requeueMap map[string]chan []byte 48 | killCh chan struct{} 49 | numWorkers int 50 | requeueRetries int 51 | registry map[string]Job 52 | timeout time.Duration 53 | wg sync.WaitGroup 54 | m sync.RWMutex 55 | running bool 56 | } 57 | 58 | // WorkerPoolOptions exists for defining things that wouldn't be possible 59 | // within a yaml configuration file. Failure is optional, but jobs are required 60 | // if you want the workers to do anything. 61 | type WorkerPoolOptions struct { 62 | Failure FailureFunc 63 | Jobs []Job 64 | } 65 | 66 | // NewWorkerPool returns a new WorkerPool. It fails when a connection to redis 67 | // cannot be established. 68 | func NewWorkerPool(cfg WorkerConfig, opts WorkerPoolOptions) (*WorkerPool, error) { 69 | redisPool, err := newRedisPool(cfg.Hostport, cfg.Password, cfg.Timeout) 70 | if err != nil { 71 | return nil, err 72 | } 73 | 74 | wp := &WorkerPool{ 75 | queues: cfg.Queues, 76 | redisPool: redisPool, 77 | workCh: make(chan qj), 78 | requeueMap: make(map[string]chan []byte), 79 | numWorkers: cfg.NumWorkers, 80 | registry: make(map[string]Job), 81 | fail: opts.Failure, 82 | timeout: cfg.Timeout, 83 | requeueRetries: cfg.RequeRetries, 84 | } 85 | 86 | for _, job := range opts.Jobs { 87 | wp.registry[job.Name()] = job 88 | } 89 | 90 | return wp, nil 91 | } 92 | 93 | // Start tells the worker pool to start pulling things off the queue to be 94 | // processed. 95 | func (wp *WorkerPool) Start() { 96 | wp.m.Lock() 97 | wp.running = true 98 | wp.killCh = make(chan struct{}) 99 | 100 | for i := 0; i < wp.numWorkers; i++ { 101 | go wp.startWorker(i) 102 | } 103 | 104 | for _, q := range wp.queues { 105 | wp.requeueMap[q] = make(chan []byte) 106 | go wp.startReqeuer(q) 107 | } 108 | 109 | go wp.startPolling() 110 | go wp.startZPolling() 111 | } 112 | 113 | func (wp *WorkerPool) startPolling() { 114 | qstr := strings.Join(wp.queues, " ") // standard queues 115 | 116 | for wp.running { 117 | conn := wp.redisPool.Get() 118 | res, err := redis.ByteSlices(conn.Do(blpop, qstr, wp.timeout.Seconds())) 119 | conn.Close() 120 | if err != nil { 121 | continue 122 | } 123 | wp.workCh <- qj{string(res[0]), res[1]} 124 | } 125 | } 126 | 127 | func (wp *WorkerPool) startZPolling() { 128 | var zqstrs []string 129 | for _, qname := range wp.queues { 130 | zqstrs = append(zqstrs, scheduledQueue(qname)) 131 | } 132 | 133 | for wp.running { 134 | now := time.Now().UTC().Unix() 135 | for _, zset := range zqstrs { 136 | conn := wp.redisPool.Get() 137 | res, err := redis.ByteSlices(conn.Do(zrangebyscore, zset, 0, now)) 138 | conn.Close() 139 | if err != nil { 140 | continue 141 | } 142 | 143 | for _, jsn := range res { 144 | wp.workCh <- qj{zset, jsn} 145 | conn := wp.redisPool.Get() 146 | _, err := conn.Do(zremrangebyscore, zset, 0, now) 147 | if err != nil { 148 | // TODO -- try again 149 | } 150 | conn.Close() 151 | } 152 | } 153 | } 154 | } 155 | 156 | // Stop waits for all jobs to finish executing, and then returns. 157 | func (wp *WorkerPool) Stop() { 158 | wp.running = false 159 | close(wp.killCh) 160 | wp.wg.Wait() 161 | wp.m.Unlock() 162 | } 163 | 164 | type qj struct { 165 | queue string 166 | jsn []byte 167 | } 168 | 169 | func (wp *WorkerPool) startWorker(n int) { 170 | for { 171 | select { 172 | case <-wp.killCh: 173 | return 174 | case qj := <-wp.workCh: 175 | job, err := wp.getJob(qj.jsn) 176 | if err != nil { 177 | wp.requeueMap[qj.queue] <- qj.jsn 178 | continue 179 | } 180 | wp.doWork(job, n) 181 | } 182 | } 183 | } 184 | 185 | func (wp *WorkerPool) doWork(job Job, n int) { 186 | if wp.fail != nil { 187 | defer func() { 188 | if r := recover(); r != nil { 189 | wp.fail(n, job, r) 190 | } 191 | }() 192 | } 193 | 194 | wp.wg.Add(1) 195 | defer wp.wg.Done() 196 | if err := job.Execute(); err != nil { 197 | wp.fail(n, job, err) 198 | } 199 | } 200 | 201 | // getJob converts a json payload into a a Job with populated fields 202 | func (wp *WorkerPool) getJob(jsn []byte) (Job, error) { 203 | var j marshalledJob 204 | if err := json.Unmarshal(jsn, &j); err != nil { 205 | return nil, err 206 | } 207 | 208 | emptyJob, ok := wp.registry[j.N] 209 | if !ok { 210 | return nil, ErrInvalidJob 211 | } 212 | 213 | rt := reflect.TypeOf(emptyJob) 214 | nj := reflect.New(rt).Elem() 215 | 216 | for k, v := range j.A { 217 | field := nj.FieldByName(k) 218 | if field.CanSet() { 219 | if f, ok := v.(float64); ok { 220 | v = convertFloat(field.Kind(), f) 221 | } 222 | field.Set(reflect.ValueOf(v)) 223 | } 224 | } 225 | 226 | job, ok := nj.Interface().(Job) 227 | if !ok { 228 | return nil, ErrInvalidJob 229 | } 230 | return job, nil 231 | } 232 | 233 | func (wp *WorkerPool) startReqeuer(qn string) { 234 | ch := wp.requeueMap[qn] 235 | for { 236 | select { 237 | case <-wp.killCh: 238 | return 239 | case jsn := <-ch: 240 | wp.wg.Add(1) 241 | for i := 0; i < wp.requeueRetries; i++ { 242 | conn := wp.redisPool.Get() 243 | _, err := conn.Do(rpush, qn, jsn) 244 | conn.Close() 245 | if err == nil { 246 | break 247 | } 248 | } 249 | wp.wg.Done() 250 | } 251 | } 252 | } 253 | --------------------------------------------------------------------------------