diff --git a/d2cli/main.go b/d2cli/main.go index 74a24c17e..0bfa4edf4 100644 --- a/d2cli/main.go +++ b/d2cli/main.go @@ -31,6 +31,7 @@ import ( "oss.terrastruct.com/d2/d2themes" "oss.terrastruct.com/d2/d2themes/d2themescatalog" "oss.terrastruct.com/d2/lib/background" + "oss.terrastruct.com/d2/lib/env" "oss.terrastruct.com/d2/lib/imgbundler" ctxlog "oss.terrastruct.com/d2/lib/log" "oss.terrastruct.com/d2/lib/pdf" @@ -298,8 +299,16 @@ func Run(ctx context.Context, ms *xmain.State) (err error) { return w.run() } - ctx, cancel := context.WithTimeout(ctx, time.Minute*2) - defer cancel() + timeout := time.Minute * 2 + 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() + } _, 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 4a5836f35..b357267af 100644 --- a/d2plugin/exec.go +++ b/d2plugin/exec.go @@ -15,6 +15,7 @@ import ( "oss.terrastruct.com/util-go/xmain" "oss.terrastruct.com/d2/d2graph" + "oss.terrastruct.com/d2/lib/env" ) // execPlugin uses the binary at pathname with the plugin protocol to implement @@ -147,8 +148,15 @@ func (p *execPlugin) Info(ctx context.Context) (_ *PluginInfo, err error) { } func (p *execPlugin) Layout(ctx context.Context, g *d2graph.Graph) error { - ctx, cancel := context.WithTimeout(ctx, time.Minute) - defer cancel() + 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() + } graphBytes, err := d2graph.SerializeGraph(g) if err != nil { diff --git a/e2etests/report/main.go b/e2etests/report/main.go index a2ef1896a..905aa0800 100644 --- a/e2etests/report/main.go +++ b/e2etests/report/main.go @@ -15,6 +15,7 @@ import ( "text/template" "time" + "oss.terrastruct.com/d2/lib/env" "oss.terrastruct.com/d2/lib/log" ) @@ -69,8 +70,17 @@ func main() { if !*skipTests { ctx := log.Stderr(context.Background()) - ctx, cancel := context.WithTimeout(ctx, 2*time.Minute) - defer cancel() + + 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() + } + // don't want to pass empty args to CommandContext args := []string{"test", testDir, testMatchString} if cpuProfileStr != "" { diff --git a/lib/env/env.go b/lib/env/env.go index 07ba22eef..c37eea66d 100644 --- a/lib/env/env.go +++ b/lib/env/env.go @@ -2,6 +2,7 @@ package env import ( "os" + "strconv" ) func Test() bool { @@ -25,3 +26,13 @@ func DevOnly() bool { func SkipGraphDiffTests() bool { return os.Getenv("SKIP_GRAPH_DIFF_TESTS") != "" } + +func Timeout() (int, bool) { + if s := os.Getenv("D2_TIMEOUT"); s != "" { + i, err := strconv.ParseInt(s, 10, 64) + if err == nil { + return int(i), true + } + } + return -1, false +}