new tracing system

internal/testenv/snippets/shutdown.go

@@ -0,0 +1,40 @@
+package snippets
+
+import (
+	"context"
+	"errors"
+	"net/http"
+
+	"github.com/pomerium/pomerium/internal/log"
+	"github.com/pomerium/pomerium/internal/testenv"
+)
+
+func RunWithDelayedShutdown(ctx context.Context, serve func() error, stop func()) func() error {
+	env := testenv.EnvFromContext(ctx)
+
+	stopping := make(chan struct{})
+	serveExited := make(chan error, 1)
+	env.OnStateChanged(testenv.Stopping, func() {
+		close(stopping)
+	})
+	cancel := env.OnStateChanged(testenv.Stopped, func() {
+		stop()
+		if err := <-serveExited; err != nil && !errors.Is(err, http.ErrServerClosed) {
+			log.Ctx(ctx).Err(err).Msg("error stopping server")
+		}
+	})
+	go func() {
+		serveExited <- serve()
+		close(serveExited)
+	}()
+
+	return func() error {
+		select {
+		case <-stopping:
+			return nil
+		case err := <-serveExited:
+			cancel()
+			return err
+		}
+	}
+}

internal/testenv/snippets/wait.go

@@ -4,6 +4,7 @@ import (
 	"context"
 	"time"
 
+	"github.com/pomerium/pomerium/internal/telemetry/trace"
 	"github.com/pomerium/pomerium/internal/testenv"
 	"github.com/pomerium/pomerium/pkg/grpcutil"
 	"google.golang.org/grpc"
@@ -12,6 +13,11 @@ import (
 )
 
 func WaitStartupComplete(env testenv.Environment, timeout ...time.Duration) time.Duration {
+	if env.GetState() == testenv.NotRunning {
+		panic("test bug: WaitStartupComplete called before starting the test environment")
+	}
+	_, span := trace.Continue(env.Context(), "snippets.WaitStartupComplete")
+	defer span.End()
 	start := time.Now()
 	recorder := env.NewLogRecorder()
 	if len(timeout) == 0 {
@@ -31,5 +37,6 @@ func WaitStartupComplete(env testenv.Environment, timeout ...time.Duration) time
 	)
 	env.Require().NoError(err)
 	env.Require().True(cc.WaitForStateChange(ctx, connectivity.Ready))
+	cc.Close()
 	return time.Since(start)
 }