Add a 'logging' section to config.yaml supporting:
- output: stderr (default), stdout, or a file path
- level: debug, info, warn, error (default: info)
Implementation:
- New internal/logging package with Setup() for output destination
and Logger struct with level-aware Debug/Info/Warn/Error methods
- Config struct extended with Logging section (output + level fields)
- Hub refactored to accept *logging.Logger via constructor injection
- main.go initializes logging early after config load
The leveled logger suppresses messages below the configured threshold
while maintaining the stdlib log format. File output uses append mode
with 0644 permissions for safe log rotation.
🤖 Assisted by the code-assist SOP
116 lines
2.8 KiB
Go
116 lines
2.8 KiB
Go
package main
|
|
|
|
import (
|
|
"context"
|
|
"flag"
|
|
"fmt"
|
|
"log"
|
|
"net/http"
|
|
"os"
|
|
"os/signal"
|
|
"syscall"
|
|
"time"
|
|
|
|
"websocket-relay/internal/config"
|
|
"websocket-relay/internal/hub"
|
|
"websocket-relay/internal/logging"
|
|
|
|
"github.com/prometheus/client_golang/prometheus/promhttp"
|
|
)
|
|
|
|
func main() {
|
|
configFile := flag.String("config-file", "config.yaml", "Path to configuration file")
|
|
flag.Parse()
|
|
|
|
cfg, err := config.Load(*configFile)
|
|
if err != nil {
|
|
log.Fatal("Failed to load config:", err)
|
|
}
|
|
|
|
// Setup log output destination
|
|
logFile, err := logging.Setup(cfg.Logging.Output)
|
|
if err != nil {
|
|
log.Fatal("Failed to setup logging:", err)
|
|
}
|
|
if logFile != nil {
|
|
defer logFile.Close()
|
|
}
|
|
|
|
// Create leveled logger
|
|
logger := logging.NewLogger(cfg.Logging.Level, log.Writer())
|
|
|
|
h := hub.New(logger)
|
|
go h.Run()
|
|
|
|
// Start metrics server if enabled
|
|
var metricsServer *http.Server
|
|
if cfg.Metrics.Enabled {
|
|
metricsMux := http.NewServeMux()
|
|
metricsMux.Handle("/metrics", promhttp.Handler())
|
|
metricsAddr := fmt.Sprintf(":%d", cfg.Metrics.Port)
|
|
metricsServer = &http.Server{
|
|
Addr: metricsAddr,
|
|
Handler: metricsMux,
|
|
}
|
|
go func() {
|
|
logger.Infof("Metrics server starting on %s", metricsAddr)
|
|
if err := metricsServer.ListenAndServe(); err != nil && err != http.ErrServerClosed {
|
|
logger.Errorf("Metrics server error: %v", err)
|
|
}
|
|
}()
|
|
}
|
|
|
|
mux := http.NewServeMux()
|
|
mux.HandleFunc("/", h.HandleWebSocket)
|
|
|
|
addr := fmt.Sprintf(":%d", cfg.Server.Port)
|
|
server := &http.Server{
|
|
Addr: addr,
|
|
Handler: mux,
|
|
}
|
|
|
|
// Start the main server in a goroutine
|
|
go func() {
|
|
if cfg.Server.TLS.Enabled {
|
|
logger.Infof("WebSocket relay server starting on %s (TLS)", addr)
|
|
if err := server.ListenAndServeTLS(cfg.Server.TLS.CertFile, cfg.Server.TLS.KeyFile); err != nil && err != http.ErrServerClosed {
|
|
logger.Errorf("Server error: %v", err)
|
|
os.Exit(1)
|
|
}
|
|
} else {
|
|
logger.Infof("WebSocket relay server starting on %s (HTTP)", addr)
|
|
if err := server.ListenAndServe(); err != nil && err != http.ErrServerClosed {
|
|
logger.Errorf("Server error: %v", err)
|
|
os.Exit(1)
|
|
}
|
|
}
|
|
}()
|
|
|
|
// Wait for interrupt signal
|
|
quit := make(chan os.Signal, 1)
|
|
signal.Notify(quit, syscall.SIGINT, syscall.SIGTERM)
|
|
sig := <-quit
|
|
logger.Infof("Received signal %v, shutting down gracefully...", sig)
|
|
|
|
// Create a deadline for the shutdown
|
|
ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
|
|
defer cancel()
|
|
|
|
// Shut down the main HTTP server (stops accepting new connections)
|
|
if err := server.Shutdown(ctx); err != nil {
|
|
logger.Errorf("HTTP server shutdown error: %v", err)
|
|
}
|
|
|
|
// Shut down the metrics server
|
|
if metricsServer != nil {
|
|
if err := metricsServer.Shutdown(ctx); err != nil {
|
|
logger.Errorf("Metrics server shutdown error: %v", err)
|
|
}
|
|
}
|
|
|
|
// Stop the hub and close all WebSocket connections
|
|
h.Shutdown()
|
|
|
|
logger.Info("Server stopped")
|
|
}
|