ci daily fix

d2cli/watch.go

@@ -15,6 +15,7 @@ import (
 	"sort"
 	"strings"
 	"sync"
+	"sync/atomic"
 	"time"
 
 	"github.com/fsnotify/fsnotify"
@@ -219,6 +220,60 @@ func (w *watcher) goFunc(fn func(context.Context) error) {
  *
  * TODO: Abstract out file system and fsnotify to test this with 100% coverage. See comment in main_test.go
  */
+
+type safeServer struct {
+	server  *http.Server
+	running atomic.Bool
+	mu      sync.Mutex
+}
+
+func newSafeServer(handler http.Handler) *safeServer {
+	return &safeServer{
+		server: &http.Server{Handler: handler},
+	}
+}
+
+func (s *safeServer) ListenAndServe(l net.Listener) error {
+	s.mu.Lock()
+	defer s.mu.Unlock()
+
+	if !s.running.CompareAndSwap(false, true) {
+		return errors.New("server is already running")
+	}
+	defer s.running.Store(false)
+
+	return s.server.Serve(l)
+}
+
+func (s *safeServer) Shutdown(ctx context.Context) error {
+	s.mu.Lock()
+	defer s.mu.Unlock()
+
+	if !s.running.Load() {
+		return nil
+	}
+
+	return s.server.Shutdown(ctx)
+}
+
+func SafeServe(ctx context.Context, timeout time.Duration, handler http.Handler, l net.Listener) error {
+	s := newSafeServer(handler)
+
+	errCh := make(chan error, 1)
+	go func() {
+		errCh <- s.ListenAndServe(l)
+	}()
+
+	select {
+	case err := <-errCh:
+		return err
+	case <-ctx.Done():
+		shutdownCtx, cancel := context.WithTimeout(context.Background(), timeout)
+		defer cancel()
+		return s.Shutdown(shutdownCtx)
+	}
+}
+
 func (w *watcher) watchLoop(ctx context.Context) error {
 	lastModified := make(map[string]time.Time)
 
@@ -482,9 +537,10 @@ func (w *watcher) goServe() error {
 	m.Handle("/static/", http.StripPrefix("/static", w.staticFileServer))
 	m.Handle("/watch", xhttp.HandlerFuncAdapter{Log: w.ms.Log, Func: w.handleWatch})
 
-	s := xhttp.NewServer(w.ms.Log.Warn, xhttp.Log(w.ms.Log, m))
+	handler := xhttp.Log(w.ms.Log, m)
+
 	w.goFunc(func(ctx context.Context) error {
-		return xhttp.Serve(ctx, time.Second*30, s, w.l)
+		return SafeServe(ctx, time.Second*30, handler, w.l)
 	})
 
 	return nil