anian
283 lines
7.3 KiB
Go
283 lines
7.3 KiB
Go
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")
|
|
}
|
|
}
|
|
}
|