Add message stats to reporting (#2748)

Since we're now listening on the `OutputRoomEvent` stream, we are able
to store messages stats.

userapi/consumers/roomserver_test.go

@@ -2,7 +2,10 @@ package consumers
 
 import (
 	"context"
+	"reflect"
+	"sync"
 	"testing"
+	"time"
 
 	"github.com/matrix-org/gomatrixserverlib"
 	"github.com/stretchr/testify/assert"
@@ -12,6 +15,7 @@ import (
 	"github.com/matrix-org/dendrite/test"
 	"github.com/matrix-org/dendrite/test/testrig"
 	"github.com/matrix-org/dendrite/userapi/storage"
+	userAPITypes "github.com/matrix-org/dendrite/userapi/types"
 )
 
 func mustCreateDatabase(t *testing.T, dbType test.DBType) (storage.Database, func()) {
@@ -132,3 +136,122 @@ func Test_evaluatePushRules(t *testing.T) {
 		}
 	})
 }
+
+func TestMessageStats(t *testing.T) {
+	type args struct {
+		eventType   string
+		eventSender string
+		roomID      string
+	}
+	tests := []struct {
+		name           string
+		args           args
+		ourServer      gomatrixserverlib.ServerName
+		lastUpdate     time.Time
+		initRoomCounts map[gomatrixserverlib.ServerName]map[string]bool
+		wantStats      userAPITypes.MessageStats
+	}{
+		{
+			name:      "m.room.create does not count as a message",
+			ourServer: "localhost",
+			args: args{
+				eventType:   "m.room.create",
+				eventSender: "@alice:localhost",
+			},
+		},
+		{
+			name:      "our server - message",
+			ourServer: "localhost",
+			args: args{
+				eventType:   "m.room.message",
+				eventSender: "@alice:localhost",
+				roomID:      "normalRoom",
+			},
+			wantStats: userAPITypes.MessageStats{Messages: 1, SentMessages: 1},
+		},
+		{
+			name:      "our server - E2EE message",
+			ourServer: "localhost",
+			args: args{
+				eventType:   "m.room.encrypted",
+				eventSender: "@alice:localhost",
+				roomID:      "encryptedRoom",
+			},
+			wantStats: userAPITypes.MessageStats{Messages: 1, SentMessages: 1, MessagesE2EE: 1, SentMessagesE2EE: 1},
+		},
+
+		{
+			name:      "remote server - message",
+			ourServer: "localhost",
+			args: args{
+				eventType:   "m.room.message",
+				eventSender: "@alice:remote",
+				roomID:      "normalRoom",
+			},
+			wantStats: userAPITypes.MessageStats{Messages: 2, SentMessages: 1, MessagesE2EE: 1, SentMessagesE2EE: 1},
+		},
+		{
+			name:      "remote server - E2EE message",
+			ourServer: "localhost",
+			args: args{
+				eventType:   "m.room.encrypted",
+				eventSender: "@alice:remote",
+				roomID:      "encryptedRoom",
+			},
+			wantStats: userAPITypes.MessageStats{Messages: 2, SentMessages: 1, MessagesE2EE: 2, SentMessagesE2EE: 1},
+		},
+		{
+			name:       "day change creates a new room map",
+			ourServer:  "localhost",
+			lastUpdate: time.Now().Add(-time.Hour * 24),
+			initRoomCounts: map[gomatrixserverlib.ServerName]map[string]bool{
+				"localhost": {"encryptedRoom": true},
+			},
+			args: args{
+				eventType:   "m.room.encrypted",
+				eventSender: "@alice:remote",
+				roomID:      "someOtherRoom",
+			},
+			wantStats: userAPITypes.MessageStats{Messages: 2, SentMessages: 1, MessagesE2EE: 3, SentMessagesE2EE: 1},
+		},
+	}
+
+	test.WithAllDatabases(t, func(t *testing.T, dbType test.DBType) {
+		db, close := mustCreateDatabase(t, dbType)
+		defer close()
+
+		for _, tt := range tests {
+			t.Run(tt.name, func(t *testing.T) {
+				if tt.lastUpdate.IsZero() {
+					tt.lastUpdate = time.Now()
+				}
+				if tt.initRoomCounts == nil {
+					tt.initRoomCounts = map[gomatrixserverlib.ServerName]map[string]bool{}
+				}
+				s := &OutputRoomEventConsumer{
+					db:         db,
+					msgCounts:  map[gomatrixserverlib.ServerName]userAPITypes.MessageStats{},
+					roomCounts: tt.initRoomCounts,
+					countsLock: sync.Mutex{},
+					lastUpdate: tt.lastUpdate,
+					serverName: tt.ourServer,
+				}
+				s.storeMessageStats(context.Background(), tt.args.eventType, tt.args.eventSender, tt.args.roomID)
+				t.Logf("%+v", s.roomCounts)
+				gotStats, activeRooms, activeE2EERooms, err := db.DailyRoomsMessages(context.Background(), tt.ourServer)
+				if err != nil {
+					t.Fatalf("unexpected error: %s", err)
+				}
+				if !reflect.DeepEqual(gotStats, tt.wantStats) {
+					t.Fatalf("expected %+v, got %+v", tt.wantStats, gotStats)
+				}
+				if tt.args.eventType == "m.room.encrypted" && activeE2EERooms != 1 {
+					t.Fatalf("expected room to be activeE2EE")
+				}
+				if tt.args.eventType == "m.room.message" && activeRooms != 1 {
+					t.Fatalf("expected room to be active")
+				}
+			})
+		}
+	})
+}