func (m *beanstalkdPubSub) watchIncomingMessages() {
if m.connSub == nil {
return
}
// create tubeset for topic
tubeset := beanstalk.NewTubeSet(m.connSub, m.topic)
watchLoop:
for {
select {
// watch for close signal
case <-m.c.Done():
return
default:
// get the message
id, body, err := tubeset.Reserve(time.Minute)
if err == nil {
// broadcast it to all subscribers
m.broadcast(body)
// simply delete the message
m.connSub.Delete(id)
continue
}
if err.(beanstalk.ConnError).Err == beanstalk.ErrTimeout {
// re-reserve
continue
} else if err.(beanstalk.ConnError).Err == beanstalk.ErrDeadline {
time.Sleep(time.Second)
// re-reserve
continue
} else if m.possibleNetworkError(err.(beanstalk.ConnError).Err) {
// try reconnecting
for {
select {
case <-m.c.Done():
return
default:
<-time.After(time.Second * 3)
if err := m.dialSubSocket(); err != nil {
continue
}
tubeset = beanstalk.NewTubeSet(m.connSub, m.topic)
goto watchLoop
}
}
}
}
}
}
func (m *mainFrame) buryJobs() error {
tubeSet := beanstalk.NewTubeSet(m.c, m.currentTubeName())
n := 0
var lastError error
for {
id, _, err := tubeSet.Reserve(time.Second)
if err != nil {
lastError = lastError
break
}
stats, err := m.c.StatsJob(id)
if err != nil {
lastError = err
break
}
pri := strToInt(stats["pri"])
if err := m.c.Bury(id, uint32(pri)); err != nil {
lastError = err
break
}
n++
}
m.showStatus(fmt.Sprintf("%s: %d jobs buried", m.currentTubeName(), n))
return lastError
}
func (c *TailCommand) Tail() error {
ts := beanstalk.NewTubeSet(c.conn, c.Tube)
errors := 0
for {
if errors > 100 {
return TooManyErrorsError
}
id, body, err := ts.Reserve(time.Hour * 24)
if err != nil {
if err.Error() != "reserve-with-timeout: deadline soon" {
errors++
fmt.Println("Error", err)
}
continue
}
if err := c.PrintJob(id, body); err != nil {
errors++
fmt.Println("Error", err)
continue
}
if err := c.postPrintAction(id); err != nil {
return err
}
fmt.Println(strings.Repeat("-", 80))
}
return nil
}
func runGet(cmd *Command) {
conn := DialBeanstalk()
ts := beanstalk.NewTubeSet(conn, strings.Split(*getTubes, ",")...)
n := *getNum
w := time.Duration(*getTimeout) * time.Second
var ok bool
var action func(*beanstalk.Conn, uint64, []byte)
if action, ok = Actions[*getAction]; !ok {
fatal(2, "'%s' isn't a valid action", *getAction)
}
if *getAction == "r" && n == 0 {
// Protect users from themselves
fatal(2, "Using -n 0 and -x r together causes a tight loop and is disallowed")
}
for i := uint64(0); n == 0 || i < n; i++ {
id, body, err := ts.Reserve(w)
if err != nil {
if cerr, ok := err.(beanstalk.ConnError); ok && cerr.Err == beanstalk.ErrTimeout {
// Only write message if no jobs at all, but exit w/ 0
if i == 0 {
writeStderr("No jobs")
}
return
}
fatal(2, "Error getting job:\n%v", err)
}
fmt.Printf("%s\n", body)
action(conn, id, body)
}
}
func main() {
flag.Parse()
var tb *beanstalk.TubeSet
var conn_bs *beanstalk.Conn
rs_timeout := time.Duration(Settings.BeanstalkdReserveTimeout)
fail_wait := time.Duration(Settings.FailWait) * time.Second
conn_bs, e := beanstalk.Dial("tcp", Settings.BeanstalkdAddr)
if e != nil {
log.Fatal("failed to connected to beanstalkd", e)
}
tb = beanstalk.NewTubeSet(conn_bs, Settings.BeanstalkdTube)
for {
// reserve a job
id, job, e := tb.Reserve(rs_timeout)
// timeout is valid, anything else is fatal
if cerr, ok := e.(beanstalk.ConnError); ok && cerr.Err == beanstalk.ErrTimeout {
time.Sleep(fail_wait)
continue
} else if e != nil {
log.Fatal("failed to reserve job", e)
} else {
log.Println("read job id", id, "size", len(job), "bytes")
}
// connect to the gor replay server
conn_gr, e := net.Dial("tcp", Settings.GorReplayAddr)
if e != nil {
log.Fatal("failed to connected to gor replay server", e)
time.Sleep(fail_wait)
}
// write to gor replay server
w, e := conn_gr.Write(job)
if e != nil {
log.Fatal("failed to write to", Settings.GorReplayAddr, "error", e)
} else {
log.Println("wrote", w, "bytes to", Settings.GorReplayAddr)
}
// close connection to gor replay server
conn_gr.Close()
// delete the job
e = conn_bs.Delete(id)
if e != nil {
log.Println("failed to delete job id", id, "error", e)
}
}
}
func Example_reserveOtherTubeSet() {
tubeSet := beanstalk.NewTubeSet(conn, "mytube1", "mytube2")
id, body, err := tubeSet.Reserve(10 * time.Hour)
if err != nil {
panic(err)
}
fmt.Println("job", id)
fmt.Println(string(body))
}
// newTubeSet creates a new tubeSet for a tube name
func newTubeSet(conn *beanstalk.Conn, name string) *tubeSet {
return &tubeSet{
consume: beanstalk.NewTubeSet(conn, name),
publish: &beanstalk.Tube{
Conn: conn,
Name: name,
},
}
}
func (t *BeanWorker) Open(tube string) error {
conn, err := beanstalk.Dial("tcp", t.address)
if err != nil {
return err
}
conn.Tube = beanstalk.Tube{conn, tube}
conn.TubeSet = *beanstalk.NewTubeSet(conn, tube)
t.conn = conn
return nil
}
func (c *Client) tubes(conn *beanstalk.Conn) map[string]*beanstalk.TubeSet {
names := []string{"default"}
if mux, isMux := c.Handler.(*WorkMux); isMux {
names = mux.Tubes()
}
tubes := make(map[string]*beanstalk.TubeSet, len(names))
for _, name := range names {
tubes[name] = beanstalk.NewTubeSet(conn, name)
}
return tubes
}
// Reserve and return a job from one of the tubes. If no job is available before time timeout has passed, Reserve returns a ConnError recording ErrTimeout.
func (c *Client) Reserve(timeout time.Duration, tubes ...string) (JobID, []byte, error) {
err := c.Connect()
if err != nil {
return 0, nil, err
}
if len(tubes) == 0 {
tubes = []string{"default"}
}
ts := beanstalk.NewTubeSet(c.Conn, tubes...)
jid, body, err := ts.Reserve(timeout)
if err != nil {
return 0, nil, unwrap(err)
}
return JobID(jid), body, nil
}
func deleteAllFromTube(t *beanstalk.Tube) {
queue := beanstalk.NewTubeSet(t.Conn, t.Name)
deleted := 0
for {
id, _, err := queue.Reserve(3 * time.Second)
if err != nil {
break
}
if err := t.Conn.Delete(id); err != nil {
fmt.Println(err)
continue
}
deleted++
}
fmt.Printf("%d jobs deleted from %s\n", deleted, t.Name)
}
func NewClient(addr string, tubes []string) (client *Client, err error) {
var conn *beanstalk.Conn
if conn, err = beanstalk.Dial("tcp", addr); err != nil {
return
}
conn.TubeSet = *beanstalk.NewTubeSet(conn, tubes...)
client = &Client{
conn: conn,
mu: new(sync.Mutex),
ReserveTimeout: time.Duration(5 * time.Second),
}
return
}
// Run consume available job based on tube watched and proceed the task
func (w *ApnsWorker) Run() {
fmt.Printf("[worker] Starting APNS worker\n")
tube := beanstalk.NewTubeSet(w.Conn, w.TubeName)
for {
id, body, err := tube.Reserve(24 * time.Hour)
if err != nil {
panic(err)
}
var data Job
err = json.Unmarshal(body, &data)
if w.ApnsConn == nil {
certPem, err := ioutil.ReadFile(data.Config.CertPem)
if err != nil {
panic(err)
}
keyPem, err := ioutil.ReadFile(data.Config.KeyPem)
if err != nil {
panic(err)
}
w.ApnsConn, _ = apns.NewAPNSConnection(&apns.APNSConfig{
CertificateBytes: certPem,
KeyBytes: keyPem,
})
defer w.ApnsConn.Disconnect()
}
customField := make(map[string]interface{})
customField["type"] = data.Data.PushType
payload := &apns.Payload{
Token: data.Data.DeviceToken,
AlertText: data.Data.Alert,
Sound: data.Data.Sound,
Badge: apns.NewBadgeNumber(1),
CustomFields: customField,
}
w.ApnsConn.SendChannel <- payload
w.Conn.Delete(id)
}
}
func get(timeout time.Duration, queues ...string) (*Message, error) {
conn, err := connection()
if err != nil {
return nil, err
}
ts := beanstalk.NewTubeSet(conn, queues...)
id, body, err := ts.Reserve(timeout)
if err != nil {
if timeoutRegexp.MatchString(err.Error()) {
return nil, &timeoutError{timeout: timeout}
}
return nil, err
}
defer conn.Delete(id)
r := bytes.NewReader(body)
var msg Message
if err = gob.NewDecoder(r).Decode(&msg); err != nil && err != io.EOF {
return nil, fmt.Errorf("Invalid message: %q", body)
}
return &msg, nil
}
func (c *BuryCommand) Bury() error {
if err := c.calcNum(); err != nil {
return err
}
if c.Num == 0 {
fmt.Printf("Empty ready queue at tube %q.\n", c.Tube)
return nil
}
fmt.Printf("Trying to bury %d jobs from %q ...\n", c.Num, c.Tube)
count := 0
ts := beanstalk.NewTubeSet(c.conn, c.Tube)
for count < c.Num {
id, _, err := ts.Reserve(time.Second)
if err != nil {
return err
}
s, err := c.conn.StatsJob(id)
if err != nil {
return err
}
pri, err := strconv.ParseUint(s["pri"], 10, 32)
if err != nil {
return err
}
if err := c.conn.Bury(id, uint32(pri)); err != nil {
return err
}
count++
}
fmt.Printf("Actually buried %d.\n", count)
return nil
}
// Satisfies the Source interface. Begins listening to an AWS SQS queue. If no
// message is on the queue it sleeps for a set period of time before trying again
func (this *Beanstalkd) Listen(wg sync.WaitGroup) {
sCon := fmt.Sprintf("%s:%d", this.Host, this.Port)
this.Log.Info("Start listening (beanstalkd:%s)", sCon)
c, err := beanstalk.Dial("tcp", sCon)
if err != nil {
this.Log.Error("Could not listen on %s: %s", sCon, err.Error())
wg.Done()
return
}
defer c.Close()
timeout := true
var tubeSet *beanstalk.TubeSet
for {
if timeout {
tube := this.selectTubes(c)
if tube == nil {
time.Sleep(time.Second * 5)
continue
}
tubeSet = beanstalk.NewTubeSet(c, tube.Name)
}
this.Log.Debug("Watching %v", tubeSet.Name)
id, body, err := tubeSet.Reserve(time.Second * 5)
if err != nil {
timeout = true
continue
} else {
timeout = false
}
this.handle(id, body, c)
}
wg.Done()
}
func (w *worker) run(started chan<- struct{}) {
beanConn, err := beanstalk.Dial("tcp", w.options.Host)
if err != nil {
panic(fmt.Sprintf("dial err: %s", err))
}
// worker comm channels
jobs := make(chan Request)
done := make(chan struct{})
defer func() {
// close the conn
beanConn.Close()
// shutdown the workers
close(jobs)
// wait for them to stop
for i := 0; i < w.options.Count; i++ {
select {
case <-done:
case <-w.control.dead:
}
}
close(w.control.shutdown)
}()
// start up our workers
for i := 0; i < w.options.Count; i++ {
go w.work(jobs, done)
}
// watch the worker tube
var watch = beanstalk.NewTubeSet(beanConn, w.tube)
// off we go
close(started)
running := true
jobCnt := 0
for jobCnt > 0 || running {
// check the control channels
select {
case res := <-w.control.completed:
// a worker is finished -- handle it
switch res.result {
case Success:
beanConn.Delete(res.jobID)
case BuryJob:
beanConn.Bury(res.jobID, res.priority)
log.Printf("Burying job. Id: %d\n", res.jobID)
case DeleteJob:
beanConn.Delete(res.jobID)
log.Printf("Deleting job. Id: %d\n", res.jobID)
case ReleaseJob:
beanConn.Release(res.jobID, res.priority, res.delay)
log.Printf("Releasing job for: %s Id: %d %s\n", res.delay.String(), res.jobID)
}
jobCnt--
default:
}
if !running {
<-time.After(250 * time.Millisecond)
continue
}
select {
case <-w.control.dead:
// a worker died -- start up a new one
go w.work(jobs, done)
continue
case <-w.control.shutdown:
// we need to shutdown
running = false
continue
default:
}
// get some work
id, msg, err := watch.Reserve(w.options.Reserve)
if err != nil {
cerr, ok := err.(beanstalk.ConnError)
if ok && cerr.Err == beanstalk.ErrTimeout {
continue
} else {
panic(fmt.Sprintf("conn err: %s", err))
}
}
// unmarshal the work payload
job := Request{}
err = json.Unmarshal(msg, &job)
if err != nil {
beanConn.Delete(id)
continue
}
job.id = id
job.host = w.options.Host
jobCnt++
go func(j Request) {
// send it off!
jobs <- j
}(job)
}
}
func QueueWatch(conn *beanstalk.Conn, name ...string) *beanstalk.TubeSet {
return beanstalk.NewTubeSet(conn, name...)
}
// pushJobs push the reserved job to worker
func (p *beanWorkerHandler) pushJobs(w Worker) {
// create tubeset for topic
p.conn = mustCreateConnection(p.c, p.address)
tubeset := beanstalk.NewTubeSet(p.conn, string(p.jobType))
// create value context to store the reference to broker itself
ctx := context.WithValue(p.c, BrokerKey, p.broker)
for {
select {
// watch for close signal
case <-p.c.Done():
return
default:
// get the job
id, body, err := tubeset.Reserve(p.reservationTimeout)
// if everything is fine
if err == nil {
// pass it to a worker and evaluate the response value
switch w.Do(ctx, &Job{id, body}) {
case Delete:
p.conn.Delete(id)
case Bury:
// gets the current priority of the task
stats, err := p.conn.StatsJob(id)
if err != nil {
log.Println("worker:", err.Error())
continue
}
pri, err := strconv.Atoi(stats["pri"])
if err != nil {
log.Println("worker:", err.Error())
continue
}
// bury it with its original priority
p.conn.Bury(id, uint32(pri))
case Touch:
p.conn.Touch(id)
case Release:
// supports only immediate release
p.conn.Release(id, 1, time.Second)
}
continue
}
if err.(beanstalk.ConnError).Err == beanstalk.ErrTimeout {
continue
} else if err.(beanstalk.ConnError).Err == beanstalk.ErrDeadline {
time.Sleep(time.Second)
continue
} else if possibleNetworkError(err.(beanstalk.ConnError).Err) {
// try reconnecting
p.conn = mustCreateConnection(p.c, p.address)
tubeset = beanstalk.NewTubeSet(p.conn, string(p.jobType))
} else {
log.Println("beanbroker:", err)
}
}
}
}
func (q *BeanstalkQueue) NewTubeSet(name ...string) TubeSet {
return beanstalk.NewTubeSet(q.Conn, name...)
}
// Send a unit of work to a worker. 'workerTube' determines the
// tube that will respond to incoming work. 'feedback' is lets the caller
// determine if a response is expected. 'options' is an optional parameter
// to configure the beanstalkd interaction, otherwise, the default options
// will be used.
func Send(tube string, data interface{}, feedback bool, options *Options) ([]byte, error) {
if options == nil {
options = defaultOptions
}
// put together our request data
req := &struct {
Data interface{} `json:"data"`
Feedback bool `json:"feedback"`
}{
Data: data,
Feedback: feedback,
}
// marshal the data into a payload
jsonReq, err := json.Marshal(req)
if err != nil {
return nil, ErrJSONMarshal
}
// connect to beanstalkd
beanConn, err := beanstalk.Dial("tcp", options.Host)
if err != nil {
return nil, ErrBeanstalkConnect
}
defer beanConn.Close()
// configure conn for send tube
workerTube := beanstalk.Tube{beanConn, tube}
// send it
jobID, err := workerTube.Put(jsonReq, options.Priority, options.Delay, options.TTR)
if err != nil {
return nil, ErrUnableToSend
}
// no response -- all done with the send
if !feedback {
return nil, nil
}
var (
resTube = tube + "_" + strconv.FormatUint(jobID, 10)
watch = beanstalk.NewTubeSet(beanConn, resTube)
id uint64
msg []byte
start = time.Now()
)
// wait for a response from the worker
for {
id, msg, err = watch.Reserve(options.Reserve)
if err != nil {
cerr, ok := err.(beanstalk.ConnError)
if ok && cerr.Err == beanstalk.ErrTimeout {
if time.Since(start) > options.Wait {
return nil, ErrNoResponse
}
continue
} else {
return nil, err
}
}
break
}
// delete the job
beanConn.Delete(id)
// handle the response
resp := &struct {
Error string `json:"error"`
Data json.RawMessage `json:"data"`
}{}
err = json.Unmarshal(msg, resp)
if err != nil {
return nil, err
} else if len(resp.Error) > 0 {
return nil, errors.New(resp.Error)
}
// success!
return resp.Data, nil
}
// Run connects to beanstalkd and starts broking.
// If ticks channel is present, one job is processed per tick.
func (b *Broker) Run(ticks chan bool) {
b.log.Println("command:", b.Cmd)
b.log.Println("connecting to", b.Address)
conn, err := beanstalk.Dial("tcp", b.Address)
if err != nil {
panic(err)
}
b.log.Println("watching", b.Tube)
ts := beanstalk.NewTubeSet(conn, b.Tube)
for {
if ticks != nil {
if _, ok := <-ticks; !ok {
break
}
}
b.log.Println("reserve (waiting for job)")
id, body := bs.MustReserveWithoutTimeout(ts)
job := bs.NewJob(id, body, conn)
t, err := job.Timeouts()
if err != nil {
b.log.Panic(err)
}
if t >= TimeoutTries {
b.log.Printf("job %d has %d timeouts, burying", job.Id, t)
job.Bury()
if b.results != nil {
b.results <- &JobResult{JobId: job.Id, Buried: true}
}
continue
}
releases, err := job.Releases()
if err != nil {
b.log.Panic(err)
}
if releases >= ReleaseTries {
b.log.Printf("job %d has %d releases, burying", job.Id, releases)
job.Bury()
if b.results != nil {
b.results <- &JobResult{JobId: job.Id, Buried: true}
}
continue
}
b.log.Printf("executing job %d", job.Id)
result, err := b.executeJob(job, b.Cmd)
if err != nil {
log.Panic(err)
}
err = b.handleResult(job, result)
if err != nil {
log.Panic(err)
}
if result.Error != nil {
b.log.Println("result had error:", result.Error)
}
if b.results != nil {
b.results <- result
}
}
b.log.Println("broker finished")
}
