fix: Resolve race condition in HTTP server handling

diff/diff.go

@@ -169,31 +169,32 @@ func formatRunes(s string) string {
 //
 // Here's an example that you can play with to better understand the behaviour:
 //
-//     err = diff.TestdataJSON(filepath.Join("testdata", t.Name()), "change me")
-//     if err != nil {
-//     	t.Fatal(err)
-//     }
+//	err = diff.TestdataJSON(filepath.Join("testdata", t.Name()), "change me")
+//	if err != nil {
+//		t.Fatal(err)
+//	}
 //
 // Normally you want to use t.Name() as path for clarity but you can pass in any string.
 // e.g. a single test could persist two json objects into testdata with:
 //
-//     err = diff.TestdataJSON(filepath.Join("testdata", t.Name(), "1"), "change me 1")
-//     if err != nil {
-//     	t.Fatal(err)
-//     }
-//     err = diff.TestdataJSON(filepath.Join("testdata", t.Name(), "2"), "change me 2")
-//     if err != nil {
-//     	t.Fatal(err)
-//     }
+//	err = diff.TestdataJSON(filepath.Join("testdata", t.Name(), "1"), "change me 1")
+//	if err != nil {
+//		t.Fatal(err)
+//	}
+//	err = diff.TestdataJSON(filepath.Join("testdata", t.Name(), "2"), "change me 2")
+//	if err != nil {
+//		t.Fatal(err)
+//	}
 //
 // These would persist in testdata/${t.Name()}/1.exp.json and testdata/${t.Name()}/2.exp.json
 //
 // It uses Files under the hood.
 //
 // note: testdata is the canonical Go directory for such persistent test only files.
-//       It is unfortunately poorly documented. See https://pkg.go.dev/cmd/go/internal/test
-//       So normally you'd want path to be filepath.Join("testdata", t.Name()).
-//       This is also the reason this function is named "TestdataJSON".
+//
+//	It is unfortunately poorly documented. See https://pkg.go.dev/cmd/go/internal/test
+//	So normally you'd want path to be filepath.Join("testdata", t.Name()).
+//	This is also the reason this function is named "TestdataJSON".
 func TestdataJSON(path string, got interface{}) error {
 	gotb := xjson.Marshal(got)
 	gotb = append(gotb, '\n')

xhttp/serve.go

@@ -3,9 +3,12 @@ package xhttp
 
 import (
 	"context"
+	"errors"
 	"log"
 	"net"
 	"net/http"
+	"sync"
+	"sync/atomic"
 	"time"
 
 	"oss.terrastruct.com/util-go/xcontext"
@@ -22,15 +25,52 @@ func NewServer(log *log.Logger, h http.Handler) *http.Server {
 	}
 }
 
+type safeServer struct {
+	*http.Server
+	running int32
+	mu      sync.Mutex
+}
+
+func newSafeServer(s *http.Server) *safeServer {
+	return &safeServer{
+		Server: s,
+	}
+}
+
+func (s *safeServer) ListenAndServe(l net.Listener) error {
+	s.mu.Lock()
+	defer s.mu.Unlock()
+
+	if !atomic.CompareAndSwapInt32(&s.running, 0, 1) {
+		return errors.New("server is already running")
+	}
+	defer atomic.StoreInt32(&s.running, 0)
+
+	return s.Serve(l)
+}
+
+func (s *safeServer) Shutdown(ctx context.Context) error {
+	s.mu.Lock()
+	defer s.mu.Unlock()
+
+	if atomic.LoadInt32(&s.running) == 0 {
+		return nil
+	}
+
+	return s.Server.Shutdown(ctx)
+}
+
 func Serve(ctx context.Context, shutdownTimeout time.Duration, s *http.Server, l net.Listener) error {
 	s.BaseContext = func(net.Listener) context.Context {
 		return ctx
 	}
 
+	ss := newSafeServer(s)
+
 	serverClosed := make(chan struct{})
 	var serverError error
 	go func() {
-		serverError = s.Serve(l)
+		serverError = ss.ListenAndServe(l)
 		close(serverClosed)
 	}()
 
@@ -41,11 +81,12 @@ func Serve(ctx context.Context, shutdownTimeout time.Duration, s *http.Server, l
 		shutdownCtx, cancel := context.WithTimeout(xcontext.WithoutCancel(ctx), shutdownTimeout)
 		defer cancel()
 
-		err := s.Shutdown(shutdownCtx)
+		err := ss.Shutdown(shutdownCtx)
 		<-serverClosed // Wait for server to exit
 		if err != nil {
 			return err
 		}
 		return serverError
 	}
+
 }