Inside Terraform: Expanding for_each and count

Disclaimer: I am working at HashiCorp (now IBM) as part of the Terraform Core team. The postings on this site are my own and don’t necessarily represent IBM’s positions, strategies or opinions.
Since I am involved in Terraform my opinions can sometimes be (unconsciously) biased. I hope you enjoy the post anyway.

This is part of my Inside Terraform series where I deep dive into different parts of Terraform and explain how they work under the hood. This is the fifth post in the series and the start of the more graph-related posts. We will take a look at how Terraform expands blocks with for_each and count into multiple nodes in the graph.

This post won’t require any prior knowledge, but reading through the post on addrs, references and evaluation to clarify some of the concepts and terminology I will be using here might be helpful.

What does expansion mean?

Expansion means moving from a single block to a set of nodes in the graph. This is necessary to support for_each and count, but it is also necessary to have a new block in a module since modules support for_each and count.

resource "example" "foo" {
  # Can be a set
  for_each = toset(["a", "b", "c"])

  name = each.value # each.key is the same
}

resource "example" "bar" {
  # Can be a map
  for_each = { a = "foo", b = "bar", c = "baz" }

  name = each.key
  content = each.value
}

resource "example" "baz" {
  count = 3

  name = "example-${count.index}"
}

for_each and count can also use references of values from other resources, but the values must be known at plan time. One of the fundamental principle of Terraform is that we want to be able to plan as exactly as possible what will happen when we apply.

When a reference to a value that is not known at plan time is used in a for_each or count, Terraform will normally throw an error. There is an ongoing experiment (shipped to alphas only) to support deferring changes to a later plan / apply cycle that is also used in Stacks, but I’ll cover that in a later post.

Expansion in the Graph

Let’s first take a look at when we expand in the lifecycle of a Terraform run.

flowchart TD
A[Command]-->B[Config Loading]
B-->D
subgraph D[**Graph Walk**]
D1[Config Nodes]-->|**Expansion**|D2[Instance Nodes]
end
D2-->E[Providers]
E-->E2[Cloud APIs]
D2-->F[Plan / State]

This is simplified, but explains what systems are involved. The DynamicExpand method is called on the config nodes to expand the config nodes to instance nodes. This happens every time, it does not matter if for_each or count is used or not. If none is set the instance key will be addrs.NoKey and there will be only one instance node created. If for_each or count is used, there will be multiple instance nodes created with different instance keys.

Registering an Expansion

There are two parts to expanding a resource. First we need to know to which instances a resource (or top-level block supporting for_each / count in general) should be expanded. This entails both the module instances and the instance keys for the resource itself. The dynamic expand step takes care of both. When we reach this point, we will already have registered the expansion for the module, so we just need to register the expansion for the resource itself.

The context provides an instance expander under moduleCtx.InstanceExpander(). The expander has two main functionalities (in this context): Registering the instance keys for a resource and returning the list of instance addresses for a resource.

In the context of a resource we run this:

// Source: https://github.com/hashicorp/terraform/blob/43716508b7b09ed0cf99658bb25d24b660d256c8/internal/terraform/node_resource_abstract.go#L447-L479
	switch {
	case n.Config != nil && n.Config.Count != nil:
		count, countDiags := evaluateCountExpression(n.Config.Count, ctx, allowUnknown)
		diags = diags.Append(countDiags)
		if countDiags.HasErrors() {
			return diags
		}

		if count >= 0 {
			expander.SetResourceCount(addr.Module, n.Addr.Resource, count)
		} else {
			// -1 represents "unknown"
			expander.SetResourceCountUnknown(addr.Module, n.Addr.Resource)
		}

	case n.Config != nil && n.Config.ForEach != nil:
		forEach, known, forEachDiags := evaluateForEachExpression(n.Config.ForEach, ctx, allowUnknown)
		diags = diags.Append(forEachDiags)
		if forEachDiags.HasErrors() {
			return diags
		}

		// This method takes care of all of the business logic of updating this
		// while ensuring that any existing instances are preserved, etc.
		if known {
			expander.SetResourceForEach(addr.Module, n.Addr.Resource, forEach)
		} else {
			expander.SetResourceForEachUnknown(addr.Module, n.Addr.Resource)
		}

	default:
		expander.SetResourceSingle(addr.Module, n.Addr.Resource)
	}

The evaluate* functions are essentially what we discussed in evaluation plus a few sanity checks on the type / marks. With this being done we can now easily expand the resource.

Expanding a Node

When it comes to expansion there are multiple ways and the “right” way depends on the constraints of the node type you want to expand. We will go from simplest to most complex here. All these code snippets will be part of your DynamicExpand implementation.

`forEachModuleInstance` - No `for_each` or `count` on the block itself`, but can be used in a module

One could also use this helper if your block supports for_each or count but in practice we haven’t used it in those cases. Consistency in a codebase this big is hard.

// Source: https://github.com/hashicorp/terraform/blob/43716508b7b09ed0cf99658bb25d24b660d256c8/internal/terraform/node_local.go#L66-L86
func (n *nodeExpandLocal) DynamicExpand(ctx EvalContext) (*Graph, tfdiags.Diagnostics) {
	var g Graph
	expander := ctx.InstanceExpander()
	forEachModuleInstance(expander, n.Module, false, func(module addrs.ModuleInstance) {
		o := &NodeLocal{
			Addr:   n.Addr.Absolute(module),
			Config: n.Config,
		}
		log.Printf("[TRACE] Expanding local: adding %s as %T", o.Addr.String(), o)
		g.Add(o)
	}, func(pem addrs.PartialExpandedModule) {
		o := &nodeLocalInPartialModule{
			Addr:   addrs.ObjectInPartialExpandedModule(pem, n.Addr),
			Config: n.Config,
		}
		log.Printf("[TRACE] Expanding local: adding placeholder for all %s as %T", o.Addr.String(), o)
		g.Add(o)
	})
	addRootNodeToGraph(&g)
	return &g, nil
}

The PartialExpandedModule part refers to a module with an unknown value in for_each or count, as mentioned in the beginning I won’t be covering that part in depth but safe it for another post.

Classic expansion - not too much complexity in the node, but `for_each` and `count` are supported

This is an example from the action block. First we get the expander and the module instances:

expander := ctx.InstanceExpander()
moduleInstances := expander.ExpandModule(n.Addr.Module, false)

Then there is some special handling for Partial modules I am skipping over. Lastly we loop through the module instances and register the instance keys for the resource:

// Source: https://github.com/hashicorp/terraform/blob/43716508b7b09ed0cf99658bb25d24b660d256c8/internal/terraform/node_action.go#L60-L97
	for _, module := range moduleInstances {
		absActAddr := n.Addr.Absolute(module)

		// Check if the actions language experiment is enabled for this module.
		moduleCtx := evalContextForModuleInstance(ctx, module)

		// recordActionData is responsible for informing the expander of what
		// repetition mode this resource has, which allows expander.ExpandAction
		// to work below.
		moreDiags := n.recordActionData(moduleCtx, absActAddr)

		diags = diags.Append(moreDiags)
		if moreDiags.HasErrors() {
			return nil, diags
		}

		_, knownInstKeys, haveUnknownKeys := expander.ActionInstanceKeys(absActAddr)
		if haveUnknownKeys {
			// this should never happen, n.recordActionData explicitly sets
			// allowUnknown to be false, so we should pick up diagnostics
			// during that call instance reaching this branch.
			panic("found unknown keys in action instance")
		}

		// Expand the action instances for this module.
		for _, knownInstKey := range knownInstKeys {
			node := NodeActionDeclarationInstance{
				Addr:             absActAddr.Instance(knownInstKey),
				Config:           &n.Config,
				Schema:           n.Schema,
				ResolvedProvider: n.ResolvedProvider,
				Dependencies:     n.Dependencies,
			}

			g.Add(&node)
		}
	}
	addRootNodeToGraph(&g)

Expansion through a subgraph - when you need the power of Transformers

I haven’t touched on how we build the graph yet, so this is a bit of foreshadowing here. We basically have a set of transformers that encapsulate certain behaviors nodes can subscribe to (like being able to reference them). When your nodes might have more complex relationships / lifecycles you might want to use a subgraph to do the expansion. Resources is one example where this is done:

// Source: https://github.com/hashicorp/terraform/blob/43716508b7b09ed0cf99658bb25d24b660d256c8/internal/terraform/node_resource_plan.go#L489-L546
func (n *nodeExpandPlannableResource) resourceInstanceSubgraph(ctx EvalContext, addr addrs.AbsResource, instanceAddrs []addrs.AbsResourceInstance, imports addrs.Map[addrs.AbsResourceInstance, cty.Value]) (*Graph, tfdiags.Diagnostics) {
	var diags tfdiags.Diagnostics

	if n.Config == nil && n.generateConfigPath != "" && imports.Len() == 0 {
		// We're generating configuration, but there's nothing to import, which
		// means the import block must have expanded to zero instances.
		// the instance expander will always return a single instance because
		// we have assumed there will eventually be a configuration for this
		// resource, so return here before we add that to the graph.
		return &Graph{}, diags
	}

	// Our graph transformers require access to the full state, so we'll
	// temporarily lock it while we work on this.
	state := ctx.State().Lock()
	defer ctx.State().Unlock()

	// Start creating the steps
	steps := []GraphTransformer{
		// Expand the count or for_each (if present)
		&ResourceCountTransformer{
			Concrete:      n.concreteResource(ctx, imports, addrs.MakeMap[addrs.PartialExpandedResource, addrs.Set[addrs.AbsResourceInstance]](), n.skipPlanChanges),
			Schema:        n.Schema,
			Addr:          n.ResourceAddr(),
			InstanceAddrs: instanceAddrs,
		},

		// Add the count/for_each orphans
		&OrphanResourceInstanceCountTransformer{
			Concrete:      n.concreteResourceOrphan,
			Addr:          addr,
			InstanceAddrs: instanceAddrs,
			State:         state,
		},

		// Attach the state
		&AttachStateTransformer{State: state},

		// Targeting
		&TargetsTransformer{Targets: n.Targets},

		// Connect references so ordering is correct
		&ReferenceTransformer{},

		// Make sure there is a single root
		&RootTransformer{},
	}

	// Build the graph
	b := &BasicGraphBuilder{
		Steps: steps,
		Name:  "nodeExpandPlannableResource",
	}
	graph, graphDiags := b.Build(addr.Module)
	diags = diags.Append(graphDiags)

	return graph, diags
}

I will make a follow-up post about the graph building, until then the best resource to learn more about it is the Architecture doc.

Please mind the gaps

My main goal with this series is to give you glimpse into the internals of Terraform and enable you to contribute to Terraform itself. I would do myself and everyone reading this a disservice if I didn’t cover the footgun I ran into when I first implemented a node with expansion.

This particular footgun has to do with references and their timing in relationship to expansion.

It turns out that if you try to reference a node that does not exist yet (e.g. an instance type node used in a reference in your config) the graph does not create an edge to this currently non-existing node. Also, if you try to reference the config-level node the edge does work because these nodes are there from the start. The downside is that as soon as the dynamic expansion is done, the config-level node is done and your referencing node is ready to continue work. And that in turn means that it will run before the instance nodes that were created and that you intended to target. So think about the expanding nodes not as containers for the instance nodes, but rather as nodes that will create the expanded nodes.

flowchart TD
subgraph A[Initial]
A1[ConfigResourceA]
A2[ConfigResourceB]
end

subgraph B[References]
B1[ConfigResourceA]
B2[ConfigResourceB]
B1-->B2
end

subgraph C[Expanded]
C1[ConfigResourceA]-->
C2[ConfigResourceB]
C11[InstanceResourceA.0]
C12[InstanceResourceA.1]
C21[InstanceResourceB.0]
C22[InstanceResourceB.1]
end

subgraph D[Expanded + References]
D1[ConfigResourceA]-->
D2[ConfigResourceB]
D11[InstanceResourceA.0]
D12[InstanceResourceA.1]
D21[InstanceResourceB.0]
D22[InstanceResourceB.1]
D11-->D21
D12-->D22
end

A-->|References|B
B-->|DynamicExpansion|C
C-->|References|D

What this means in practice is that your config-level nodes should reference config-level nodes and your instance-level nodes should reference instance-level nodes. If you have a reference from an instance-level node to a config-level node, you will run into the footgun I described above.

In the next post in this series we will look at how planning in general works and how plans are rendered. Stay tuned!