Repository: hashicorp/nomad
File: client/allocrunner/consul_hook.go (Go)

Code before:
				"requested_by": fmt.Sprintf("nomad_task_%s", task.Name),
			},
		}

		token, err = h.getConsulToken(consulConfig.Name, req)
		if err != nil {
			return fmt.Errorf("failed to derive Consul token for task %s: %v", task.Name, err)
		}

		// Store token in client state so it can be reused in case of desconnection
		if swi.ACLAccessTokensB64 == nil {
			// Avoid panics for tokens created before this field was added.
			swi.ACLAccessTokensB64 = map[string]string{}
		}

		// If this steps fail, we should log and continue allowing the task to start
		// if by any chance it restarts, it wont be able to reuse the stored token
		// and it will create a new one.
		t, err := writeACLToken(token)
		if err != nil {
			h.logger.Warn("error processing access token for", "task", task.Name, "error", err)
		}

		swi.ACLAccessTokensB64[tokenName] = t

		err = h.widmgr.Set(swi)
		if err != nil {
			h.logger.Warn("error updating access token for", "task", task.Name, "error", err)
		}
	}

	// Store token in results.
	if _, ok = tokens[clusterName]; !ok {
		tokens[clusterName] = make(map[string]*consulapi.ACLToken)
	}

	tokens[clusterName][tokenName] = token

	return nil
}

func (h *consulHook) prepareConsulTokensForServices(services []*structs.Service, tg *structs.TaskGroup, tokens map[string]map[string]*consulapi.ACLToken, env *taskenv.TaskEnv) error {
	if len(services) == 0 {
		return nil
	}

	var mErr *multierror.Error
	for _, service := range services {
		// Exit early if service doesn't need a Consul token.
		if service == nil || !service.IsConsul() || service.Identity == nil {
			continue

Review comment (performance):
Even if we change this to use its own state and not the WI's state, I think we end up potentially making a bunch of database writes here because we can have multiple tokens (ex. one for each task with a `template` block and one for each service). Can we batch these up?

Code after:
func (h *consulHook) prepareConsulTokensForServices(services []*structs.Service, tg *structs.TaskGroup, tokens map[string]map[string]*consulapi.ACLToken, env *taskenv.TaskEnv) error {
	var mErr *multierror.Error
	for _, service := range services {
		// Exit early if service doesn't need a Consul token.
		if service == nil || !service.IsConsul() || service.Identity == nil {
			continue
		}

		clusterName := service.GetConsulClusterName(tg)
		consulConfig, ok := h.consulConfigs[clusterName]
		if !ok {
			return fmt.Errorf("no such consul cluster: %s", clusterName)
		}

		// Find signed identity workload.
		ti := *service.IdentityHandle(env.ReplaceEnv)
		tokenName := service.Identity.Name
		token := tokens[clusterName][tokenName]

		// If no token was previously stored, create one.
		if token == nil {
			swi, err := h.widmgr.Get(ti)
			if err != nil {
				mErr = multierror.Append(mErr, fmt.Errorf(
					"error getting signed identity for service %s: %v",
					service.Name, err,
				))
				continue
			}

			h.logger.Debug("logging into consul", "name", ti.IdentityName, "type", ti.WorkloadType)
			req := consul.JWTLoginRequest{
				JWT:            swi.JWT,
				AuthMethodName: consulConfig.ServiceIdentityAuthMethod,
				Meta: map[string]string{
					"requested_by": fmt.Sprintf("nomad_service_%s", ti.InterpolatedWorkloadIdentifier),
				},
			}

			token, err = h.getConsulToken(clusterName, req)
			if err != nil {
				mErr = multierror.Append(mErr, fmt.Errorf(
					"failed to derive Consul token for service %s: %v",
					service.Name, err,
				))
				continue
			}

		}

		// Store token in results.