func (c *Client) req(m protocol.Message) (protocol.Message, error) {
if c.err != nil {
return nil, c.err
}
retries := c.retries
if m.GetKey() == "" {
d := make([]byte, 8)
if _, err := io.ReadFull(rand.Reader, d); err != nil {
return nil, err
}
m.SetKey(hex.EncodeToString(d))
}
ch := make(chan protocol.Message, 1)
c.m.Lock()
c.pending[m.GetKey()] = ch
c.m.Unlock()
defer func() {
c.m.Lock()
delete(c.pending, m.GetKey())
c.m.Unlock()
}()
for {
if _, err := c.socket.Write(protocol.Serialise(m)); err != nil {
return nil, err
}
select {
case r := <-ch:
return r, nil
case <-time.After(c.timeout):
if retries == 0 {
return nil, ErrTimeout
}
retries--
}
}
}
func main() {
kingpin.MustParse(app.Parse(os.Args[1:]))
ll, lerr := logrus.ParseLevel(*logLevel)
if lerr != nil {
panic(lerr)
}
logrus.SetLevel(ll)
logrus.WithFields(logrus.Fields{
"db_path": *dbPath,
"addr": *addr,
"log_level": *logLevel,
}).Info("starting up")
logrus.WithField("db_path", *dbPath).Debug("opening database")
db, dberr := sql.Open("sqlite3", *dbPath)
if dberr != nil {
panic(dberr)
}
defer db.Close()
logrus.Debug("opened database")
logrus.Debug("ensuring tables exist")
if _, err := db.Exec(createTableQuery); err != nil {
panic(err)
}
logrus.Debug("tables created")
logrus.Debug("opening listening socket")
s, serr := net.ListenPacket("udp4", *addr)
if serr != nil {
panic(serr)
}
logrus.Info("listening")
snum := 1
for {
logrus.Debug("waiting for incoming message")
b := make([]byte, protocol.MessageSize)
n, r, err := s.ReadFrom(b)
if err != nil {
panic(err)
}
before := time.Now()
mnum := snum
snum++
l := logrus.WithField("seq", mnum)
l.WithFields(logrus.Fields{
"size": n,
"remote": r.String(),
}).Debug("got message")
func() {
defer func() {
l := l.WithField("measure#duration", time.Now().Sub(before).Seconds()*1000)
if e := recover(); e != nil {
if err, ok := e.(error); ok {
l.WithField("error", err.Error()).Error("error processing message")
} else {
l.WithField("error", e).Error("error processing message")
}
} else {
l.Debug("processed message successfully")
}
}()
m, err := protocol.Parse(bytes.TrimSpace(b[0:n]))
if err != nil {
panic(err)
}
l = l.WithField("message_key", m.GetKey())
l.WithField("message_type", fmt.Sprintf("%T", m)).Debug("processing message")
switch m := m.(type) {
case *protocol.PingMessage:
if _, err := s.WriteTo(m.Serialise(), r); err != nil {
panic(err)
}
case *protocol.JobMessage:
maybePanic(withTx(db, func(tx *sql.Tx) error {
if m.HoldUntil == 0 {
m.HoldUntil = time.Now().Unix()
}
if m.TTR == 0 {
m.TTR = uint64(time.Hour / time.Second)
}
var queue, content string
var priority float64
var holdUntil int64
var ttr uint64
var found bool
if err := tx.QueryRow(fetchJobQuery, m.ID).Scan(&queue, &priority, &holdUntil, &ttr, &content); err != nil && err != sql.ErrNoRows {
return err
} else if err == nil {
found = true
}
if found == false {
if _, err := tx.Exec(putJobQuery, m.ID, m.Queue, m.Priority, m.HoldUntil, m.TTR, m.Content); err != nil {
return err
}
d := protocol.Serialise(&protocol.SuccessMessage{Key: m.Key})
if _, err := s.WriteTo(d, r); err != nil {
return err
}
l.WithFields(logrus.Fields{
"queue": m.Queue,
"job_id": m.ID,
"measure#duration_ms": time.Now().Sub(before).Seconds() * 1000,
}).Info("created job")
} else {
if m.HoldUntil > holdUntil {
m.HoldUntil = holdUntil
}
if _, err := tx.Exec(updateJobQuery, m.Priority, m.HoldUntil, m.TTR, m.ID); err != nil {
return err
}
d := protocol.Serialise(&protocol.SuccessMessage{Key: m.Key})
if _, err := s.WriteTo(d, r); err != nil {
return err
}
l.WithFields(logrus.Fields{
"queue": m.Queue,
"job_id": m.ID,
"measure#duration_ms": time.Now().Sub(before).Seconds() * 1000,
}).Info("updated job")
}
return nil
}))
case *protocol.ReserveMessage:
maybePanic(withTx(db, func(tx *sql.Tx) error {
var id, queue, content string
var priority float64
var holdUntil int64
var ttr uint64
if err := tx.QueryRow(getTopJobQuery, m.Queue, time.Now().Unix()).Scan(&id, &queue, &priority, &holdUntil, &ttr, &content); err != nil {
if err == sql.ErrNoRows {
d := protocol.Serialise(&protocol.ErrorMessage{Key: m.Key, Reason: "empty"})
if _, werr := s.WriteTo(d, r); werr != nil {
return werr
}
return nil
}
return err
}
if _, err := tx.Exec(reserveJobQuery, time.Now().Unix(), id); err != nil {
return err
}
d := protocol.Serialise(&protocol.JobMessage{Key: m.Key, ID: id, Queue: queue, Priority: priority, HoldUntil: holdUntil, TTR: ttr, Content: content})
if _, err := s.WriteTo(d, r); err != nil {
return err
}
l.WithFields(logrus.Fields{
"queue": m.Queue,
"job_id": id,
"measure#duration_ms": time.Now().Sub(before).Seconds() * 1000,
}).Info("dispatched job")
return nil
}))
case *protocol.PeekMessage:
maybePanic(withTx(db, func(tx *sql.Tx) error {
var id, queue, content string
var priority float64
var holdUntil int64
var ttr uint64
if err := tx.QueryRow(getTopJobQuery, m.Queue, time.Now().Unix()).Scan(&id, &queue, &priority, &holdUntil, &ttr, &content); err != nil {
if err == sql.ErrNoRows {
d := protocol.Serialise(&protocol.ErrorMessage{Key: m.Key, Reason: "empty"})
if _, werr := s.WriteTo(d, r); werr != nil {
return werr
}
return nil
}
return err
}
d := protocol.Serialise(&protocol.JobMessage{Key: m.Key, ID: id, Queue: queue, Priority: priority, HoldUntil: holdUntil, TTR: ttr, Content: content})
if _, err := s.WriteTo(d, r); err != nil {
return err
}
return nil
}))
case *protocol.DeleteMessage:
maybePanic(withTx(db, func(tx *sql.Tx) error {
qr, err := db.Exec(deleteJobQuery, m.Queue, m.ID)
if err != nil {
return err
}
n, err := qr.RowsAffected()
if err != nil {
return err
}
if n == 0 {
d := protocol.Serialise(&protocol.ErrorMessage{Key: m.Key, Reason: "not found"})
if _, err := s.WriteTo(d, r); err != nil {
return err
}
} else {
d := protocol.Serialise(&protocol.SuccessMessage{Key: m.Key})
if _, err := s.WriteTo(d, r); err != nil {
return err
}
}
l.WithFields(logrus.Fields{
"queue": m.Queue,
"job_id": m.ID,
"measure#duration_ms": time.Now().Sub(before).Seconds() * 1000,
}).Info("deleted job")
return nil
}))
}
}()
}
}