func Serve(ctx scope.Context, addr string) {
http.Handle("/metrics", prometheus.Handler())
listener, err := net.Listen("tcp", addr)
if err != nil {
ctx.Terminate(err)
}
closed := false
m := sync.Mutex{}
closeListener := func() {
m.Lock()
if !closed {
listener.Close()
closed = true
}
m.Unlock()
}
// Spin off goroutine to watch ctx and close listener if shutdown requested.
go func() {
<-ctx.Done()
closeListener()
}()
if err := http.Serve(listener, nil); err != nil {
fmt.Printf("http[%s]: %s\n", addr, err)
ctx.Terminate(err)
}
closeListener()
ctx.WaitGroup().Done()
}
func (cmd *serveEmbedCmd) run(ctx scope.Context, args []string) error {
listener, err := net.Listen("tcp", cmd.addr)
if err != nil {
return err
}
closed := false
m := sync.Mutex{}
closeListener := func() {
m.Lock()
if !closed {
listener.Close()
closed = true
}
m.Unlock()
}
// Spin off goroutine to watch ctx and close listener if shutdown requested.
go func() {
<-ctx.Done()
closeListener()
}()
if err := http.Serve(listener, cmd); err != nil {
fmt.Printf("http[%s]: %s\n", cmd.addr, err)
return err
}
closeListener()
ctx.WaitGroup().Done()
return ctx.Err()
}
func (ctrl *Controller) terminal(ctx scope.Context, ch ssh.Channel) {
defer ch.Close()
lines := make(chan string)
term := terminal.NewTerminal(ch, "> ")
go func() {
for ctx.Err() == nil {
line, err := term.ReadLine()
if err != nil {
ctx.Terminate(err)
return
}
lines <- line
}
}()
for {
var line string
select {
case <-ctx.Done():
return
case line = <-lines:
}
cmd := parse(line)
fmt.Printf("[control] > %v\n", cmd)
switch cmd[0] {
case "":
continue
case "quit":
return
case "shutdown":
ctrl.ctx.Terminate(fmt.Errorf("shutdown initiated from console"))
default:
runCommand(ctx.Fork(), ctrl, cmd[0], term, cmd[1:])
}
}
}
func (cmd *retentionCmd) run(ctx scope.Context, args []string) error {
heim, b, err := getHeimWithPsqlBackend(ctx)
if err != nil {
return err
}
defer func() {
ctx.Cancel()
ctx.WaitGroup().Wait()
heim.Backend.Close()
}()
// start metrics server
ctx.WaitGroup().Add(1)
go retention.Serve(ctx, cmd.addr)
// start metrics scanner
ctx.WaitGroup().Add(1)
go retention.ExpiredScanLoop(ctx, heim.Cluster, b, cmd.interval)
// start delete scanner
ctx.WaitGroup().Add(1)
retention.DeleteScanLoop(ctx, heim.Cluster, b, cmd.interval)
return nil
}
func (c *Controller) background(ctx scope.Context) {
defer ctx.WaitGroup().Done()
var lastStatCheck time.Time
for {
logging.Logger(ctx).Printf("[%s] background loop", c.w.QueueName())
if time.Now().Sub(lastStatCheck) > StatsInterval {
logging.Logger(ctx).Printf("[%s] collecting stats", c.w.QueueName())
stats, err := c.jq.Stats(ctx)
if err != nil {
logging.Logger(ctx).Printf("error: %s stats: %s", c.w.QueueName(), err)
return
}
lastStatCheck = time.Now()
labels := map[string]string{"queue": c.w.QueueName()}
claimedGauge.With(labels).Set(float64(stats.Claimed))
dueGauge.With(labels).Set(float64(stats.Due))
waitingGauge.With(labels).Set(float64(stats.Waiting))
}
if err := c.processOne(ctx); err != nil {
// TODO: retry a couple times before giving up
logging.Logger(ctx).Printf("error: %s: %s", c.w.QueueName(), err)
return
}
}
}
func (c *Controller) processOne(ctx scope.Context) error {
job, err := c.claimOrSteal(ctx.ForkWithTimeout(StatsInterval))
if err != nil {
if err == scope.TimedOut {
return nil
}
return err
}
if job.Type != jobs.EmailJobType {
return jobs.ErrInvalidJobType
}
payload, err := job.Payload()
if err != nil {
return err
}
return job.Exec(ctx, func(ctx scope.Context) error {
labels := prometheus.Labels{"queue": c.jq.Name()}
defer processedCounter.With(labels).Inc()
if err := c.w.Work(ctx, job, payload); err != nil {
failedCounter.With(labels).Inc()
return err
}
completedCounter.With(labels).Inc()
return nil
})
}
// Starts a new bot, blocking until the bot finishes.
func (bcfg *BotConfig) Start(ctx scope.Context, heimDialer proto.ConnectionDialer) error {
forkedCtx := ctx.Fork()
conn, err := RetryingConnectionDialer(heimDialer).Dial(forkedCtx, bcfg.RoomName)
if err != nil {
return err
}
sessions := make(map[string]*Session)
bot := &Bot{
ctx: forkedCtx,
conn: conn,
sessions: sessions,
nick: bcfg.Nick,
}
for _, cmd := range bcfg.Cmds {
err := bot.NewSession(RetryingPluginDialer(&cmd))
if err != nil {
// If we can't make a session we just failed at start up
bot.Kill(err)
return err
}
}
go bot.forwardPacketLoop()
_ = <-bot.ctx.Done()
return bot.ctx.Err()
}
func (cmd *retentionCmd) run(ctx scope.Context, args []string) error {
c, err := getCluster(ctx)
if err != nil {
return err
}
b, err := getBackend(ctx, c)
if err != nil {
return fmt.Errorf("backend error: %s", err)
}
defer b.Close()
defer func() {
ctx.Cancel()
ctx.WaitGroup().Wait()
}()
// start metrics server
ctx.WaitGroup().Add(1)
go retention.Serve(ctx, cmd.addr)
// start metrics scanner
ctx.WaitGroup().Add(1)
go retention.ExpiredScanLoop(ctx, c, b, cmd.interval)
// start delete scanner
ctx.WaitGroup().Add(1)
retention.DeleteScanLoop(ctx, c, b, cmd.interval)
return nil
}
func controller(
ctx scope.Context, addr string, b proto.Backend, kms security.KMS, c cluster.Cluster) error {
if addr != "" {
ctrl, err := console.NewController(ctx, addr, b, kms, c)
if err != nil {
return err
}
if backend.Config.Console.HostKey != "" {
if err := ctrl.AddHostKey(backend.Config.Console.HostKey); err != nil {
return err
}
} else {
if err := ctrl.AddHostKeyFromCluster(backend.Config.Cluster.ServerID); err != nil {
return err
}
}
for _, authKey := range backend.Config.Console.AuthKeys {
if authKey == "" {
continue
}
if err := ctrl.AddAuthorizedKeys(authKey); err != nil {
return err
}
}
ctx.WaitGroup().Add(1)
go ctrl.Serve()
}
return nil
}
func (j *Job) Exec(ctx scope.Context, f func(scope.Context) error) error {
if j.JobClaim == nil {
return ErrJobNotClaimed
}
w := io.MultiWriter(os.Stdout, j)
prefix := fmt.Sprintf("[%s-%s] ", j.Queue.Name(), j.HandlerID)
deadline := time.Now().Add(j.MaxWorkDuration)
child := logging.LoggingContext(ctx.ForkWithTimeout(j.MaxWorkDuration), w, prefix)
if err := f(child); err != nil {
if err != scope.TimedOut {
delay := time.Duration(j.AttemptsMade+1) * BackoffDuration
if time.Now().Add(delay).After(deadline) {
delay = deadline.Sub(time.Now())
}
time.Sleep(delay)
}
if ferr := j.Fail(ctx, err.Error()); ferr != nil {
return ferr
}
return err
}
if err := j.Complete(ctx); err != nil {
return err
}
return nil
}
func (jq *JobQueueBinding) WaitForJob(ctx scope.Context) error {
ch := make(chan error)
// background goroutine to wait on condition
go func() {
// synchronize with caller
<-ch
jq.m.Unlock()
jq.Backend.jobQueueListener().wait(jq.Name())
ch <- nil
}()
// synchronize with background goroutine
jq.m.Lock()
ch <- nil
jq.m.Lock()
jq.m.Unlock()
select {
case <-ctx.Done():
jq.Backend.jobQueueListener().wakeAll(jq.Name())
<-ch
return ctx.Err()
case err := <-ch:
return err
}
}
func (et *EmailTracker) Send(
ctx scope.Context, js jobs.JobService, templater *templates.Templater, deliverer emails.Deliverer,
account proto.Account, to, templateName string, data interface{}) (
*emails.EmailRef, error) {
if to == "" {
to, _ = account.Email()
}
sf, err := snowflake.New()
if err != nil {
return nil, err
}
msgID := fmt.Sprintf("<%s@%s>", sf, deliverer.LocalName())
ref, err := emails.NewEmail(templater, msgID, to, templateName, data)
if err != nil {
return nil, err
}
ref.AccountID = account.ID()
jq, err := js.GetQueue(ctx, jobs.EmailQueue)
if err != nil {
return nil, err
}
payload := &jobs.EmailJob{
AccountID: account.ID(),
EmailID: ref.ID,
}
job, err := jq.AddAndClaim(ctx, jobs.EmailJobType, payload, "immediate", jobs.EmailJobOptions...)
if err != nil {
return nil, err
}
ref.JobID = job.ID
et.m.Lock()
if et.emailsByAccount == nil {
et.emailsByAccount = map[snowflake.Snowflake][]*emails.EmailRef{}
}
et.emailsByAccount[account.ID()] = append(et.emailsByAccount[account.ID()], ref)
et.m.Unlock()
child := ctx.Fork()
child.WaitGroup().Add(1)
go job.Exec(child, func(ctx scope.Context) error {
defer ctx.WaitGroup().Done()
logging.Logger(ctx).Printf("delivering to %s\n", to)
if err := deliverer.Deliver(ctx, ref); err != nil {
return err
}
return nil
})
return ref, nil
}
func (c *Client) FromRequest(ctx scope.Context, r *http.Request) {
c.UserAgent = r.Header.Get("User-Agent")
c.Connected = time.Now()
c.IP = getIP(r)
var k clientKey
ctx.Set(k, c)
}
func (cmd *serveCmd) run(ctx scope.Context, args []string) error {
listener, err := net.Listen("tcp", cmd.addr)
if err != nil {
return err
}
m := sync.Mutex{}
closed := false
closeListener := func() {
m.Lock()
if !closed {
closed = true
listener.Close()
}
m.Unlock()
}
defer closeListener()
heim, err := getHeim(ctx)
if err != nil {
return fmt.Errorf("configuration error: %s", err)
}
defer heim.Backend.Close()
if cmd.static != "" {
heim.StaticPath = cmd.static
}
if err := controller(heim, cmd.consoleAddr); err != nil {
return fmt.Errorf("controller error: %s", err)
}
serverDesc := backend.Config.Cluster.DescribeSelf()
server, err := backend.NewServer(heim, serverDesc.ID, serverDesc.Era)
if err != nil {
return fmt.Errorf("server error: %s", err)
}
server.SetInsecureCookies(backend.Config.SetInsecureCookies)
server.AllowRoomCreation(backend.Config.AllowRoomCreation)
server.NewAccountMinAgentAge(backend.Config.NewAccountMinAgentAge)
// Spin off goroutine to watch ctx and close listener if shutdown requested.
go func() {
<-ctx.Done()
closeListener()
}()
fmt.Printf("serving era %s on %s\n", serverDesc.Era, cmd.addr)
if err := http.Serve(listener, newVersioningHandler(server)); err != nil {
if strings.HasSuffix(err.Error(), "use of closed network connection") {
return nil
}
return err
}
return nil
}
func (c *Client) FromContext(ctx scope.Context) bool {
var k clientKey
src, ok := ctx.Get(k).(*Client)
if !ok || src == nil {
return false
}
*c = *src
return true
}
func WSConnectionFactory(ctx scope.Context, roomName string) proto.Connection {
return &Connection{
ctx: ctx.Fork(),
testToClient: make(chan *hproto.Packet),
clientToTest: make(chan *hproto.Packet),
incoming: make(chan *hproto.Packet),
outgoing: make(chan *hproto.Packet),
}
}
func (s *session) Send(ctx scope.Context, cmdType proto.PacketType, payload interface{}) error {
// Special case: certain events have privileged info that may need to be stripped from them
switch event := payload.(type) {
case *proto.PresenceEvent:
switch s.privilegeLevel() {
case proto.Staff:
case proto.Host:
event.RealClientAddress = ""
default:
event.RealClientAddress = ""
event.ClientAddress = ""
}
case *proto.Message:
if s.privilegeLevel() == proto.General {
event.Sender.ClientAddress = ""
}
case *proto.EditMessageEvent:
if s.privilegeLevel() == proto.General {
event.Sender.ClientAddress = ""
}
}
var err error
payload, err = proto.DecryptPayload(payload, &s.client.Authorization, s.privilegeLevel())
if err != nil {
return err
}
encoded, err := json.Marshal(payload)
if err != nil {
return err
}
cmd := &proto.Packet{
Type: cmdType,
Data: encoded,
}
// Add to outgoing channel. If channel is full, defer to goroutine so as not to block
// the caller (this may result in deliveries coming out of order).
select {
case <-ctx.Done():
// Session is closed, return error.
return ctx.Err()
case s.outgoing <- cmd:
// Packet delivered to queue.
default:
// Queue is full.
logging.Logger(s.ctx).Printf("outgoing channel full, ordering cannot be guaranteed")
go func() { s.outgoing <- cmd }()
}
return nil
}
func (d *RetryPluginDialer) Dial(ctx scope.Context) (proto.Plugin, error) {
retry := &RetryPlugin{
ctx: ctx.Fork(),
dialer: d.Dialer,
}
err := retry.start(nil)
if err != nil {
return nil, err
}
return retry, err
}
func (d *RetryConnectionDialer) Dial(ctx scope.Context, roomName string) (proto.Connection, error) {
retry := &RetryConnection{
ctx: ctx.Fork(),
dialer: d.Dialer,
roomName: roomName,
}
err := retry.start(nil)
if err != nil {
return nil, err
}
return retry, nil
}
func (d *CmdPluginDialer) Dial(ctx scope.Context) (proto.Plugin, error) {
cmdPlugin := &CmdPlugin{
ctx: ctx.Fork(),
cmd: d.Cmd,
args: d.Args,
}
err := cmdPlugin.connect()
if err != nil {
proto.GetLogger(cmdPlugin.ctx).Errorf("Failed to connect CmdPlugin: %s", err)
return nil, err
}
return cmdPlugin, nil
}
func (cmd *presenceCmd) run(ctx scope.Context, args []string) error {
c, err := getCluster(ctx)
if err != nil {
return err
}
b, err := getBackend(ctx, c)
if err != nil {
return fmt.Errorf("backend error: %s", err)
}
defer b.Close()
defer func() {
ctx.Cancel()
ctx.WaitGroup().Wait()
}()
// Start metrics server.
ctx.WaitGroup().Add(1)
go presence.Serve(ctx, cmd.addr)
// Start scanner.
ctx.WaitGroup().Add(1)
presence.ScanLoop(ctx, c, b, cmd.interval)
return nil
}
func (cmd *activityCmd) run(ctx scope.Context, args []string) error {
// Get cluster in order to load config.
_, err := getCluster(ctx)
if err != nil {
return fmt.Errorf("cluster error: %s", err)
}
listener := pq.NewListener(backend.Config.DB.DSN, 200*time.Millisecond, 5*time.Second, nil)
if err := listener.Listen("broadcast"); err != nil {
return fmt.Errorf("pq listen error: %s", err)
}
defer func() {
ctx.Cancel()
ctx.WaitGroup().Wait()
}()
// Start metrics server.
ctx.WaitGroup().Add(1)
go activity.Serve(ctx, cmd.addr)
// Start scanner.
ctx.WaitGroup().Add(1)
activity.ScanLoop(ctx, listener)
return nil
}
func (d *WSConnectionDialer) Dial(ctx scope.Context, roomName string) (proto.Connection, error) {
ws := WSConnection{
roomName: roomName,
ctx: ctx.Fork(),
incoming: make(chan *hproto.Packet, 8),
outgoing: make(chan *hproto.Packet, 8)}
dialer := websocket.Dialer{
HandshakeTimeout: 5 * time.Second,
}
url := fmt.Sprintf(d.UrlFormat, ws.roomName)
proto.GetLogger(ws.ctx).Infof("Connecting to: %s", url)
conn, _, err := dialer.Dial(url, nil)
if err != nil {
return nil, err
}
ws.conn = conn
ws.start()
return &ws, nil
}
func (jq *JobQueue) WaitForJob(ctx scope.Context) error {
ch := make(chan error)
go func() {
jq.m.Lock()
jq.c.Wait()
jq.m.Unlock()
ch <- nil
}()
select {
case <-ctx.Done():
jq.m.Lock()
jq.c.Broadcast()
jq.m.Unlock()
<-ch
return ctx.Err()
case err := <-ch:
return err
}
}
func Claim(ctx scope.Context, jq JobQueue, handlerID string, pollTime time.Duration, stealChance float64) (*Job, error) {
for ctx.Err() == nil {
if rand.Float64() < stealChance {
job, err := jq.TrySteal(ctx, handlerID)
if err != nil && err != ErrJobNotFound {
return nil, err
}
if err == nil {
return job, nil
}
}
job, err := jq.TryClaim(ctx, handlerID)
if err != nil {
if err == ErrJobNotFound {
child := ctx.ForkWithTimeout(pollTime)
if err = jq.WaitForJob(child); err != nil && err != scope.TimedOut {
return nil, err
}
continue
}
return nil, err
}
return job, nil
}
return nil, ctx.Err()
}
func (cmd *workerCmd) run(ctx scope.Context, args []string) error {
if len(args) < 1 {
fmt.Printf("Usage: %s\r\n", cmd.usage())
// TODO: list queues
return nil
}
fmt.Printf("getting config\n")
cfg, err := getConfig(ctx)
if err != nil {
return err
}
fmt.Printf("getting heim\n")
heim, err := cfg.Heim(ctx)
if err != nil {
fmt.Printf("heim error: %s\n", err)
return err
}
defer func() {
ctx.Cancel()
ctx.WaitGroup().Wait()
}()
// Start metrics server.
fmt.Printf("starting server\n")
ctx.WaitGroup().Add(1)
go worker.Serve(ctx, cmd.addr)
// Start scanner.
return worker.Loop(ctx, heim, cmd.worker, args[0])
}
func (ctrl *Controller) interact(ctx scope.Context, conn net.Conn) {
defer ctx.WaitGroup().Done()
_, nchs, reqs, err := ssh.NewServerConn(conn, ctrl.config)
if err != nil {
return
}
go ssh.DiscardRequests(reqs)
for nch := range nchs {
if nch.ChannelType() != "session" {
nch.Reject(ssh.UnknownChannelType, "unknown channel type")
continue
}
ch, reqs, err := nch.Accept()
if err != nil {
return
}
go ctrl.filterClientRequests(reqs)
go ctrl.terminal(ctx, ch)
}
}
func ScanLoop(ctx scope.Context, listener *pq.Listener) {
defer ctx.WaitGroup().Done()
logger := logging.Logger(ctx)
for {
select {
case <-ctx.Done():
logger.Printf("received cancellation signal, shutting down")
return
case notice := <-listener.Notify:
if notice == nil {
logger.Printf("received nil from listener")
continue
}
var msg psql.BroadcastMessage
if err := json.Unmarshal([]byte(notice.Extra), &msg); err != nil {
logger.Printf("error: pq listen: invalid broadcast: %s", err)
logger.Printf(" payload: %#v", notice.Extra)
continue
}
switch msg.Event.Type {
case proto.BounceEventType:
bounceActivity.WithLabelValues(msg.Room).Inc()
case proto.JoinEventType:
joinActivity.WithLabelValues(msg.Room).Inc()
case proto.PartEventType:
partActivity.WithLabelValues(msg.Room).Inc()
case proto.SendEventType:
messageActivity.WithLabelValues(msg.Room).Inc()
}
}
}
}
func Loop(ctx scope.Context, heim *proto.Heim, workerName, queueName string) error {
fmt.Printf("Loop\n")
ctrl, err := NewController(ctx, heim, workerName, queueName)
if err != nil {
fmt.Printf("error: %s\n", err)
return err
}
ctx.WaitGroup().Add(1)
go ctrl.background(ctx)
ctx.WaitGroup().Wait()
return ctx.Err()
}
func ScanLoop(ctx scope.Context, c cluster.Cluster, pb *psql.Backend, interval time.Duration) {
defer ctx.WaitGroup().Done()
errCount := 0
for {
t := time.After(interval)
select {
case <-ctx.Done():
return
case <-t:
if err := scan(ctx.Fork(), c, pb); err != nil {
errCount++
fmt.Printf("scan error [%d/%d]: %s", errCount, maxErrors, err)
if errCount > maxErrors {
fmt.Printf("maximum scan errors exceeded, terminating\n")
ctx.Terminate(fmt.Errorf("maximum scan errors exceeded"))
return
}
continue
}
errCount = 0
}
}
}
