Warp/server/internal/pushserver/worker.go

package pushserver

import (
	"context"
	"encoding/json"

	"github.com/dariuszjaniak/warp/server/internal/matrix/event"
	"github.com/dariuszjaniak/warp/server/internal/matrix/id"
	"github.com/dariuszjaniak/warp/server/internal/pushrules"
	"github.com/dariuszjaniak/warp/server/internal/roomserver"
	"github.com/dariuszjaniak/warp/server/internal/storage"
)

// EventPersisted implements roomserver.Notifier. It only queues the event:
// see the package comment for why nothing is evaluated here.
func (a *API) EventPersisted(_ context.Context, n roomserver.Notification) {
	if n.Event == nil {
		return
	}
	a.mu.Lock()
	if a.ctx.Err() != nil {
		a.mu.Unlock()
		return
	}
	a.incoming = append(a.incoming, n)
	a.mu.Unlock()
	select {
	case a.wake <- struct{}{}:
	default:
	}
}

// WaitIdle blocks until every event queued so far has been evaluated, or
// ctx ends. Tests use it to look at counts without sleeping, and it is the
// honest way for anything else to ask "has the push server caught up".
func (a *API) WaitIdle(ctx context.Context) error {
	stop := context.AfterFunc(ctx, func() {
		a.mu.Lock()
		a.idle.Broadcast()
		a.mu.Unlock()
	})
	defer stop()

	a.mu.Lock()
	defer a.mu.Unlock()
	for len(a.incoming) > 0 || a.busy {
		if err := ctx.Err(); err != nil {
			return err
		}
		if a.ctx.Err() != nil {
			return a.ctx.Err()
		}
		a.idle.Wait()
	}
	return nil
}

// work drains the event list, one event at a time and in the order they
// were persisted, so that a user's notifications are numbered in the
// order their events happened.
func (a *API) work() {
	defer a.wg.Done()
	for {
		select {
		case <-a.ctx.Done():
			return
		case <-a.wake:
		}
		for {
			a.mu.Lock()
			if len(a.incoming) == 0 {
				a.busy = false
				a.idle.Broadcast()
				a.mu.Unlock()
				break
			}
			batch := a.incoming
			a.incoming = nil
			a.busy = true
			a.mu.Unlock()

			for _, n := range batch {
				if a.ctx.Err() != nil {
					return
				}
				if err := a.process(a.ctx, n); err != nil && a.ctx.Err() == nil {
					a.log.Error("evaluating push rules failed",
						"room_id", n.RoomID, "event_id", n.Event.EventID(), "error", err)
				}
			}
		}
	}
}

// process evaluates one event for every local user it may concern.
func (a *API) process(ctx context.Context, n roomserver.Notification) error {
	e := n.Event
	sender := e.Sender()

	// A local user leaving a room has nothing more to read in it, and a
	// count left behind there would sit in their badge for good. The event
	// is still evaluated for everybody else below.
	if e.Type() == event.TypeMember {
		if target, ok := e.StateKey(); ok && a.isLocal(target) {
			switch membership, _ := e.Membership(); membership {
			case event.MembershipLeave, event.MembershipBan:
				if err := a.markRead(ctx, target, n.RoomID, n.StreamOrdering); err != nil {
					return err
				}
			}
		}
	}

	room, err := a.rooms.PushRoomFor(ctx, n.RoomID)
	if err != nil {
		return err
	}

	// Who the event is for: the room's local members, and the user an
	// invite is addressed to, who is not a member yet. Never the sender —
	// nobody is notified of what they did themselves.
	var candidates []string
	for userID := range room.Members {
		if userID != sender && a.isLocal(userID) {
			candidates = append(candidates, userID)
		}
	}
	if e.Type() == event.TypeMember {
		if target, ok := e.StateKey(); ok && target != sender && a.isLocal(target) {
			if _, member := room.Members[target]; !member {
				candidates = append(candidates, target)
			}
		}
	}
	if len(candidates) == 0 {
		return nil
	}

	ev, err := pushrules.NewEvent(e.JSON())
	if err != nil {
		return err
	}

	var (
		notifications []storage.NewNotification
		outcomes      = map[string]pushrules.Outcome{}
	)
	now := a.now().UnixMilli()
	for _, userID := range candidates {
		localpart, _, err := id.SplitUserID(userID)
		if err != nil {
			continue
		}
		rs, err := a.Ruleset(ctx, localpart, userID)
		if err != nil {
			return err
		}
		rule, _, ok := rs.Evaluate(ev, pushrules.Context{
			UserID:      userID,
			DisplayName: room.Members[userID],
			MemberCount: len(room.Members),
			SenderMayNotify: func(key string) bool {
				return room.SenderMayNotify(sender, key)
			},
		})
		if !ok {
			continue
		}
		outcome := pushrules.ParseActions(rule.Actions)
		if !outcome.Notify {
			continue
		}
		actions, err := json.Marshal(rule.Actions)
		if err != nil {
			return err
		}
		notifications = append(notifications, storage.NewNotification{
			UserID:      userID,
			RoomID:      n.RoomID,
			EventID:     e.EventID(),
			EventPos:    n.StreamOrdering,
			TimestampMS: now,
			Highlight:   outcome.Highlight,
			Actions:     actions,
		})
		outcomes[userID] = outcome
	}
	if len(notifications) == 0 {
		return nil
	}

	if _, err := a.store.AddNotifications(ctx, notifications); err != nil {
		return err
	}
	notified := make([]string, 0, len(notifications))
	for _, nn := range notifications {
		notified = append(notified, nn.UserID)
	}
	a.notifyUsers(notified...)

	if a.pushEnabled {
		for _, userID := range notified {
			a.push(ctx, userID, room, n, outcomes[userID])
		}
	}
	return nil
}

// push hands one notification to the gateways of a user's pushers.
func (a *API) push(ctx context.Context, userID string, room roomserver.PushRoom, n roomserver.Notification, outcome pushrules.Outcome) {
	localpart, _, err := id.SplitUserID(userID)
	if err != nil {
		return
	}
	pushers, err := a.store.Pushers(ctx, localpart)
	if err != nil {
		a.log.Warn("reading pushers failed", "user_id", userID, "error", err)
		return
	}
	if len(pushers) == 0 {
		return
	}
	unread, err := a.store.UnreadNotifications(ctx, userID)
	if err != nil {
		a.log.Warn("counting unread notifications failed", "user_id", userID, "error", err)
	}

	e := n.Event
	note := notificationPush{
		EventID:           e.EventID(),
		RoomID:            n.RoomID,
		Type:              e.Type(),
		Sender:            e.Sender(),
		SenderDisplayName: room.Members[e.Sender()],
		RoomName:          room.Name,
		RoomAlias:         room.CanonicalAlias,
		Content:           e.Content(),
		Counts:            pushCounts{Unread: unread},
		Prio:              priority(outcome),
	}
	if stateKey, ok := e.StateKey(); ok && e.Type() == event.TypeMember {
		isTarget := stateKey == userID
		note.UserIsTarget = &isTarget
		if membership, err := e.Membership(); err == nil {
			note.Membership = membership
		}
	}
	for _, p := range pushers {
		if p.Kind != PusherKindHTTP {
			continue
		}
		a.gateway.enqueue(pushJob{
			localpart: localpart,
			pusher:    p,
			note:      note,
			tweaks:    outcome.Tweaks,
		})
	}
}

// priority is a push's prio: high for anything that makes a sound or is a
// highlight, which is what a phone should wake its user for, and low for
// the rest, which a gateway may batch to save the battery.
func priority(o pushrules.Outcome) string {
	if o.Highlight {
		return "high"
	}
	if _, ok := o.Tweaks[pushrules.TweakSound]; ok {
		return "high"
	}
	return "low"
}

// isLocal reports whether a user ID belongs to this server.
func (a *API) isLocal(userID string) bool {
	server, err := id.ServerNameOf(userID)
	return err == nil && server == a.serverName
}