manager.go

/*
   Copyright 2014 CoreOS, Inc.

   Licensed under the Apache License, Version 2.0 (the "License");
   you may not use this file except in compliance with the License.
   You may obtain a copy of the License at

       http://www.apache.org/licenses/LICENSE-2.0

   Unless required by applicable law or agreed to in writing, software
   distributed under the License is distributed on an "AS IS" BASIS,
   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
   See the License for the specific language governing permissions and
   limitations under the License.
*/

package systemd

import (
	"fmt"
	"io/ioutil"
	"os"
	"path"
	"sync"

	"github.com/coreos/fleet/Godeps/_workspace/src/github.com/coreos/go-systemd/dbus"

	"github.com/coreos/fleet/log"
	"github.com/coreos/fleet/pkg"
	"github.com/coreos/fleet/unit"
)

const (
	DefaultUnitsDirectory = "/run/fleet/units/"
)

type systemdUnitManager struct {
	systemd  *dbus.Conn
	unitsDir string

	hashes map[string]unit.Hash
	mutex  sync.RWMutex
}

func NewSystemdUnitManager(uDir string) (*systemdUnitManager, error) {
	systemd, err := dbus.New()
	if err != nil {
		return nil, err
	}

	if err := os.MkdirAll(uDir, os.FileMode(0755)); err != nil {
		return nil, err
	}

	hashes, err := hashUnitFiles(uDir)
	if err != nil {
		return nil, err
	}

	mgr := systemdUnitManager{
		systemd:  systemd,
		unitsDir: uDir,
		hashes:   hashes,
		mutex:    sync.RWMutex{},
	}
	return &mgr, nil
}

func hashUnitFiles(dir string) (map[string]unit.Hash, error) {
	uNames, err := lsUnitsDir(dir)
	if err != nil {
		return nil, err
	}

	hMap := make(map[string]unit.Hash)
	for _, uName := range uNames {
		h, err := hashUnitFile(path.Join(dir, uName))
		if err != nil {
			return nil, err
		}

		hMap[uName] = h
	}

	return hMap, nil
}

func hashUnitFile(loc string) (unit.Hash, error) {
	b, err := ioutil.ReadFile(loc)
	if err != nil {
		return unit.Hash{}, err
	}

	uf, err := unit.NewUnitFile(string(b))
	if err != nil {
		return unit.Hash{}, err
	}

	return uf.Hash(), nil
}

// Load writes the given Unit to disk, subscribing to relevant dbus
// events, caching the Unit's Hash, and, if necessary, instructing the systemd
// daemon to reload.
func (m *systemdUnitManager) Load(name string, u unit.UnitFile) error {
	m.mutex.Lock()
	defer m.mutex.Unlock()
	err := m.writeUnit(name, u.String())
	if err != nil {
		return err
	}
	m.hashes[name] = u.Hash()
	if m.unitRequiresDaemonReload(name) {
		return m.daemonReload()
	}
	return nil
}

// Unload removes the indicated unit from the filesystem, deletes its
// associated Hash from the cache, clears its unit status in systemd, and
// performs a systemd daemon-reload
func (m *systemdUnitManager) Unload(name string) {
	m.mutex.Lock()
	defer m.mutex.Unlock()
	delete(m.hashes, name)
	m.removeUnit(name)
}

// TriggerStart asynchronously starts the unit identified by the given name.
// This function does not block for the underlying unit to actually start.
func (m *systemdUnitManager) TriggerStart(name string) {
	jobID, err := m.systemd.StartUnit(name, "replace", nil)
	if err == nil {
		log.Infof("Triggered systemd unit %s start: job=%d", name, jobID)
	} else {
		log.Errorf("Failed to trigger systemd unit %s start: %v", name, err)
	}
}

// TriggerStop asynchronously starts the unit identified by the given name.
// This function does not block for the underlying unit to actually stop.
func (m *systemdUnitManager) TriggerStop(name string) {
	jobID, err := m.systemd.StopUnit(name, "replace", nil)
	if err == nil {
		log.Infof("Triggered systemd unit %s stop: job=%d", name, jobID)
	} else {
		log.Errorf("Failed to trigger systemd unit %s stop: %v", name, err)
	}
}

// GetUnitState generates a UnitState object representing the
// current state of a Unit
func (m *systemdUnitManager) GetUnitState(name string) (*unit.UnitState, error) {
	m.mutex.Lock()
	defer m.mutex.Unlock()
	us, err := m.getUnitState(name)
	if err != nil {
		return nil, err
	}
	if h, ok := m.hashes[name]; ok {
		us.UnitHash = h.String()
	}
	return us, nil
}

func (m *systemdUnitManager) getUnitState(name string) (*unit.UnitState, error) {
	info, err := m.systemd.GetUnitProperties(name)
	if err != nil {
		return nil, err
	}
	us := unit.UnitState{
		LoadState:   info["LoadState"].(string),
		ActiveState: info["ActiveState"].(string),
		SubState:    info["SubState"].(string),
	}
	return &us, nil
}

func (m *systemdUnitManager) readUnit(name string) (string, error) {
	path := m.getUnitFilePath(name)
	contents, err := ioutil.ReadFile(path)
	if err == nil {
		return string(contents), nil
	}
	return "", fmt.Errorf("no unit file at local path %s", path)
}

func (m *systemdUnitManager) unitRequiresDaemonReload(name string) bool {
	prop, err := m.systemd.GetUnitProperty(name, "NeedDaemonReload")
	if prop == nil || err != nil {
		return false
	}

	return prop.Value.Value().(bool)
}

func (m *systemdUnitManager) daemonReload() error {
	log.Infof("Instructing systemd to reload units")
	return m.systemd.Reload()
}

// Units enumerates all files recognized as valid systemd units in
// this manager's units directory.
func (m *systemdUnitManager) Units() ([]string, error) {
	return lsUnitsDir(m.unitsDir)

}

func (m *systemdUnitManager) GetUnitStates(filter pkg.Set) (map[string]*unit.UnitState, error) {
	// Unfortunately we need to lock for the entire operation to ensure we
	// have a consistent view of the hashes. Otherwise, Load/Unload
	// operations could mutate the hashes before we've retrieved the state
	// for every unit in the filter, since they won't necessarily all be
	// present in the initial ListUnits() call.
	m.mutex.Lock()
	defer m.mutex.Unlock()
	dbusStatuses, err := m.systemd.ListUnits()

	if err != nil {
		return nil, err
	}

	states := make(map[string]*unit.UnitState)
	for _, dus := range dbusStatuses {
		if !filter.Contains(dus.Name) {
			continue
		}

		us := &unit.UnitState{
			LoadState:   dus.LoadState,
			ActiveState: dus.ActiveState,
			SubState:    dus.SubState,
		}
		if h, ok := m.hashes[dus.Name]; ok {
			us.UnitHash = h.String()
		}
		states[dus.Name] = us
	}

	// grab data on subscribed units that didn't show up in ListUnits, most
	// likely due to being inactive
	for _, name := range filter.Values() {
		if _, ok := states[name]; ok {
			continue
		}

		us, err := m.getUnitState(name)
		if err != nil {
			return nil, err
		}
		if h, ok := m.hashes[name]; ok {
			us.UnitHash = h.String()
		}
		states[name] = us
	}

	return states, nil
}

func (m *systemdUnitManager) writeUnit(name string, contents string) error {
	bContents := []byte(contents)
	log.Infof("Writing systemd unit %s (%db)", name, len(bContents))

	ufPath := m.getUnitFilePath(name)
	err := ioutil.WriteFile(ufPath, bContents, os.FileMode(0644))
	if err != nil {
		return err
	}

	_, err = m.systemd.LinkUnitFiles([]string{ufPath}, true, true)
	return err
}

func (m *systemdUnitManager) removeUnit(name string) {
	log.Infof("Removing systemd unit %s", name)

	m.systemd.DisableUnitFiles([]string{name}, true)
	m.systemd.ResetFailedUnit(name)

	ufPath := m.getUnitFilePath(name)
	os.Remove(ufPath)
}

func (m *systemdUnitManager) getUnitFilePath(name string) string {
	return path.Join(m.unitsDir, name)
}

func lsUnitsDir(dir string) ([]string, error) {
	filterFunc := func(name string) bool {
		if !unit.RecognizedUnitType(name) {
			log.Warningf("Found unrecognized file in %s, ignoring", path.Join(dir, name))
			return true
		}

		return false
	}

	return pkg.ListDirectory(dir, filterFunc)
}