mirror of https://github.com/usememos/memos
fix: reduce list memo query overhead (#5880)
boojack
1 month ago
committed by
GitHub
parent
ebc0e10f07
commit
50638040f6
@ -0,0 +1,282 @@
|
||||
package test
|
||||
|
||||
import (
|
||||
"context"
|
||||
"fmt"
|
||||
"path/filepath"
|
||||
"testing"
|
||||
|
||||
"github.com/usememos/memos/internal/profile"
|
||||
"github.com/usememos/memos/internal/version"
|
||||
v1pb "github.com/usememos/memos/proto/gen/api/v1"
|
||||
"github.com/usememos/memos/server/auth"
|
||||
apiv1 "github.com/usememos/memos/server/router/api/v1"
|
||||
"github.com/usememos/memos/store"
|
||||
"github.com/usememos/memos/store/db"
|
||||
)
|
||||
|
||||
const (
|
||||
benchmarkTopLevelMemoCount = 5000
|
||||
benchmarkPageSize = 16
|
||||
)
|
||||
|
||||
type benchmarkService struct {
|
||||
*TestService
|
||||
hostUser *store.User
|
||||
authenticatedCtx context.Context
|
||||
publicCtx context.Context
|
||||
pageTenToken string
|
||||
commentParentName string
|
||||
}
|
||||
|
||||
func newBenchmarkService(tb testing.TB) *benchmarkService {
|
||||
tb.Helper()
|
||||
|
||||
ctx := context.Background()
|
||||
testService := newTestingServiceForTB(tb)
|
||||
|
||||
hostUser, err := testService.CreateHostUser(ctx, "bench-host")
|
||||
if err != nil {
|
||||
tb.Fatalf("failed to create host user: %v", err)
|
||||
}
|
||||
|
||||
commentParentName, err := seedListMemosBenchmarkData(ctx, testService.Store, hostUser)
|
||||
if err != nil {
|
||||
tb.Fatalf("failed to seed benchmark data: %v", err)
|
||||
}
|
||||
|
||||
authenticatedCtx := context.WithValue(context.Background(), auth.UserIDContextKey, hostUser.ID)
|
||||
pageTenToken, err := getListMemosPageToken(authenticatedCtx, testService.Service, 10, benchmarkPageSize)
|
||||
if err != nil {
|
||||
tb.Fatalf("failed to build page token: %v", err)
|
||||
}
|
||||
|
||||
return &benchmarkService{
|
||||
TestService: testService,
|
||||
hostUser: hostUser,
|
||||
authenticatedCtx: authenticatedCtx,
|
||||
publicCtx: context.Background(),
|
||||
pageTenToken: pageTenToken,
|
||||
commentParentName: commentParentName,
|
||||
}
|
||||
}
|
||||
|
||||
func newTestingServiceForTB(tb testing.TB) *TestService {
|
||||
tb.Helper()
|
||||
|
||||
ctx := context.Background()
|
||||
dataDir := tb.TempDir()
|
||||
testProfile := getBenchmarkProfile(dataDir)
|
||||
dbDriver, err := db.NewDBDriver(testProfile)
|
||||
if err != nil {
|
||||
tb.Fatalf("failed to create db driver: %v", err)
|
||||
}
|
||||
|
||||
testStore := store.New(dbDriver, testProfile)
|
||||
if err := testStore.Migrate(ctx); err != nil {
|
||||
tb.Fatalf("failed to migrate db: %v", err)
|
||||
}
|
||||
tb.Cleanup(func() {
|
||||
testStore.Close()
|
||||
})
|
||||
|
||||
service := newServiceWithProfile(testProfile, testStore)
|
||||
return &TestService{
|
||||
Service: service,
|
||||
Store: testStore,
|
||||
Profile: testProfile,
|
||||
Secret: service.Secret,
|
||||
}
|
||||
}
|
||||
|
||||
func getBenchmarkProfile(dataDir string) *profile.Profile {
|
||||
return &profile.Profile{
|
||||
Demo: true,
|
||||
Version: version.GetCurrentVersion(),
|
||||
InstanceURL: "http://localhost:8080",
|
||||
Driver: "sqlite",
|
||||
DSN: filepath.Join(dataDir, "bench.db"),
|
||||
Data: dataDir,
|
||||
}
|
||||
}
|
||||
|
||||
func newServiceWithProfile(testProfile *profile.Profile, testStore *store.Store) *apiv1.APIV1Service {
|
||||
service := apiv1.NewAPIV1Service("bench-secret", testProfile, testStore)
|
||||
return service
|
||||
}
|
||||
|
||||
func seedListMemosBenchmarkData(ctx context.Context, stores *store.Store, hostUser *store.User) (string, error) {
|
||||
topLevelMemos := make([]*store.Memo, 0, benchmarkTopLevelMemoCount)
|
||||
commentParentName := ""
|
||||
|
||||
for i := 0; i < benchmarkTopLevelMemoCount; i++ {
|
||||
visibility := store.Private
|
||||
if i%4 == 0 {
|
||||
visibility = store.Public
|
||||
}
|
||||
memo, err := stores.CreateMemo(ctx, &store.Memo{
|
||||
UID: fmt.Sprintf("memo-%06d", i),
|
||||
CreatorID: hostUser.ID,
|
||||
Content: benchmarkMemoContent(i),
|
||||
Visibility: visibility,
|
||||
})
|
||||
if err != nil {
|
||||
return "", err
|
||||
}
|
||||
topLevelMemos = append(topLevelMemos, memo)
|
||||
|
||||
if i%3 == 0 {
|
||||
if _, err := stores.CreateAttachment(ctx, &store.Attachment{
|
||||
UID: fmt.Sprintf("att-%06d", i),
|
||||
CreatorID: hostUser.ID,
|
||||
Filename: fmt.Sprintf("memo-%06d.png", i),
|
||||
Type: "image/png",
|
||||
Size: 2048,
|
||||
MemoID: &memo.ID,
|
||||
}); err != nil {
|
||||
return "", err
|
||||
}
|
||||
}
|
||||
|
||||
if i%5 == 0 {
|
||||
if _, err := stores.UpsertReaction(ctx, &store.Reaction{
|
||||
CreatorID: hostUser.ID,
|
||||
ContentID: "memos/" + memo.UID,
|
||||
ReactionType: "thumbs-up",
|
||||
}); err != nil {
|
||||
return "", err
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
for i, memo := range topLevelMemos {
|
||||
if i+1 < len(topLevelMemos) && i%4 == 0 {
|
||||
if _, err := stores.UpsertMemoRelation(ctx, &store.MemoRelation{
|
||||
MemoID: memo.ID,
|
||||
RelatedMemoID: topLevelMemos[i+1].ID,
|
||||
Type: store.MemoRelationReference,
|
||||
}); err != nil {
|
||||
return "", err
|
||||
}
|
||||
}
|
||||
|
||||
if i%6 == 0 {
|
||||
commentMemo, err := stores.CreateMemo(ctx, &store.Memo{
|
||||
UID: fmt.Sprintf("comment-%06d", i),
|
||||
CreatorID: hostUser.ID,
|
||||
Content: fmt.Sprintf("Comment for memo %06d", i),
|
||||
Visibility: store.Private,
|
||||
})
|
||||
if err != nil {
|
||||
return "", err
|
||||
}
|
||||
if _, err := stores.UpsertMemoRelation(ctx, &store.MemoRelation{
|
||||
MemoID: commentMemo.ID,
|
||||
RelatedMemoID: memo.ID,
|
||||
Type: store.MemoRelationComment,
|
||||
}); err != nil {
|
||||
return "", err
|
||||
}
|
||||
if commentParentName == "" {
|
||||
commentParentName = "memos/" + memo.UID
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
return commentParentName, nil
|
||||
}
|
||||
|
||||
func benchmarkMemoContent(i int) string {
|
||||
return fmt.Sprintf("# Bench Memo %06d\n\nThis is benchmark memo %06d with enough content to exercise snippet generation.\n\n- task one\n- task two\n", i, i)
|
||||
}
|
||||
|
||||
func getListMemosPageToken(ctx context.Context, service *apiv1.APIV1Service, page int, pageSize int32) (string, error) {
|
||||
pageToken := ""
|
||||
for range page - 1 {
|
||||
resp, err := service.ListMemos(ctx, &v1pb.ListMemosRequest{
|
||||
PageSize: pageSize,
|
||||
PageToken: pageToken,
|
||||
})
|
||||
if err != nil {
|
||||
return "", err
|
||||
}
|
||||
pageToken = resp.NextPageToken
|
||||
if pageToken == "" {
|
||||
break
|
||||
}
|
||||
}
|
||||
return pageToken, nil
|
||||
}
|
||||
|
||||
func BenchmarkListMemos(b *testing.B) {
|
||||
bench := newBenchmarkService(b)
|
||||
|
||||
b.Run("authenticated_first_page", func(b *testing.B) {
|
||||
req := &v1pb.ListMemosRequest{PageSize: benchmarkPageSize}
|
||||
b.ReportAllocs()
|
||||
b.ResetTimer()
|
||||
for i := 0; i < b.N; i++ {
|
||||
resp, err := bench.Service.ListMemos(bench.authenticatedCtx, req)
|
||||
if err != nil {
|
||||
b.Fatalf("ListMemos failed: %v", err)
|
||||
}
|
||||
if len(resp.Memos) == 0 {
|
||||
b.Fatal("expected memos in authenticated benchmark response")
|
||||
}
|
||||
}
|
||||
})
|
||||
|
||||
b.Run("authenticated_page_ten", func(b *testing.B) {
|
||||
req := &v1pb.ListMemosRequest{PageSize: benchmarkPageSize, PageToken: bench.pageTenToken}
|
||||
b.ReportAllocs()
|
||||
b.ResetTimer()
|
||||
for i := 0; i < b.N; i++ {
|
||||
resp, err := bench.Service.ListMemos(bench.authenticatedCtx, req)
|
||||
if err != nil {
|
||||
b.Fatalf("ListMemos failed: %v", err)
|
||||
}
|
||||
if len(resp.Memos) == 0 {
|
||||
b.Fatal("expected memos in paged benchmark response")
|
||||
}
|
||||
}
|
||||
})
|
||||
|
||||
b.Run("public_first_page", func(b *testing.B) {
|
||||
req := &v1pb.ListMemosRequest{PageSize: benchmarkPageSize}
|
||||
b.ReportAllocs()
|
||||
b.ResetTimer()
|
||||
for i := 0; i < b.N; i++ {
|
||||
resp, err := bench.Service.ListMemos(bench.publicCtx, req)
|
||||
if err != nil {
|
||||
b.Fatalf("ListMemos failed: %v", err)
|
||||
}
|
||||
if len(resp.Memos) == 0 {
|
||||
b.Fatal("expected memos in public benchmark response")
|
||||
}
|
||||
}
|
||||
})
|
||||
}
|
||||
|
||||
func BenchmarkListMemoCommentsPreview(b *testing.B) {
|
||||
bench := newBenchmarkService(b)
|
||||
if bench.commentParentName == "" {
|
||||
b.Fatal("expected seeded memo with comments")
|
||||
}
|
||||
|
||||
req := &v1pb.ListMemoCommentsRequest{
|
||||
Name: bench.commentParentName,
|
||||
PageSize: 3,
|
||||
}
|
||||
|
||||
b.ReportAllocs()
|
||||
b.ResetTimer()
|
||||
for i := 0; i < b.N; i++ {
|
||||
resp, err := bench.Service.ListMemoComments(bench.authenticatedCtx, req)
|
||||
if err != nil {
|
||||
b.Fatalf("ListMemoComments failed: %v", err)
|
||||
}
|
||||
if len(resp.Memos) == 0 {
|
||||
b.Fatal("expected comments in benchmark response")
|
||||
}
|
||||
}
|
||||
}
|
||||
Loading…
Reference in New Issue