server/server.go
Rama main · réplica de lectura del repositorio original.
package server
import (
"context"
"fmt"
"log/slog"
"net"
"net/http"
"os"
"path"
"sync"
"time"
"uuid"
"github.com/pkg/errors"
"github.com/usememos/memos/core/memopayload"
"github.com/usememos/memos/internal/clientip"
"github.com/usememos/memos/internal/httphandler"
"github.com/usememos/memos/internal/profile"
storepb "github.com/usememos/memos/proto/gen/store"
"github.com/usememos/memos/server/api"
"github.com/usememos/memos/server/fileserver"
"github.com/usememos/memos/server/frontend"
"github.com/usememos/memos/server/mcp"
"github.com/usememos/memos/store"
)
const (
shutdownTimeout = 10 * time.Second
// readHeaderTimeout bounds how long a client may take to send request
// headers, so idle or slow connections cannot pin a worker forever. Bodies
// and responses are not bounded here: uploads and SSE streams are
// legitimately long, and the request context still ends on disconnect.
readHeaderTimeout = 15 * time.Second
// idleTimeout closes keep-alive connections that send nothing.
idleTimeout = 2 * time.Minute
)
// healthzDatabaseTimeout bounds the database ping behind /healthz.
const healthzDatabaseTimeout = 2 * time.Second
type Server struct {
Secret string
Profile *profile.Profile
Store *store.Store
handler http.Handler
httpServer *http.Server
apiService *api.APIService
// backgroundCancel stops the startup jobs, and backgroundJobs waits for
// them, before the store closes.
backgroundCancel context.CancelFunc
backgroundJobs sync.WaitGroup
}
func NewServer(ctx context.Context, profile *profile.Profile, store *store.Store) (*Server, error) {
s := &Server{
Store: store,
Profile: profile,
}
mux := http.NewServeMux()
// Resolve the client address once per request, before anything that keys
// on it: rate limits, session records, and the file server.
clientIPResolver, err := clientip.ParseTrustedProxies(profile.TrustedProxies)
if err != nil {
return nil, errors.Wrap(err, "invalid trusted proxies")
}
instanceBasicSetting, err := s.getOrUpsertInstanceBasicSetting(ctx)
if err != nil {
return nil, errors.Wrap(err, "failed to get instance basic setting")
}
secret := "usememos"
if !profile.Demo {
secret = instanceBasicSetting.SecretKey
}
s.Secret = secret
// Register healthz endpoint. It reports ready only while the database
// answers, so a load balancer stops routing to an instance that lost it.
mux.Handle("/healthz", httphandler.GET(httphandler.Func(func(w http.ResponseWriter, r *http.Request) error {
pingCtx, cancel := context.WithTimeout(r.Context(), healthzDatabaseTimeout)
defer cancel()
if err := s.Store.GetDriver().GetDB().PingContext(pingCtx); err != nil {
slog.Warn("health check failed to reach the database", slog.String("error", err.Error()))
return httphandler.Text(w, http.StatusServiceUnavailable, "Database unavailable.")
}
return httphandler.Text(w, http.StatusOK, "Service ready.")
}), nil))
apiService := api.NewAPIService(s.Secret, profile, store)
s.apiService = apiService
apiHandler, err := apiService.NewHandler(ctx)
if err != nil {
return nil, errors.Wrap(err, "failed to create API handler")
}
frontendService := frontend.NewFrontendService(profile, store)
frontendService.Serve(ctx, mux, apiHandler)
fileServerService := fileserver.NewFileServerService(s.Profile, s.Store, s.Secret)
fileServerService.MemoExporter = apiService
fileServerService.RateLimiter = apiService.RateLimiter
fileServerService.RegisterRoutes(mux, apiHandler)
// Preserve unclean-path routing without ServeMux's automatic redirects.
fallback := frontendService.Handler(apiHandler)
router := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
frontend.SetCacheHeaders(w, r)
routingRequest := httphandler.RoutingRequest(r)
if path.Clean(routingRequest.URL.Path) != routingRequest.URL.Path && routingRequest.URL.Path != "/" {
fallback.ServeHTTP(w, r)
return
}
mux.ServeHTTP(w, routingRequest)
})
s.handler = httphandler.Recover(clientip.Middleware(clientIPResolver)(newCORSMiddleware(profile, mux)(router)))
mcpService, err := mcp.NewMCPService(profile, s.handler)
if err != nil {
return nil, errors.Wrap(err, "failed to create MCP service")
}
mcpService.RegisterRoutes(mux)
return s, nil
}
func (s *Server) Start() error {
var address, network string
if len(s.Profile.UNIXSock) == 0 {
address = fmt.Sprintf("%s:%d", s.Profile.Addr, s.Profile.Port)
network = "tcp"
} else {
address = s.Profile.UNIXSock
network = "unix"
}
listener, err := net.Listen(network, address)
if err != nil {
return errors.Wrap(err, "failed to listen")
}
if network == "unix" {
if err := os.Chmod(address, 0660); err != nil {
_ = listener.Close()
return errors.Wrap(err, "failed to chmod socket")
}
}
// All transports share the standard HTTP server.
s.httpServer = &http.Server{
Handler: s.handler,
ReadHeaderTimeout: readHeaderTimeout,
IdleTimeout: idleTimeout,
}
go func() {
if err := s.httpServer.Serve(listener); err != nil && err != http.ErrServerClosed {
slog.Error("failed to start HTTP server", "error", err)
}
}()
s.startBackgroundJobs()
return nil
}
// startBackgroundJobs runs one-off data jobs that must not delay serving.
func (s *Server) startBackgroundJobs() {
ctx, cancel := context.WithCancel(context.Background())
s.backgroundCancel = cancel
runner := memopayload.NewRunner(s.Store, s.apiService.MarkdownService)
s.backgroundJobs.Go(func() {
if err := runner.Backfill(ctx); err != nil && ctx.Err() == nil {
slog.Error("failed to rebuild memo payloads; retrying on next start", slog.String("error", err.Error()))
}
})
}
func (s *Server) stopBackgroundJobs() {
if s.backgroundCancel != nil {
s.backgroundCancel()
}
s.backgroundJobs.Wait()
}
func (s *Server) Shutdown(ctx context.Context) {
ctx, cancel := context.WithTimeout(ctx, shutdownTimeout)
defer cancel()
slog.Info("server shutting down")
s.closeLongLivedConnections()
s.shutdownHTTPServer(ctx)
s.apiService.CloseUploads()
s.stopBackgroundJobs()
// Close database connection.
if err := s.Store.Close(); err != nil {
slog.Error("failed to close database", slog.String("error", err.Error()))
}
slog.Info("memos stopped properly")
}
func (s *Server) closeLongLivedConnections() {
// Long-lived SSE requests do not finish on their own during http.Server.Shutdown.
s.apiService.SSEHub.Close()
}
func (s *Server) shutdownHTTPServer(ctx context.Context) {
if s.httpServer == nil {
return
}
if err := s.httpServer.Shutdown(ctx); err != nil {
slog.Error("failed to shutdown server", slog.String("error", err.Error()))
if closeErr := s.httpServer.Close(); closeErr != nil && closeErr != http.ErrServerClosed {
slog.Error("failed to close server", slog.String("error", closeErr.Error()))
}
}
}
func (s *Server) getOrUpsertInstanceBasicSetting(ctx context.Context) (*storepb.InstanceBasicSetting, error) {
instanceBasicSetting, err := s.Store.GetInstanceBasicSetting(ctx)
if err != nil {
return nil, errors.Wrap(err, "failed to get instance basic setting")
}
modified := false
if instanceBasicSetting.SecretKey == "" {
instanceBasicSetting.SecretKey = uuid.NewV4().String()
modified = true
}
if modified {
instanceSetting, err := s.Store.UpsertInstanceSetting(ctx, &storepb.InstanceSetting{
Key: storepb.InstanceSettingKey_BASIC,
Value: &storepb.InstanceSetting_BasicSetting{BasicSetting: instanceBasicSetting},
})
if err != nil {
return nil, errors.Wrap(err, "failed to upsert instance setting")
}
instanceBasicSetting = instanceSetting.GetBasicSetting()
}
return instanceBasicSetting, nil
}