Ejemplo n.º 1
0
func (c *Client) run() {
	for {
		d := make([]byte, protocol.MessageSize)
		if _, _, err := c.socket.ReadFrom(d); err != nil {
			c.err = err
			return
		}

		m, err := protocol.Parse(d)
		if err != nil {
			continue
		}

		c.m.RLock()
		ch, ok := c.pending[m.GetKey()]
		c.m.RUnlock()
		if ok {
			ch <- m
		}
	}
}
Ejemplo n.º 2
0
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
				}))
			}
		}()
	}
}