diff --git a/d2cli/main.go b/d2cli/main.go index 2a7425141..9ae124b20 100644 --- a/d2cli/main.go +++ b/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 - if timeoutFlag != nil { - timeout = time.Duration(*timeoutFlag) * time.Second - } - if timeout > 0 { - var cancel func() - ctx, cancel = context.WithTimeout(ctx, timeout) - defer cancel() - } + ctx, cancel := log.WithTimeout(ctx, time.Minute*2) + defer cancel() _, written, err := compile(ctx, ms, plugin, renderOpts, fontFamily, *animateIntervalFlag, inputPath, outputPath, *bundleFlag, *forceAppendixFlag, pw.Page) if err != nil { diff --git a/d2plugin/exec.go b/d2plugin/exec.go index b357267af..bc9835846 100644 --- a/d2plugin/exec.go +++ b/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 - 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() - } + ctx, cancel := log.WithTimeout(ctx, time.Minute) + defer cancel() graphBytes, err := d2graph.SerializeGraph(g) if err != nil { diff --git a/e2etests/report/main.go b/e2etests/report/main.go index 905aa0800..e50fca540 100644 --- a/e2etests/report/main.go +++ b/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 - 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() - } + ctx, cancel := log.WithTimeout(ctx, 2*time.Minute) + defer cancel() // don't want to pass empty args to CommandContext args := []string{"test", testDir, testMatchString} diff --git a/lib/log/log.go b/lib/log/log.go index d70edf0b8..681da7231 100644 --- a/lib/log/log.go +++ b/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) +}