github.com/axw/juju@v0.0.0-20161005053422-4bd6544d08d4/state/reboot.go (about)

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