github.com/juju/juju@v0.0.0-20240430160146-1752b71fcf00/state/distribution.go (about)

     1  // Copyright 2014 Canonical Ltd.
     2  // Licensed under the AGPLv3, see LICENCE file for details.
     3  
     4  package state
     5  
     6  import (
     7  	"fmt"
     8  
     9  	"github.com/juju/collections/set"
    10  	"github.com/juju/errors"
    11  
    12  	"github.com/juju/juju/core/instance"
    13  	"github.com/juju/juju/environs/context"
    14  )
    15  
    16  // distributeUnit takes a unit and set of clean, possibly empty, instances
    17  // and asks the InstanceDistributor policy (if any) which ones are suitable
    18  // for assigning the unit to. If there is no InstanceDistributor, or the
    19  // distribution group is empty, then all of the candidates will be returned.
    20  func distributeUnit(u *Unit, candidates []instance.Id, limitZones []string) ([]instance.Id, error) {
    21  	if len(candidates) == 0 {
    22  		return nil, nil
    23  	}
    24  	if u.st.policy == nil {
    25  		return candidates, nil
    26  	}
    27  
    28  	distributor, err := u.st.policy.InstanceDistributor()
    29  	if errors.IsNotImplemented(err) {
    30  		return candidates, nil
    31  	} else if err != nil {
    32  		return nil, err
    33  	}
    34  	if distributor == nil {
    35  		return nil, fmt.Errorf("policy returned nil instance distributor without an error")
    36  	}
    37  
    38  	distributionGroup, err := ApplicationInstances(u.st, u.doc.Application)
    39  	if err != nil {
    40  		return nil, err
    41  	}
    42  	if len(distributionGroup) == 0 {
    43  		return candidates, nil
    44  	}
    45  	return distributor.DistributeInstances(context.CallContext(u.st), candidates, distributionGroup, limitZones)
    46  }
    47  
    48  // ApplicationInstances returns the instance IDs of provisioned
    49  // machines that are assigned units of the specified application.
    50  func ApplicationInstances(st *State, application string) ([]instance.Id, error) {
    51  	units, err := allUnits(st, application)
    52  	if err != nil {
    53  		return nil, err
    54  	}
    55  	instanceIds := make([]instance.Id, 0, len(units))
    56  	for _, unit := range units {
    57  		machineId, err := unit.AssignedMachineId()
    58  		if errors.IsNotAssigned(err) {
    59  			continue
    60  		} else if err != nil {
    61  			return nil, err
    62  		}
    63  		machine, err := st.Machine(machineId)
    64  		if err != nil {
    65  			return nil, err
    66  		}
    67  		instanceId, err := machine.InstanceId()
    68  		if err == nil {
    69  			instanceIds = append(instanceIds, instanceId)
    70  		} else if errors.IsNotProvisioned(err) {
    71  			continue
    72  		} else {
    73  			return nil, err
    74  		}
    75  	}
    76  	return instanceIds, nil
    77  }
    78  
    79  // ApplicationMachines returns the machine IDs of machines which have
    80  // the specified application listed as a principal.
    81  func ApplicationMachines(st *State, application string) ([]string, error) {
    82  	machines, err := st.AllMachines()
    83  	if err != nil {
    84  		return nil, err
    85  	}
    86  	applicationName := unitAppName(application)
    87  	var machineIds []string
    88  	for _, machine := range machines {
    89  		principalSet := set.NewStrings()
    90  		for _, principal := range machine.Principals() {
    91  			principalSet.Add(unitAppName(principal))
    92  		}
    93  		if principalSet.Contains(applicationName) {
    94  			machineIds = append(machineIds, machine.Id())
    95  		}
    96  	}
    97  	return machineIds, nil
    98  }