util func for overriding default timeout according to env var

d2cli/main.go

@@ -32,6 +32,7 @@ import (
 	"oss.terrastruct.com/d2/d2themes/d2themescatalog"
 	"oss.terrastruct.com/d2/lib/background"
 	"oss.terrastruct.com/d2/lib/imgbundler"
+	"oss.terrastruct.com/d2/lib/log"
 	ctxlog "oss.terrastruct.com/d2/lib/log"
 	"oss.terrastruct.com/d2/lib/pdf"
 	"oss.terrastruct.com/d2/lib/png"
@@ -306,15 +307,8 @@ func Run(ctx context.Context, ms *xmain.State) (err error) {
 		return w.run()
 	}
 
-	timeout := time.Minute * 2
+	ctx, cancel := log.WithTimeout(ctx, time.Minute*2)
-	if timeoutFlag != nil {
-		timeout = time.Duration(*timeoutFlag) * time.Second
-	}
-	if timeout > 0 {
-		var cancel func()
-		ctx, cancel = context.WithTimeout(ctx, timeout)
 	defer cancel()
-	}
 
 	_, written, err := compile(ctx, ms, plugin, renderOpts, fontFamily, *animateIntervalFlag, inputPath, outputPath, *bundleFlag, *forceAppendixFlag, pw.Page)
 	if err != nil {

d2plugin/exec.go

@@ -15,7 +15,7 @@ import (
 	"oss.terrastruct.com/util-go/xmain"
 
 	"oss.terrastruct.com/d2/d2graph"
-	"oss.terrastruct.com/d2/lib/env"
+	"oss.terrastruct.com/d2/lib/log"
 )
 
 // execPlugin uses the binary at pathname with the plugin protocol to implement
@@ -148,15 +148,8 @@ func (p *execPlugin) Info(ctx context.Context) (_ *PluginInfo, err error) {
 }
 
 func (p *execPlugin) Layout(ctx context.Context, g *d2graph.Graph) error {
-	timeout := time.Minute
+	ctx, cancel := log.WithTimeout(ctx, time.Minute)
-	if seconds, has := env.Timeout(); has {
-		timeout = time.Duration(seconds) * time.Second
-	}
-	if timeout > 0 {
-		var cancel func()
-		ctx, cancel = context.WithTimeout(ctx, timeout)
 	defer cancel()
-	}
 
 	graphBytes, err := d2graph.SerializeGraph(g)
 	if err != nil {

e2etests/report/main.go

@@ -15,7 +15,6 @@ import (
 	"text/template"
 	"time"
 
-	"oss.terrastruct.com/d2/lib/env"
 	"oss.terrastruct.com/d2/lib/log"
 )
 
@@ -71,15 +70,8 @@ func main() {
 	if !*skipTests {
 		ctx := log.Stderr(context.Background())
 
-		timeout := 2 * time.Minute
+		ctx, cancel := log.WithTimeout(ctx, 2*time.Minute)
-		if seconds, has := env.Timeout(); has {
-			timeout = time.Duration(seconds) * time.Second
-		}
-		if timeout > 0 {
-			var cancel func()
-			ctx, cancel = context.WithTimeout(ctx, timeout)
 		defer cancel()
-		}
 
 		// don't want to pass empty args to CommandContext
 		args := []string{"test", testDir, testMatchString}

lib/log/log.go

@@ -8,6 +8,7 @@ import (
 	"os"
 	"runtime/debug"
 	"testing"
+	"time"
 
 	"cdr.dev/slog"
 	"cdr.dev/slog/sloggers/sloghuman"
@@ -112,3 +113,16 @@ func Stderr(ctx context.Context) context.Context {
 
 	return With(ctx, l)
 }
+
+// WithTimeout returns context.WithTimeout(ctx, timeout) but timeout is overridden with D2_TIMEOUT if set
+func WithTimeout(ctx context.Context, timeout time.Duration) (context.Context, context.CancelFunc) {
+	t := timeout
+	if seconds, has := env.Timeout(); has {
+		t = time.Duration(seconds) * time.Second
+	}
+	if t <= 0 {
+		return ctx, func() {}
+	}
+
+	return context.WithTimeout(ctx, t)
+}