From 559e4fc2a7f01f32fb0b69d15a9ac8829ab85173 Mon Sep 17 00:00:00 2001
From: Anton Tayanovskyy <anton@pulumi.com>
Date: Thu, 25 Jan 2024 16:12:14 -0500
Subject: [PATCH] Refactor SetTimeout into TimeoutOptions

---
 pkg/tfbridge/provider.go | 44 ++++++++++++++++++++++------------------
 pkg/tfshim/shim.go       | 24 ++++++++++++++--------
 2 files changed, 39 insertions(+), 29 deletions(-)

diff --git a/pkg/tfbridge/provider.go b/pkg/tfbridge/provider.go
index a38f86431..449dd6e44 100644
--- a/pkg/tfbridge/provider.go
+++ b/pkg/tfbridge/provider.go
@@ -19,6 +19,7 @@ import (
 	"fmt"
 	"log"
 	"strings"
+	"time"
 	"unicode"
 
 	"google.golang.org/grpc/status"
@@ -1002,23 +1003,20 @@ func (p *Provider) Create(ctx context.Context, req *pulumirpc.CreateRequest) (*p
 		return nil, errors.Wrapf(err, "preparing %s's new property state", urn)
 	}
 
-	diff, err := p.tf.Diff(ctx, res.TFName, nil, config, shim.DiffOptions{})
-	if err != nil {
-		return nil, errors.Wrapf(err, "diffing %s", urn)
-	}
-
 	// To populate default timeouts, we take the timeouts from the resource schema and insert them into the diff
 	timeouts, err := res.TF.DecodeTimeouts(config)
 	if err != nil {
 		return nil, errors.Errorf("error decoding timeout: %s", err)
 	}
-	if err = diff.EncodeTimeouts(timeouts); err != nil {
-		return nil, errors.Errorf("error setting default timeouts to diff: %s", err)
-	}
 
-	// If a custom timeout has been set for this method, overwrite the default timeout
-	if req.Timeout != 0 {
-		diff.SetTimeout(req.Timeout, shim.TimeoutCreate)
+	diff, err := p.tf.Diff(ctx, res.TFName, nil, config, shim.DiffOptions{
+		TimeoutOptions: shim.TimeoutOptions{
+			ResourceTimeout:  timeouts,
+			TimeoutOverrides: newTimeoutOverrides(shim.TimeoutCreate, req.Timeout),
+		},
+	})
+	if err != nil {
+		return nil, errors.Wrapf(err, "diffing %s", urn)
 	}
 
 	var newstate shim.InstanceState
@@ -1230,6 +1228,9 @@ func (p *Provider) Update(ctx context.Context, req *pulumirpc.UpdateRequest) (*p
 	ic := newIgnoreChanges(ctx, schema, fields, olds, news, req.GetIgnoreChanges())
 	diff, err := p.tf.Diff(ctx, res.TFName, state, config, shim.DiffOptions{
 		IgnoreChanges: ic,
+		TimeoutOptions: shim.TimeoutOptions{
+			TimeoutOverrides: newTimeoutOverrides(shim.TimeoutUpdate, req.Timeout),
+		},
 	})
 	if err != nil {
 		return nil, errors.Wrapf(err, "diffing %s", urn)
@@ -1248,10 +1249,6 @@ func (p *Provider) Update(ctx context.Context, req *pulumirpc.UpdateRequest) (*p
 			urn, diff.Destroy(), diff.RequiresNew())
 	}
 
-	if req.Timeout != 0 {
-		diff.SetTimeout(req.Timeout, shim.TimeoutUpdate)
-	}
-
 	var newstate shim.InstanceState
 	var reasons []string
 	if !req.GetPreview() {
@@ -1332,11 +1329,9 @@ func (p *Provider) Delete(ctx context.Context, req *pulumirpc.DeleteRequest) (*p
 	}
 
 	// Create a new destroy diff.
-	diff := p.tf.NewDestroyDiff(ctx, string(t))
-	if req.Timeout != 0 {
-		diff.SetTimeout(req.Timeout, shim.TimeoutDelete)
-	}
-
+	diff := p.tf.NewDestroyDiff(ctx, string(t), shim.TimeoutOptions{
+		TimeoutOverrides: newTimeoutOverrides(shim.TimeoutDelete, req.Timeout),
+	})
 	if _, err := p.tf.Apply(ctx, res.TFName, state, diff); err != nil {
 		return nil, errors.Wrapf(err, "deleting %s", urn)
 	}
@@ -1665,3 +1660,12 @@ func transformFromState(
 	}
 	return o, nil
 }
+
+// If a custom timeout has been set for this method, overwrite the default timeout.
+func newTimeoutOverrides(key shim.TimeoutKey, maybeTimeoutSeconds float64) map[shim.TimeoutKey]time.Duration {
+	timeoutOverrides := map[shim.TimeoutKey]time.Duration{}
+	if maybeTimeoutSeconds != 0 {
+		timeoutOverrides[shim.TimeoutCreate] = time.Duration(maybeTimeoutSeconds * float64(time.Second))
+	}
+	return timeoutOverrides
+}
diff --git a/pkg/tfshim/shim.go b/pkg/tfshim/shim.go
index c5c715f0d..172d9802d 100644
--- a/pkg/tfshim/shim.go
+++ b/pkg/tfshim/shim.go
@@ -41,8 +41,6 @@ type InstanceDiff interface {
 	ProposedState(res Resource, priorState InstanceState) (InstanceState, error)
 	Destroy() bool
 	RequiresNew() bool
-	EncodeTimeouts(timeouts *ResourceTimeout) error
-	SetTimeout(timeout float64, timeoutKey string)
 }
 
 type ValueType int
@@ -148,12 +146,14 @@ type SchemaMap interface {
 
 type ImportFunc func(t, id string, meta interface{}) ([]InstanceState, error)
 
+type TimeoutKey string
+
 const (
-	TimeoutCreate  = "create"
-	TimeoutRead    = "read"
-	TimeoutUpdate  = "update"
-	TimeoutDelete  = "delete"
-	TimeoutDefault = "default"
+	TimeoutCreate  TimeoutKey = "create"
+	TimeoutRead    TimeoutKey = "read"
+	TimeoutUpdate  TimeoutKey = "update"
+	TimeoutDelete  TimeoutKey = "delete"
+	TimeoutDefault TimeoutKey = "default"
 )
 
 type ResourceTimeout struct {
@@ -214,7 +214,7 @@ type Provider interface {
 	InitLogging(ctx context.Context)
 
 	// Create a Destroy diff for a resource identified by the TF token t.
-	NewDestroyDiff(ctx context.Context, t string) InstanceDiff
+	NewDestroyDiff(ctx context.Context, t string, opts TimeoutOptions) InstanceDiff
 
 	NewResourceConfig(ctx context.Context, object map[string]interface{}) ResourceConfig
 
@@ -222,8 +222,14 @@ type Provider interface {
 	IsSet(ctx context.Context, v interface{}) ([]interface{}, bool)
 }
 
+type TimeoutOptions struct {
+	ResourceTimeout  *ResourceTimeout // optional
+	TimeoutOverrides map[TimeoutKey]time.Duration
+}
+
 type DiffOptions struct {
-	IgnoreChanges IgnoreChanges
+	IgnoreChanges  IgnoreChanges
+	TimeoutOptions TimeoutOptions
 }
 
 // Supports the ignoreChanges Pulumi option.