github.com/makyo/juju@v0.0.0-20160425123129-2608902037e9/state/reboot.go (about)

     1  package state
     2  
     3  import (
     4  	"fmt"
     5  
     6  	"github.com/juju/errors"
     7  	"gopkg.in/mgo.v2"
     8  	"gopkg.in/mgo.v2/bson"
     9  	"gopkg.in/mgo.v2/txn"
    10  )
    11  
    12  var _ RebootFlagSetter = (*Machine)(nil)
    13  var _ RebootActionGetter = (*Machine)(nil)
    14  
    15  // RebootAction defines the action a machine should
    16  // take when a hook needs to reboot
    17  type RebootAction string
    18  
    19  const (
    20  	// ShouldDoNothing instructs a machine agent that no action
    21  	// is required on its part
    22  	ShouldDoNothing RebootAction = "noop"
    23  	// ShouldReboot instructs a machine to reboot
    24  	// this happens when a hook running on a machine, requests
    25  	// a reboot
    26  	ShouldReboot RebootAction = "reboot"
    27  	// ShouldShutdown instructs a machine to shut down. This usually
    28  	// happens when running inside a container, and a hook on the parent
    29  	// machine requests a reboot
    30  	ShouldShutdown RebootAction = "shutdown"
    31  )
    32  
    33  // rebootDoc will hold the reboot flag for a machine.
    34  type rebootDoc struct {
    35  	DocID     string `bson:"_id"`
    36  	Id        string `bson:"machineid"`
    37  	ModelUUID string `bson:"model-uuid"`
    38  }
    39  
    40  func (m *Machine) setFlag() error {
    41  	if m.Life() == Dead {
    42  		return mgo.ErrNotFound
    43  	}
    44  	ops := []txn.Op{
    45  		assertModelActiveOp(m.st.ModelUUID()),
    46  		{
    47  			C:      machinesC,
    48  			Id:     m.doc.DocID,
    49  			Assert: notDeadDoc,
    50  		}, {
    51  			C:      rebootC,
    52  			Id:     m.doc.DocID,
    53  			Insert: &rebootDoc{Id: m.Id()},
    54  		},
    55  	}
    56  	err := m.st.runTransaction(ops)
    57  	if err == txn.ErrAborted {
    58  		if err := checkModelActive(m.st); err != nil {
    59  			return errors.Trace(err)
    60  		}
    61  		return mgo.ErrNotFound
    62  	} else if err != nil {
    63  		return errors.Errorf("failed to set reboot flag: %v", err)
    64  	}
    65  	return nil
    66  }
    67  
    68  func removeRebootDocOp(st *State, machineId string) txn.Op {
    69  	op := txn.Op{
    70  		C:      rebootC,
    71  		Id:     st.docID(machineId),
    72  		Remove: true,
    73  	}
    74  	return op
    75  }
    76  
    77  func (m *Machine) clearFlag() error {
    78  	reboot, closer := m.st.getCollection(rebootC)
    79  	defer closer()
    80  
    81  	docID := m.doc.DocID
    82  	count, err := reboot.FindId(docID).Count()
    83  	if count == 0 {
    84  		return nil
    85  	}
    86  	ops := []txn.Op{removeRebootDocOp(m.st, m.Id())}
    87  	err = m.st.runTransaction(ops)
    88  	if err != nil {
    89  		return errors.Errorf("failed to clear reboot flag: %v", err)
    90  	}
    91  	return nil
    92  }
    93  
    94  // SetRebootFlag sets the reboot flag of a machine to a boolean value. It will also
    95  // do a lazy create of a reboot document if needed; i.e. If a document
    96  // does not exist yet for this machine, it will create it.
    97  func (m *Machine) SetRebootFlag(flag bool) error {
    98  	if flag {
    99  		return m.setFlag()
   100  	}
   101  	return m.clearFlag()
   102  }
   103  
   104  // GetRebootFlag returns the reboot flag for this machine.
   105  func (m *Machine) GetRebootFlag() (bool, error) {
   106  	rebootCol, closer := m.st.getCollection(rebootC)
   107  	defer closer()
   108  
   109  	count, err := rebootCol.FindId(m.doc.DocID).Count()
   110  	if err != nil {
   111  		return false, fmt.Errorf("failed to get reboot flag: %v", err)
   112  	}
   113  	if count == 0 {
   114  		return false, nil
   115  	}
   116  	return true, nil
   117  }
   118  
   119  func (m *Machine) machinesToCareAboutRebootsFor() []string {
   120  	var possibleIds []string
   121  	for currentId := m.Id(); currentId != ""; {
   122  		possibleIds = append(possibleIds, currentId)
   123  		currentId = ParentId(currentId)
   124  	}
   125  	return possibleIds
   126  }
   127  
   128  // ShouldRebootOrShutdown check if the current node should reboot or shutdown
   129  // If we are a container, and our parent needs to reboot, this should return:
   130  // ShouldShutdown
   131  func (m *Machine) ShouldRebootOrShutdown() (RebootAction, error) {
   132  	rebootCol, closer := m.st.getCollection(rebootC)
   133  	defer closer()
   134  
   135  	machines := m.machinesToCareAboutRebootsFor()
   136  
   137  	docs := []rebootDoc{}
   138  	sel := bson.D{{"machineid", bson.D{{"$in", machines}}}}
   139  	if err := rebootCol.Find(sel).All(&docs); err != nil {
   140  		return ShouldDoNothing, errors.Trace(err)
   141  	}
   142  
   143  	iNeedReboot := false
   144  	for _, val := range docs {
   145  		if val.Id != m.doc.Id {
   146  			return ShouldShutdown, nil
   147  		}
   148  		iNeedReboot = true
   149  	}
   150  	if iNeedReboot {
   151  		return ShouldReboot, nil
   152  	}
   153  	return ShouldDoNothing, nil
   154  }
   155  
   156  type RebootFlagSetter interface {
   157  	SetRebootFlag(flag bool) error
   158  }
   159  
   160  type RebootActionGetter interface {
   161  	ShouldRebootOrShutdown() (RebootAction, error)
   162  }