feat(stats): admin instance resource statistics

2026-05-01 23:15:56 +08:00
parent cd4f28ae10
commit ea0625da45
23 changed files with 1024 additions and 83 deletions
@@ -47,6 +47,14 @@ func (s *ConnectServiceHandler) TestInstanceEmailSetting(ctx context.Context, re
 	return connect.NewResponse(resp), nil
 }

+func (s *ConnectServiceHandler) GetInstanceStats(ctx context.Context, req *connect.Request[v1pb.GetInstanceStatsRequest]) (*connect.Response[v1pb.InstanceStats], error) {
+	resp, err := s.APIV1Service.GetInstanceStats(ctx, req.Msg)
+	if err != nil {
+		return nil, convertGRPCError(err)
+	}
+	return connect.NewResponse(resp), nil
+}
+
 // AuthService
 //
 // Auth service methods need special handling for response headers (cookies).
@@ -0,0 +1,165 @@
+package v1
+
+import (
+	"context"
+	"log/slog"
+	"os"
+	"path/filepath"
+	"sync"
+	"time"
+
+	"github.com/pkg/errors"
+	"golang.org/x/sync/errgroup"
+	"google.golang.org/grpc/codes"
+	"google.golang.org/grpc/status"
+	"google.golang.org/protobuf/types/known/timestamppb"
+
+	v1pb "github.com/usememos/memos/proto/gen/api/v1"
+	"github.com/usememos/memos/store"
+)
+
+const instanceStatsCacheTTL = 60 * time.Second
+
+// instanceStatsCache is a single-value, mutex-guarded cache for InstanceStats.
+type instanceStatsCache struct {
+	mu     sync.Mutex
+	value  *v1pb.InstanceStats
+	expiry time.Time
+}
+
+func (c *instanceStatsCache) get() (*v1pb.InstanceStats, bool) {
+	c.mu.Lock()
+	defer c.mu.Unlock()
+	if c.value == nil || time.Now().After(c.expiry) {
+		return nil, false
+	}
+	return c.value, true
+}
+
+func (c *instanceStatsCache) set(v *v1pb.InstanceStats, ttl time.Duration) {
+	c.mu.Lock()
+	defer c.mu.Unlock()
+	c.value = v
+	c.expiry = time.Now().Add(ttl)
+}
+
+// GetInstanceStats returns resource usage statistics. Admin only.
+func (s *APIV1Service) GetInstanceStats(ctx context.Context, _ *v1pb.GetInstanceStatsRequest) (*v1pb.InstanceStats, error) {
+	user, err := s.fetchCurrentUser(ctx)
+	if err != nil {
+		return nil, status.Errorf(codes.Internal, "failed to get current user: %v", err)
+	}
+	if user == nil {
+		return nil, status.Errorf(codes.Unauthenticated, "user not authenticated")
+	}
+	if user.Role != store.RoleAdmin {
+		return nil, status.Errorf(codes.PermissionDenied, "permission denied")
+	}
+
+	if cached, ok := s.instanceStatsCache.get(); ok {
+		return cached, nil
+	}
+
+	stats, err := s.computeInstanceStats(ctx)
+	if err != nil {
+		return nil, status.Errorf(codes.Internal, "failed to compute instance stats: %v", err)
+	}
+	s.instanceStatsCache.set(stats, instanceStatsCacheTTL)
+	return stats, nil
+}
+
+// computeInstanceStats runs all stat subqueries in parallel and assembles the result.
+// Per-subtask failures degrade to -1 sentinel values; only a total failure (every
+// subtask errored) is propagated as an error.
+func (s *APIV1Service) computeInstanceStats(ctx context.Context) (*v1pb.InstanceStats, error) {
+	stats := &v1pb.InstanceStats{
+		Database: &v1pb.InstanceStats_DatabaseStats{
+			Driver:    s.Profile.Driver,
+			SizeBytes: -1,
+		},
+		LocalStorageBytes: -1,
+		GeneratedTime:     timestamppb.Now(),
+	}
+
+	type result struct {
+		name string
+		err  error
+	}
+	var (
+		mu      sync.Mutex
+		results []result
+		record  = func(name string, err error) {
+			mu.Lock()
+			results = append(results, result{name, err})
+			mu.Unlock()
+		}
+	)
+
+	g, gctx := errgroup.WithContext(ctx)
+
+	g.Go(func() error {
+		size, err := s.Store.GetDriver().GetDatabaseSize(gctx)
+		if err != nil {
+			record("database_size", err)
+			return nil
+		}
+		stats.Database.SizeBytes = size
+		return nil
+	})
+
+	g.Go(func() error {
+		size, err := walkLocalStorage(s.Profile.Data)
+		if err != nil {
+			record("local_storage", err)
+			return nil
+		}
+		stats.LocalStorageBytes = size
+		return nil
+	})
+
+	_ = g.Wait()
+
+	for _, r := range results {
+		slog.Warn("instance stats subtask failed", slog.String("subtask", r.name), slog.String("err", r.err.Error()))
+	}
+
+	const totalSubtasks = 2
+	if len(results) == totalSubtasks {
+		return nil, errors.New("all instance stats subtasks failed")
+	}
+	return stats, nil
+}
+
+// walkLocalStorage returns the recursive size of dir in bytes.
+// Symlinks are not followed; per-entry errors below the root are ignored
+// (the walk continues). An error accessing the root itself is returned.
+func walkLocalStorage(dir string) (int64, error) {
+	if dir == "" {
+		return -1, errors.New("empty data directory")
+	}
+	var total int64
+	err := filepath.WalkDir(dir, func(path string, entry os.DirEntry, walkErr error) error {
+		if walkErr != nil {
+			if path == dir {
+				// Root itself is inaccessible — abort the walk.
+				return walkErr
+			}
+			// Ignore per-entry errors (e.g. permission denied on a single file).
+			return nil
+		}
+		if entry.IsDir() {
+			return nil
+		}
+		info, err := entry.Info()
+		if err != nil {
+			// Ignore stat errors on individual entries; continue the walk.
+			return nil //nolint:nilerr
+		}
+		total += info.Size()
+		return nil
+	})
+	if err != nil {
+		return -1, errors.Wrap(err, "walk failed")
+	}
+	return total, nil
+}
@@ -0,0 +1,34 @@
+package v1
+
+import (
+	"os"
+	"path/filepath"
+	"testing"
+
+	"github.com/stretchr/testify/require"
+)
+
+func TestWalkLocalStorage_SumsFileSizes(t *testing.T) {
+	dir := t.TempDir()
+	require.NoError(t, os.WriteFile(filepath.Join(dir, "a.txt"), []byte("hello"), 0o600))  // 5
+	require.NoError(t, os.WriteFile(filepath.Join(dir, "b.txt"), []byte("world!"), 0o600)) // 6
+	sub := filepath.Join(dir, "sub")
+	require.NoError(t, os.Mkdir(sub, 0o700))
+	require.NoError(t, os.WriteFile(filepath.Join(sub, "c.txt"), []byte("xx"), 0o600)) // 2
+
+	size, err := walkLocalStorage(dir)
+	require.NoError(t, err)
+	require.Equal(t, int64(13), size)
+}
+
+func TestWalkLocalStorage_EmptyDir(t *testing.T) {
+	size, err := walkLocalStorage("")
+	require.Error(t, err)
+	require.Equal(t, int64(-1), size)
+}
+
+func TestWalkLocalStorage_NonexistentDir(t *testing.T) {
+	size, err := walkLocalStorage(filepath.Join(t.TempDir(), "does-not-exist"))
+	require.Error(t, err)
+	require.Equal(t, int64(-1), size)
+}
@@ -0,0 +1,72 @@
+package test
+
+import (
+	"context"
+	"testing"
+
+	"github.com/stretchr/testify/require"
+	"google.golang.org/grpc/codes"
+	"google.golang.org/grpc/status"
+
+	v1pb "github.com/usememos/memos/proto/gen/api/v1"
+)
+
+func TestGetInstanceStats_HappyPath(t *testing.T) {
+	ctx := context.Background()
+	ts := NewTestService(t)
+	defer ts.Cleanup()
+
+	admin, err := ts.CreateHostUser(ctx, "admin1")
+	require.NoError(t, err)
+	adminCtx := ts.CreateUserContext(ctx, admin.ID)
+
+	resp, err := ts.Service.GetInstanceStats(adminCtx, &v1pb.GetInstanceStatsRequest{})
+	require.NoError(t, err)
+	require.NotNil(t, resp)
+
+	require.NotNil(t, resp.Database)
+	require.Equal(t, "sqlite", resp.Database.Driver)
+	require.Greater(t, resp.Database.SizeBytes, int64(0))
+
+	require.GreaterOrEqual(t, resp.LocalStorageBytes, int64(0))
+}
+
+func TestGetInstanceStats_NonAdminDenied(t *testing.T) {
+	ctx := context.Background()
+	ts := NewTestService(t)
+	defer ts.Cleanup()
+
+	// Need an admin to exist (otherwise instance is uninitialized).
+	admin, err := ts.CreateHostUser(ctx, "admin1")
+	require.NoError(t, err)
+	_ = admin
+
+	regular, err := ts.CreateRegularUser(ctx, "alice")
+	require.NoError(t, err)
+	regularCtx := ts.CreateUserContext(ctx, regular.ID)
+
+	_, err = ts.Service.GetInstanceStats(regularCtx, &v1pb.GetInstanceStatsRequest{})
+	require.Error(t, err)
+	st, ok := status.FromError(err)
+	require.True(t, ok)
+	require.Equal(t, codes.PermissionDenied, st.Code())
+}
+
+func TestGetInstanceStats_Cache(t *testing.T) {
+	ctx := context.Background()
+	ts := NewTestService(t)
+	defer ts.Cleanup()
+
+	admin, err := ts.CreateHostUser(ctx, "admin1")
+	require.NoError(t, err)
+	adminCtx := ts.CreateUserContext(ctx, admin.ID)
+
+	first, err := ts.Service.GetInstanceStats(adminCtx, &v1pb.GetInstanceStatsRequest{})
+	require.NoError(t, err)
+
+	second, err := ts.Service.GetInstanceStats(adminCtx, &v1pb.GetInstanceStatsRequest{})
+	require.NoError(t, err)
+
+	// Cache hit: same pointer (the cache returns the stored *InstanceStats directly).
+	require.Same(t, first, second)
+}
@@ -42,6 +42,9 @@ type APIV1Service struct {
 	// thumbnailSemaphore limits concurrent thumbnail generation to prevent memory exhaustion
 	thumbnailSemaphore       *semaphore.Weighted
 	imageProcessingSemaphore *semaphore.Weighted
+
+	// instanceStatsCache memoizes GetInstanceStats results for instanceStatsCacheTTL.
+	instanceStatsCache instanceStatsCache
 }

 func NewAPIV1Service(secret string, profile *profile.Profile, store *store.Store) *APIV1Service {
@@ -72,7 +72,7 @@ func TestFrontendService_SitemapXML(t *testing.T) {

 	require.Equal(t, http.StatusOK, rec.Code)
 	require.Contains(t, rec.Header().Get("Content-Type"), "application/xml")
-	require.Contains(t, rec.Body.String(), `<loc>https://demo.usememos.com/m/publicmemo</loc>`)
+	require.Contains(t, rec.Body.String(), `<loc>https://demo.usememos.com/memos/publicmemo</loc>`)
 	require.NotContains(t, rec.Body.String(), "privatememo")
 }