upload-proxy: add reload

Author: iSchluff

cmd/upload-proxy/config.go

@@ -10,9 +10,39 @@ import (
 type Config struct {
 	ListenAddress string
 	Sinks         []SinkConfig
+	Auth          AuthConfig
 }
 
-func (c *Config) ReadFromFile(path string) error {
+type SinkConfig struct {
+	Address   string
+	QueueSize int `toml:"queue-size"`
+	AuthConfig
+}
+
+func (s *SinkConfig) ApplyDefaults() {
+	if s.QueueSize == 0 {
+		s.QueueSize = DefaultQueueSize
+	}
+	if s.Username != "" && (s.AuthType != AuthTypeBasic && s.AuthType != AuthTypeDigest) {
+		s.AuthType = AuthTypeBasic
+	}
+}
+
+type AuthConfig struct {
+	Username string
+	Password string
+	AuthType AuthType `toml:"auth-type"`
+}
+
+type AuthType string
+
+const (
+	AuthTypeNone   AuthType = "none"
+	AuthTypeBasic  AuthType = "basic"
+	AuthTypeDigest AuthType = "digest"
+)
+
+func (c *Config) Load(path string, authPath string) error {
 	data, err := os.ReadFile(path)
 	if err != nil {
 		return fmt.Errorf("read: %w", err)
@@ -21,5 +51,23 @@ func (c *Config) ReadFromFile(path string) error {
 	if err != nil {
 		return fmt.Errorf("unmarshal: %w", err)
 	}
+	// load separate auth config if specified
+	if authPath != "" {
+		data, err := os.ReadFile(authPath)
+		if err != nil {
+			return fmt.Errorf("read: %w", err)
+		}
+		err = toml.Unmarshal(data, &c.Auth)
+		if err != nil {
+			return fmt.Errorf("unmarshal: %w", err)
+		}
+	}
+	// apply auth config to all sinks if set
+	for i := range c.Sinks {
+		c.Sinks[i].ApplyDefaults()
+		if c.Auth.Username != "" {
+			c.Sinks[i].AuthConfig = c.Auth
+		}
+	}
 	return nil
 }

cmd/upload-proxy/main.go

@@ -5,21 +5,22 @@ import (
 	"flag"
 	"fmt"
 	"os"
+	"os/signal"
+	"syscall"
 
 	"github.com/rs/zerolog"
 	"github.com/rs/zerolog/log"
-	"github.com/voc/stream-api/util"
 )
 
 func main() {
 	conf := Config{}
-	config := flag.String("config", "config.toml", "Set path to proxy config")
+	configPath := flag.String("config", "config.toml", "Set path to proxy config")
+	authConfigPath := flag.String("auth-config", "", "Set path to separate auth config (optional)")
 	debug := flag.Bool("debug", false, "sets log level to debug")
 	flag.StringVar(&conf.ListenAddress, "addr", ":8080", "Set listen address")
 	flag.Parse()
 	ctx, cancel := context.WithCancel(context.Background())
 	defer cancel()
-	util.HandleSignal(ctx, cancel)
 
 	log.Logger = log.Output(zerolog.ConsoleWriter{Out: os.Stderr})
 
@@ -28,42 +29,46 @@ func main() {
 		zerolog.SetGlobalLevel(zerolog.DebugLevel)
 	}
 
-	if err := conf.ReadFromFile(*config); err != nil {
+	if err := conf.Load(*configPath, *authConfigPath); err != nil {
 		log.Fatal().Err(err).Msg("config parse failed")
 	}
 
 	// Run proxy
-	if err := run(ctx, &conf); err != nil {
+	if err := run(ctx, conf, *configPath, *authConfigPath); err != nil {
 		log.Fatal().Err(err).Msg("proxy run failed")
 	}
 }
 
-func run(parentCtx context.Context, conf *Config) error {
+func run(parentCtx context.Context, conf Config, configPath string, authConfigPath string) error {
 	ctx, cancel := context.WithCancel(parentCtx)
 	defer cancel()
-	var sinks []*Sink
-	if len(conf.Sinks) == 0 {
-		log.Warn().Msg("no sinks configured")
-	}
-	for _, sinkConfig := range conf.Sinks {
-		sink, err := NewSink(sinkConfig)
-		if err != nil {
-			return fmt.Errorf("sink init failed: %w", err)
-		}
-		sinks = append(sinks, sink)
-		log.Info().Str("sink", sink.url.Host).Str("basePath", sink.url.Path).Str("authType", string(sink.conf.AuthType)).Msg("added sink")
-	}
-	proxy, err := NewProxy(ctx, conf.ListenAddress, sinks)
+	proxy, err := NewProxy(ctx, conf)
 	if err != nil {
 		return fmt.Errorf("proxy init failed: %w", err)
 	}
 	log.Info().Msgf("listening on %s", conf.ListenAddress)
 
-	select {
-	case <-ctx.Done():
-	case err := <-proxy.Errors():
-		log.Error().Err(err).Msg("server failed")
-		cancel()
+	signalReload := make(chan os.Signal, 1)
+	signal.Notify(signalReload, syscall.SIGHUP)
+
+outer:
+	for {
+		select {
+		case <-ctx.Done():
+			break outer
+		case err := <-proxy.Errors():
+			log.Error().Err(err).Msg("server failed")
+			cancel()
+			break outer
+		case <-signalReload:
+			log.Info().Msg("reloading config")
+			var newConf Config
+			if err := newConf.Load(configPath, authConfigPath); err != nil {
+				log.Error().Err(err).Msg("config parse failed")
+				break
+			}
+			_ = proxy.UpdateConfig(ctx, newConf)
+		}
 	}
 
 	proxy.Wait()

cmd/upload-proxy/main_test.go

@@ -10,12 +10,12 @@ func TestSmoke(t *testing.T) {
 	t.Parallel()
 	ctx, cancel := context.WithTimeout(context.Background(), 1*time.Second)
 	defer cancel()
-	err := run(ctx, &Config{
+	err := run(ctx, Config{
 		ListenAddress: "127.0.0.1:0",
 		Sinks: []SinkConfig{{
 			Address: "http://1.2.3.4:5678/upload",
 		}},
-	})
+	}, "", "")
 	if err != nil {
 		t.Error(err)
 	}

cmd/upload-proxy/proxy.go

@@ -24,6 +24,7 @@ const MaxFileSize = 50 * 1024 * 1024 // 50 MB
 
 type Proxy struct {
 	sinks     []*Sink
+	sinkMutex sync.Mutex
 	errors    chan error
 	transport *http.Transport
 	ctx       context.Context
@@ -32,7 +33,7 @@ type Proxy struct {
 	metrics   *ProxyMetrics
 }
 
-func NewProxy(ctx context.Context, addr string, sinks []*Sink) (*Proxy, error) {
+func NewProxy(ctx context.Context, conf Config) (*Proxy, error) {
 	tr := &http.Transport{
 		DialContext: (&net.Dialer{
 			Timeout:   30 * time.Second,
@@ -44,27 +45,25 @@ func NewProxy(ctx context.Context, addr string, sinks []*Sink) (*Proxy, error) {
 		TLSHandshakeTimeout:   10 * time.Second,
 		ExpectContinueTimeout: 1 * time.Second,
 	}
+	reg := prometheus.NewRegistry()
+	reg.MustRegister(
+		collectors.NewProcessCollector(collectors.ProcessCollectorOpts{}),
+	)
 	p := &Proxy{
-		sinks:     sinks,
 		transport: tr,
 		errors:    make(chan error, 1),
 		ctx:       ctx,
+		metrics:   NewProxyMetrics(reg),
 	}
 
 	mux := http.NewServeMux()
-	srv := http.Server{Addr: addr, Handler: mux}
-
-	reg := prometheus.NewRegistry()
-	reg.MustRegister(
-		collectors.NewProcessCollector(collectors.ProcessCollectorOpts{}),
-	)
-	p.metrics = NewProxyMetrics(reg)
+	srv := http.Server{Addr: conf.ListenAddress, Handler: mux}
 
 	// set routes
 	mux.HandleFunc("/", p.HandleUpload)
 	mux.Handle("/metrics", promhttp.HandlerFor(reg, promhttp.HandlerOpts{}))
 
-	ln, err := net.Listen("tcp", addr)
+	ln, err := net.Listen("tcp", conf.ListenAddress)
 	if err != nil {
 		return nil, err
 	}
@@ -91,22 +90,100 @@ func NewProxy(ctx context.Context, addr string, sinks []*Sink) (*Proxy, error) {
 		}
 	}()
 
-	// run sink uploaders
-	for _, sink := range p.sinks {
-		log.Printf("setup sink %+v\n", sink)
+	p.done.Add(1)
+	go p.runTimeout(ctx)
 
-		// if the number of workers is >1 the server would have to deal with out of order playlists
-		sink.Start(ctx, p.transport, reg, 1)
+	// initial config update
+	if err := p.UpdateConfig(ctx, conf); err != nil {
+		return nil, err
 	}
 
 	return p, nil
 }
 
+func (p *Proxy) UpdateConfig(ctx context.Context, conf Config) error {
+	p.sinkMutex.Lock()
+	defer p.sinkMutex.Unlock()
+	var newConfigs []SinkConfig
+	var err error
+	// mark removed sinks for graceful deletion
+	for _, s := range p.sinks {
+		var found bool
+		for _, sinkConfig := range conf.Sinks {
+			if s.Address() == sinkConfig.Address {
+				found = true
+				break
+			}
+		}
+		if found {
+			continue
+		}
+		s.StartGracePeriod()
+	}
+	// update existing sinks
+	for _, sinkConfig := range conf.Sinks {
+		var updated bool
+		for _, s := range p.sinks {
+			if s.Address() != sinkConfig.Address {
+				continue
+			}
+			// update existing sink
+			if err2 := s.UpdateConfig(ctx, sinkConfig); err2 != nil {
+				log.Error().Err(err2).Str("sink", s.url.Host).Msg("sink update failed")
+				firstErr(err2, &err)
+			}
+			updated = true
+			break
+		}
+		if updated {
+			continue
+		}
+		newConfigs = append(newConfigs, sinkConfig)
+	}
+	// create new sinks
+	for _, sinkConfig := range newConfigs {
+		sink, err2 := NewSink(ctx, sinkConfig, p.transport, p.metrics.reg)
+		if err2 != nil {
+			firstErr(err2, &err)
+			log.Error().Err(err2).Str("sink", sinkConfig.Address).Msg("sink init failed")
+			continue
+		}
+		p.sinks = append(p.sinks, sink)
+		log.Info().Str("sink", sink.url.Host).Str("basePath", sink.url.Path).Str("authType", string(sink.conf.AuthType)).Msg("added sink")
+	}
+	return err
+}
+
+func (p *Proxy) runTimeout(ctx context.Context) {
+	defer p.done.Done()
+	ticker := time.NewTicker(time.Minute)
+	defer ticker.Stop()
+	for {
+		select {
+		case <-ctx.Done():
+			return
+		case <-ticker.C:
+			p.sinkMutex.Lock()
+			sinks := p.sinks[:0]
+			for _, s := range p.sinks {
+				if s.IsStale() {
+					s.Stop()
+					log.Info().Str("sink", s.url.Host).Msg("removed sink after timeout")
+					continue
+				}
+				sinks = append(sinks, s)
+			}
+			p.sinks = sinks
+			p.sinkMutex.Unlock()
+		}
+	}
+}
+
 // Wait for server to finish
 func (p *Proxy) Wait() {
 	p.done.Wait()
 	for _, sink := range p.sinks {
-		sink.wait()
+		sink.Stop()
 	}
 	p.metrics.deregister()
 }
@@ -235,3 +312,9 @@ func (m *ProxyMetrics) deregister() {
 	m.reg.Unregister(m.totalNumRead)
 	m.reg.Unregister(m.totalReadDelay)
 }
+
+func firstErr(err error, errPtr *error) {
+	if *errPtr == nil && err != nil {
+		*errPtr = err
+	}
+}

cmd/upload-proxy/proxy_test.go

@@ -23,17 +23,15 @@ func TestUpload(t *testing.T) {
 	ctx, cancel := context.WithTimeout(context.Background(), 2*time.Second)
 	defer cancel()
 
-	sink, err := NewSink(SinkConfig{
-		Address: srv.URL + "/upload",
+	proxy, err := NewProxy(ctx, Config{
+		ListenAddress: "127.0.0.1:0",
+		Sinks: []SinkConfig{{
+			Address: srv.URL + "/upload",
+		}},
 	})
 	if err != nil {
 		t.Fatal(err)
 	}
-
-	proxy, err := NewProxy(ctx, "127.0.0.1:0", []*Sink{sink})
-	if err != nil {
-		t.Fatal(err)
-	}
 	go func() {
 		select {
 		case <-ctx.Done():