<commit_before>package openstack

import (
	"bytes"
	"encoding/base64"
	"encoding/json"
	"fmt"
	"strings"
	"text/template"
	"time"

	"github.com/docker/docker/pkg/namesgenerator"
	"github.com/gophercloud/gophercloud/openstack/compute/v2/servers"
	"github.com/spf13/viper"

	"github.com/ovh/cds/sdk"
	"github.com/ovh/cds/sdk/log"
)

// SpawnWorker creates a new cloud instances
// requirements are not supported
func (h *HatcheryCloud) SpawnWorker(model *sdk.Model, job *sdk.PipelineBuildJob, registerOnly bool, logInfo string) (string, error) {
	//generate a pretty cool name
	name := model.Name + "-" + strings.Replace(namesgenerator.GetRandomName(0), "_", "-", -1)
	if registerOnly {
		name = "register-" + name
	}

	if job != nil {
		log.Info("spawnWorker> spawning worker %s model:%s for job %d - %s", name, model.Name, job.ID, logInfo)
	} else {
		log.Info("spawnWorker> spawning worker %s model:%s - %s", name, model.Name, logInfo)
	}

	var omd sdk.OpenstackModelData

	if h.hatch == nil {
		return "", fmt.Errorf("hatchery disconnected from engine")
	}

	if len(h.getServers()) == viper.GetInt("max-worker") {
		log.Debug("MaxWorker limit (%d) reached", viper.GetInt("max-worker"))
		return "", nil
	}

	if err := json.Unmarshal([]byte(model.Image), &omd); err != nil {
		return "", err
	}

	// Get image ID
	imageID, erri := h.imageID(omd.Image)
	if erri != nil {
		return "", erri
	}

	// Get flavor ID
	flavorID, errf := h.flavorID(omd.Flavor)
	if errf != nil {
		return "", errf
	}

	// Decode base64 given user data
	udataModel, errd := base64.StdEncoding.DecodeString(omd.UserData)
	if errd != nil {
		return "", errd
	}

	graylog := ""
	if viper.GetString("graylog_host") != "" {
		graylog += fmt.Sprintf("export CDS_GRAYLOG_HOST=%s ", viper.GetString("graylog_host"))
	}
	if viper.GetString("graylog_port") != "" {
		graylog += fmt.Sprintf("export CDS_GRAYLOG_PORT=%s ", viper.GetString("graylog_port"))
	}
	if viper.GetString("graylog_extra_key") != "" {
		graylog += fmt.Sprintf("export CDS_GRAYLOG_EXTRA_KEY=%s ", viper.GetString("graylog_extra_key"))
	}
	if viper.GetString("graylog_extra_value") != "" {
		graylog += fmt.Sprintf("export CDS_GRAYLOG_EXTRA_VALUE=%s ", viper.GetString("graylog_extra_value"))
	}

	udataEnd := `
cd $HOME
# Download and start worker with curl
curl  "{{.API}}/download/worker/$(uname -m)" -o worker --retry 10 --retry-max-time 120 -C - >> /tmp/user_data 2>&1
chmod +x worker
export CDS_SINGLE_USE=1
export CDS_FORCE_EXIT=1
export CDS_API={{.API}}
export CDS_TOKEN={{.Key}}
export CDS_NAME={{.Name}}
export CDS_MODEL={{.Model}}
export CDS_HATCHERY={{.Hatchery}}
export CDS_HATCHERY_NAME={{.HatcheryName}}
export CDS_BOOKED_JOB_ID={{.JobID}}
export CDS_TTL={{.TTL}}
{{.Graylog}}
./worker`

	if registerOnly {
		udataEnd += " register"
	}
	udataEnd += " ; sudo shutdown -h now;"

	var jobID int64
	if job != nil {
		jobID = job.ID
	}

	var withExistingImage bool
	if !model.NeedRegistration && !registerOnly {
		start := time.Now()
		imgs := h.getImages()
		log.Debug("spawnWorker> call images.List on openstack took %fs, nbImages:%d", time.Since(start).Seconds(), len(imgs))
		for _, img := range imgs {
			workerModelName, _ := img.Metadata["worker_model_name"]
			if workerModelName == model.Name {
				withExistingImage = true
				var jobInfo string
				if jobID != 0 {
					jobInfo = fmt.Sprintf(" job:%d", jobID)
				}
				log.Info("spawnWorker> existing image found for worker:%s model:%s img:%s %s %s", name, model.Name, img.ID, jobInfo, logInfo)
				imageID = img.ID
				break
			}
		}
	}

	tmpl, errt := template.New("udata").Parse(string(udataEnd))
	if errt != nil {
		return "", errt
	}
	udataParam := struct {
		API          string
		Name         string
		Key          string
		Model        int64
		Hatchery     int64
		HatcheryName string
		JobID        int64
		TTL          int
		Graylog      string
	}{
		API:          viper.GetString("api"),
		Name:         name,
		Key:          viper.GetString("token"),
		Model:        model.ID,
		Hatchery:     h.hatch.ID,
		HatcheryName: h.hatch.Name,
		JobID:        jobID,
		TTL:          h.workerTTL,
		Graylog:      graylog,
	}
	var buffer bytes.Buffer
	if err := tmpl.Execute(&buffer, udataParam); err != nil {
		return "", err
	}

	var udataBegin, udata string

	if withExistingImage {
		log.Debug("spawnWorker> using userdata from existing image")
		udataBegin = `#!/bin/bash
set +e
export CDS_FROM_WORKER_IMAGE="true";
`
	} else {
		log.Debug("spawnWorker> using userdata from worker model")
		udataBegin = `#!/bin/bash
set +e
export CDS_FROM_WORKER_IMAGE="false";
`
	}
	udata = udataBegin + string(udataModel) + buffer.String()

	// Encode again
	udata64 := base64.StdEncoding.EncodeToString([]byte(udata))

	// Create openstack vm
	meta := map[string]string{
		"worker":                     name,
		"hatchery_name":              h.Hatchery().Name,
		"register_only":              fmt.Sprintf("%t", registerOnly),
		"flavor":                     omd.Flavor,
		"model":                      omd.Image,
		"worker_model_name":          model.Name,
		"worker_model_last_modified": fmt.Sprintf("%d", model.UserLastModified.Unix()),
	}

	// Ip len(ipsInfos.ips) > 0, specify one of those
	var ip string
	if len(ipsInfos.ips) > 0 {
		var errai error
		ip, errai = h.findAvailableIP(name)
		if errai != nil {
			return "", errai
		}
		log.Debug("Found %s as available IP", ip)
	}

	networks := []servers.Network{{UUID: h.networkID, FixedIP: ip}}
	r := servers.Create(h.client, servers.CreateOpts{
		Name:      name,
		FlavorRef: flavorID,
		ImageRef:  imageID,
		Metadata:  meta,
		UserData:  []byte(udata64),
		Networks:  networks,
	})

	server, err := r.Extract()
	if err != nil {
		return "", fmt.Errorf("SpawnWorker> Unable to create server: name:%s flavor:%s image:%s metadata:%v networks:%s err:%s body:%s", name, flavorID, imageID, meta, networks, err, r.Body)
	}
	log.Debug("SpawnWorker> Created Server ID: %s", server.ID)
	return name, nil
}
<commit_msg>fix (hatchery/openstack): force get worker (#937)<commit_after>package openstack

import (
	"bytes"
	"encoding/base64"
	"encoding/json"
	"fmt"
	"strings"
	"text/template"
	"time"

	"github.com/docker/docker/pkg/namesgenerator"
	"github.com/gophercloud/gophercloud/openstack/compute/v2/servers"
	"github.com/spf13/viper"

	"github.com/ovh/cds/sdk"
	"github.com/ovh/cds/sdk/log"
)

// SpawnWorker creates a new cloud instances
// requirements are not supported
func (h *HatcheryCloud) SpawnWorker(model *sdk.Model, job *sdk.PipelineBuildJob, registerOnly bool, logInfo string) (string, error) {
	//generate a pretty cool name
	name := model.Name + "-" + strings.Replace(namesgenerator.GetRandomName(0), "_", "-", -1)
	if registerOnly {
		name = "register-" + name
	}

	if job != nil {
		log.Info("spawnWorker> spawning worker %s model:%s for job %d - %s", name, model.Name, job.ID, logInfo)
	} else {
		log.Info("spawnWorker> spawning worker %s model:%s - %s", name, model.Name, logInfo)
	}

	var omd sdk.OpenstackModelData

	if h.hatch == nil {
		return "", fmt.Errorf("hatchery disconnected from engine")
	}

	if len(h.getServers()) == viper.GetInt("max-worker") {
		log.Debug("MaxWorker limit (%d) reached", viper.GetInt("max-worker"))
		return "", nil
	}

	if err := json.Unmarshal([]byte(model.Image), &omd); err != nil {
		return "", err
	}

	// Get image ID
	imageID, erri := h.imageID(omd.Image)
	if erri != nil {
		return "", erri
	}

	// Get flavor ID
	flavorID, errf := h.flavorID(omd.Flavor)
	if errf != nil {
		return "", errf
	}

	// Decode base64 given user data
	udataModel, errd := base64.StdEncoding.DecodeString(omd.UserData)
	if errd != nil {
		return "", errd
	}

	graylog := ""
	if viper.GetString("graylog_host") != "" {
		graylog += fmt.Sprintf("export CDS_GRAYLOG_HOST=%s ", viper.GetString("graylog_host"))
	}
	if viper.GetString("graylog_port") != "" {
		graylog += fmt.Sprintf("export CDS_GRAYLOG_PORT=%s ", viper.GetString("graylog_port"))
	}
	if viper.GetString("graylog_extra_key") != "" {
		graylog += fmt.Sprintf("export CDS_GRAYLOG_EXTRA_KEY=%s ", viper.GetString("graylog_extra_key"))
	}
	if viper.GetString("graylog_extra_value") != "" {
		graylog += fmt.Sprintf("export CDS_GRAYLOG_EXTRA_VALUE=%s ", viper.GetString("graylog_extra_value"))
	}

	udataEnd := `
cd $HOME
# Download and start worker with curl
rm -f worker
curl  "{{.API}}/download/worker/$(uname -m)" -o worker --retry 10 --retry-max-time 120 -C - >> /tmp/user_data 2>&1
chmod +x worker
export CDS_SINGLE_USE=1
export CDS_FORCE_EXIT=1
export CDS_API={{.API}}
export CDS_TOKEN={{.Key}}
export CDS_NAME={{.Name}}
export CDS_MODEL={{.Model}}
export CDS_HATCHERY={{.Hatchery}}
export CDS_HATCHERY_NAME={{.HatcheryName}}
export CDS_BOOKED_JOB_ID={{.JobID}}
export CDS_TTL={{.TTL}}
{{.Graylog}}
./worker`

	if registerOnly {
		udataEnd += " register"
	}
	udataEnd += " ; sudo shutdown -h now;"

	var jobID int64
	if job != nil {
		jobID = job.ID
	}

	var withExistingImage bool
	if !model.NeedRegistration && !registerOnly {
		start := time.Now()
		imgs := h.getImages()
		log.Debug("spawnWorker> call images.List on openstack took %fs, nbImages:%d", time.Since(start).Seconds(), len(imgs))
		for _, img := range imgs {
			workerModelName, _ := img.Metadata["worker_model_name"]
			if workerModelName == model.Name {
				withExistingImage = true
				var jobInfo string
				if jobID != 0 {
					jobInfo = fmt.Sprintf(" job:%d", jobID)
				}
				log.Info("spawnWorker> existing image found for worker:%s model:%s img:%s %s %s", name, model.Name, img.ID, jobInfo, logInfo)
				imageID = img.ID
				break
			}
		}
	}

	tmpl, errt := template.New("udata").Parse(string(udataEnd))
	if errt != nil {
		return "", errt
	}
	udataParam := struct {
		API          string
		Name         string
		Key          string
		Model        int64
		Hatchery     int64
		HatcheryName string
		JobID        int64
		TTL          int
		Graylog      string
	}{
		API:          viper.GetString("api"),
		Name:         name,
		Key:          viper.GetString("token"),
		Model:        model.ID,
		Hatchery:     h.hatch.ID,
		HatcheryName: h.hatch.Name,
		JobID:        jobID,
		TTL:          h.workerTTL,
		Graylog:      graylog,
	}
	var buffer bytes.Buffer
	if err := tmpl.Execute(&buffer, udataParam); err != nil {
		return "", err
	}

	var udataBegin, udata string

	if withExistingImage {
		log.Debug("spawnWorker> using userdata from existing image")
		udataBegin = `#!/bin/bash
set +e
export CDS_FROM_WORKER_IMAGE="true";
`
	} else {
		log.Debug("spawnWorker> using userdata from worker model")
		udataBegin = `#!/bin/bash
set +e
export CDS_FROM_WORKER_IMAGE="false";
`
	}
	udata = udataBegin + string(udataModel) + buffer.String()

	// Encode again
	udata64 := base64.StdEncoding.EncodeToString([]byte(udata))

	// Create openstack vm
	meta := map[string]string{
		"worker":                     name,
		"hatchery_name":              h.Hatchery().Name,
		"register_only":              fmt.Sprintf("%t", registerOnly),
		"flavor":                     omd.Flavor,
		"model":                      omd.Image,
		"worker_model_name":          model.Name,
		"worker_model_last_modified": fmt.Sprintf("%d", model.UserLastModified.Unix()),
	}

	// Ip len(ipsInfos.ips) > 0, specify one of those
	var ip string
	if len(ipsInfos.ips) > 0 {
		var errai error
		ip, errai = h.findAvailableIP(name)
		if errai != nil {
			return "", errai
		}
		log.Debug("Found %s as available IP", ip)
	}

	networks := []servers.Network{{UUID: h.networkID, FixedIP: ip}}
	r := servers.Create(h.client, servers.CreateOpts{
		Name:      name,
		FlavorRef: flavorID,
		ImageRef:  imageID,
		Metadata:  meta,
		UserData:  []byte(udata64),
		Networks:  networks,
	})

	server, err := r.Extract()
	if err != nil {
		return "", fmt.Errorf("SpawnWorker> Unable to create server: name:%s flavor:%s image:%s metadata:%v networks:%s err:%s body:%s", name, flavorID, imageID, meta, networks, err, r.Body)
	}
	log.Debug("SpawnWorker> Created Server ID: %s", server.ID)
	return name, nil
}
<|endoftext|>


<commit_before>package diskconfig

import (
	"errors"

	"github.com/rackspace/rack/internal/github.com/rackspace/gophercloud/openstack/compute/v2/servers"
)

// DiskConfig represents one of the two possible settings for the DiskConfig option when creating,
// rebuilding, or resizing servers: Auto or Manual.
type DiskConfig string

const (
	// Auto builds a server with a single partition the size of the target flavor disk and
	// automatically adjusts the filesystem to fit the entire partition. Auto may only be used with
	// images and servers that use a single EXT3 partition.
	Auto DiskConfig = "AUTO"

	// Manual builds a server using whatever partition scheme and filesystem are present in the source
	// image. If the target flavor disk is larger, the remaining space is left unpartitioned. This
	// enables images to have non-EXT3 filesystems, multiple partitions, and so on, and enables you
	// to manage the disk configuration. It also results in slightly shorter boot times.
	Manual DiskConfig = "MANUAL"
)

// ErrInvalidDiskConfig is returned if an invalid string is specified for a DiskConfig option.
var ErrInvalidDiskConfig = errors.New("DiskConfig must be either diskconfig.Auto or diskconfig.Manual.")

// Validate ensures that a DiskConfig contains an appropriate value.
func (config DiskConfig) validate() error {
	switch config {
	case Auto, Manual:
		return nil
	default:
		return ErrInvalidDiskConfig
	}
}

// CreateOptsExt adds a DiskConfig option to the base CreateOpts.
type CreateOptsExt struct {
	servers.CreateOptsBuilder

	// DiskConfig [optional] controls how the created server's disk is partitioned.
	DiskConfig DiskConfig `json:"OS-DCF:diskConfig,omitempty"`
}

// ToServerCreateMap adds the diskconfig option to the base server creation options.
func (opts CreateOptsExt) ToServerCreateMap() (map[string]interface{}, error) {
	base, err := opts.CreateOptsBuilder.ToServerCreateMap()
	if err != nil {
		return nil, err
	}

	if string(opts.DiskConfig) == "" {
		return base, nil
	}

	serverMap := base["server"].(map[string]interface{})
	serverMap["OS-DCF:diskConfig"] = string(opts.DiskConfig)

	return base, nil
}

// RebuildOptsExt adds a DiskConfig option to the base RebuildOpts.
type RebuildOptsExt struct {
	servers.RebuildOptsBuilder

	// DiskConfig [optional] controls how the rebuilt server's disk is partitioned.
	DiskConfig DiskConfig
}

// ToServerRebuildMap adds the diskconfig option to the base server rebuild options.
func (opts RebuildOptsExt) ToServerRebuildMap() (map[string]interface{}, error) {
	err := opts.DiskConfig.validate()
	if err != nil {
		return nil, err
	}

	base, err := opts.RebuildOptsBuilder.ToServerRebuildMap()
	if err != nil {
		return nil, err
	}

	serverMap := base["rebuild"].(map[string]interface{})
	serverMap["OS-DCF:diskConfig"] = string(opts.DiskConfig)

	return base, nil
}

// ResizeOptsExt adds a DiskConfig option to the base server resize options.
type ResizeOptsExt struct {
	servers.ResizeOptsBuilder

	// DiskConfig [optional] controls how the resized server's disk is partitioned.
	DiskConfig DiskConfig
}

// ToServerResizeMap adds the diskconfig option to the base server creation options.
func (opts ResizeOptsExt) ToServerResizeMap() (map[string]interface{}, error) {
	err := opts.DiskConfig.validate()
	if err != nil {
		return nil, err
	}

	base, err := opts.ResizeOptsBuilder.ToServerResizeMap()
	if err != nil {
		return nil, err
	}

	serverMap := base["resize"].(map[string]interface{})
	serverMap["OS-DCF:diskConfig"] = string(opts.DiskConfig)

	return base, nil
}
<commit_msg>omit empty disk config<commit_after>package diskconfig

import (
	"errors"

	"github.com/rackspace/rack/internal/github.com/rackspace/gophercloud/openstack/compute/v2/servers"
)

// DiskConfig represents one of the two possible settings for the DiskConfig option when creating,
// rebuilding, or resizing servers: Auto or Manual.
type DiskConfig string

const (
	// Auto builds a server with a single partition the size of the target flavor disk and
	// automatically adjusts the filesystem to fit the entire partition. Auto may only be used with
	// images and servers that use a single EXT3 partition.
	Auto DiskConfig = "AUTO"

	// Manual builds a server using whatever partition scheme and filesystem are present in the source
	// image. If the target flavor disk is larger, the remaining space is left unpartitioned. This
	// enables images to have non-EXT3 filesystems, multiple partitions, and so on, and enables you
	// to manage the disk configuration. It also results in slightly shorter boot times.
	Manual DiskConfig = "MANUAL"
)

// ErrInvalidDiskConfig is returned if an invalid string is specified for a DiskConfig option.
var ErrInvalidDiskConfig = errors.New("DiskConfig must be either diskconfig.Auto or diskconfig.Manual.")

// Validate ensures that a DiskConfig contains an appropriate value.
func (config DiskConfig) validate() error {
	switch config {
	case Auto, Manual, "":
		return nil
	default:
		return ErrInvalidDiskConfig
	}
}

// CreateOptsExt adds a DiskConfig option to the base CreateOpts.
type CreateOptsExt struct {
	servers.CreateOptsBuilder

	// DiskConfig [optional] controls how the created server's disk is partitioned.
	DiskConfig DiskConfig `json:"OS-DCF:diskConfig,omitempty"`
}

// ToServerCreateMap adds the diskconfig option to the base server creation options.
func (opts CreateOptsExt) ToServerCreateMap() (map[string]interface{}, error) {
	base, err := opts.CreateOptsBuilder.ToServerCreateMap()
	if err != nil {
		return nil, err
	}

	if string(opts.DiskConfig) == "" {
		return base, nil
	}

	serverMap := base["server"].(map[string]interface{})
	serverMap["OS-DCF:diskConfig"] = string(opts.DiskConfig)

	return base, nil
}

// RebuildOptsExt adds a DiskConfig option to the base RebuildOpts.
type RebuildOptsExt struct {
	servers.RebuildOptsBuilder

	// DiskConfig [optional] controls how the rebuilt server's disk is partitioned.
	DiskConfig DiskConfig `json:"OS-DCF:diskConfig,omitempty"`
}

// ToServerRebuildMap adds the diskconfig option to the base server rebuild options.
func (opts RebuildOptsExt) ToServerRebuildMap() (map[string]interface{}, error) {
	err := opts.DiskConfig.validate()
	if err != nil {
		return nil, err
	}

	base, err := opts.RebuildOptsBuilder.ToServerRebuildMap()
	if err != nil {
		return nil, err
	}

	serverMap := base["rebuild"].(map[string]interface{})
	if opts.DiskConfig != "" {
		serverMap["OS-DCF:diskConfig"] = string(opts.DiskConfig)
	}

	return base, nil
}

// ResizeOptsExt adds a DiskConfig option to the base server resize options.
type ResizeOptsExt struct {
	servers.ResizeOptsBuilder

	// DiskConfig [optional] controls how the resized server's disk is partitioned.
	DiskConfig DiskConfig `json:"OS-DCF:diskConfig,omitempty"`
}

// ToServerResizeMap adds the diskconfig option to the base server creation options.
func (opts ResizeOptsExt) ToServerResizeMap() (map[string]interface{}, error) {
	err := opts.DiskConfig.validate()
	if err != nil {
		return nil, err
	}

	base, err := opts.ResizeOptsBuilder.ToServerResizeMap()
	if err != nil {
		return nil, err
	}

	serverMap := base["resize"].(map[string]interface{})
	if opts.DiskConfig != "" {
		serverMap["OS-DCF:diskConfig"] = string(opts.DiskConfig)
	}

	return base, nil
}
<|endoftext|>


<commit_before>package mongo

import (
	"gopkg.in/mgo.v2"
	"gopkg.in/mgo.v2/bson"
)

// InitHandler initialize the Handler
func InitHandler(url string) (*Handler, error) {
	return InitHandlerWithDB(url, "")
}

// InitHandlerWithDB initialize the Handler with db name
func InitHandlerWithDB(url, dbName string) (*Handler, error) {
	session, err := mgo.Dial(url)
	if err != nil {
		return nil, err
	}
	return &Handler{
		session: session,
		dbName:  dbName,
	}, nil
}

// Handler Provide the mongo's operation
type Handler struct {
	session *mgo.Session
	dbName  string
}

// Session get current session
func (h *Handler) Session() *mgo.Session {
	return h.session
}

// CloneSession clone session
func (h *Handler) CloneSession() *mgo.Session {
	return h.session.Clone()
}

// DB get database
func (h *Handler) DB() *mgo.Database {
	return h.session.DB(h.dbName)
}

// DBWithName get database with name
func (h *Handler) DBWithName(dbName string) *mgo.Database {
	return h.session.DB(dbName)
}

// CHandle use the default DB processing with a set of callback functions,
// the session will be closed down after processing is completed
func (h *Handler) CHandle(cName string, handle func(c *mgo.Collection)) {
	session := h.CloneSession()
	defer session.Close()
	handle(session.DB(h.dbName).C(cName))
}

// CHandleWithDB use the new DB processing with a set of callback functions,
// the session will be closed down after processing is completed
func (h *Handler) CHandleWithDB(dbName, cName string, handle func(c *mgo.Collection)) {
	session := h.CloneSession()
	defer session.Close()
	handle(session.DB(dbName).C(cName))
}

// C use the default db get collection
func (h *Handler) C(cName string) *mgo.Collection {
	return h.session.DB(h.dbName).C(cName)
}

// CWithDB use the new db get collection
func (h *Handler) CWithDB(dbName, cName string) *mgo.Collection {
	return h.session.DB(dbName).C(cName)
}

// IncrID get increase id
func (h *Handler) IncrID(cName string, storeCName ...string) (id int64, err error) {
	sCName := "counters"
	if len(storeCName) > 0 {
		sCName = storeCName[0]
	}
	h.CHandle(sCName, func(c *mgo.Collection) {
		var result struct {
			Seq int64 `bson:"seq"`
		}
		_, err = c.Find(bson.M{"_id": cName}).Apply(mgo.Change{
			Update:    bson.M{"$inc": bson.M{"seq": 1}},
			ReturnNew: true,
			Upsert:    true,
		}, &result)
		if err == nil {
			id = result.Seq
		}
	})
	return
}
<commit_msg>mongo add close<commit_after>package mongo

import (
	"gopkg.in/mgo.v2"
	"gopkg.in/mgo.v2/bson"
)

// InitHandler initialize the Handler
func InitHandler(url string) (*Handler, error) {
	return InitHandlerWithDB(url, "")
}

// InitHandlerWithDB initialize the Handler with db name
func InitHandlerWithDB(url, dbName string) (*Handler, error) {
	session, err := mgo.Dial(url)
	if err != nil {
		return nil, err
	}
	return &Handler{
		session: session,
		dbName:  dbName,
	}, nil
}

// Handler Provide the mongo's operation
type Handler struct {
	session *mgo.Session
	dbName  string
}

// Close terminates the session
func (h *Handler) Close() {
	h.session.Close()
}

// Session get current session
func (h *Handler) Session() *mgo.Session {
	return h.session
}

// CloneSession clone session
func (h *Handler) CloneSession() *mgo.Session {
	return h.session.Clone()
}

// DB get database
func (h *Handler) DB() *mgo.Database {
	return h.session.DB(h.dbName)
}

// DBWithName get database with name
func (h *Handler) DBWithName(dbName string) *mgo.Database {
	return h.session.DB(dbName)
}

// CHandle use the default DB processing with a set of callback functions,
// the session will be closed down after processing is completed
func (h *Handler) CHandle(cName string, handle func(c *mgo.Collection)) {
	session := h.CloneSession()
	defer session.Close()
	handle(session.DB(h.dbName).C(cName))
}

// CHandleWithDB use the new DB processing with a set of callback functions,
// the session will be closed down after processing is completed
func (h *Handler) CHandleWithDB(dbName, cName string, handle func(c *mgo.Collection)) {
	session := h.CloneSession()
	defer session.Close()
	handle(session.DB(dbName).C(cName))
}

// C use the default db get collection
func (h *Handler) C(cName string) *mgo.Collection {
	return h.session.DB(h.dbName).C(cName)
}

// CWithDB use the new db get collection
func (h *Handler) CWithDB(dbName, cName string) *mgo.Collection {
	return h.session.DB(dbName).C(cName)
}

// IncrID get increase id
func (h *Handler) IncrID(cName string, storeCName ...string) (id int64, err error) {
	sCName := "counters"
	if len(storeCName) > 0 {
		sCName = storeCName[0]
	}
	h.CHandle(sCName, func(c *mgo.Collection) {
		var result struct {
			Seq int64 `bson:"seq"`
		}
		_, err = c.Find(bson.M{"_id": cName}).Apply(mgo.Change{
			Update:    bson.M{"$inc": bson.M{"seq": 1}},
			ReturnNew: true,
			Upsert:    true,
		}, &result)
		if err == nil {
			id = result.Seq
		}
	})
	return
}
<|endoftext|>


<commit_before>// Copyright 2014 Canonical Ltd.
// Licensed under the AGPLv3, see LICENCE file for details.

package mongo

import (
	"crypto/rand"
	"encoding/base64"
	"fmt"
	"net"
	"os"
	"os/exec"
	"path"
	"path/filepath"
	"strings"

	"github.com/juju/errors"
	"github.com/juju/loggo"
	"github.com/juju/replicaset"
	"github.com/juju/utils"
	"github.com/juju/utils/packaging/config"
	"github.com/juju/utils/packaging/manager"
	"gopkg.in/mgo.v2"

	"github.com/juju/juju/network"
	"github.com/juju/juju/service"
	"github.com/juju/juju/version"
)

var (
	logger          = loggo.GetLogger("juju.mongo")
	mongoConfigPath = "/etc/default/mongodb"

	// JujuMongodPath holds the default path to the juju-specific
	// mongod.
	JujuMongodPath = "/usr/lib/juju/bin/mongod"

	// This is NUMACTL package name for apt-get
	numaCtlPkg = "numactl"
)

// WithAddresses represents an entity that has a set of
// addresses. e.g. a state Machine object
type WithAddresses interface {
	Addresses() []network.Address
}

// IsMaster returns a boolean that represents whether the given
// machine's peer address is the primary mongo host for the replicaset
func IsMaster(session *mgo.Session, obj WithAddresses) (bool, error) {
	addrs := obj.Addresses()

	masterHostPort, err := replicaset.MasterHostPort(session)

	// If the replica set has not been configured, then we
	// can have only one master and the caller must
	// be that master.
	if err == replicaset.ErrMasterNotConfigured {
		return true, nil
	}
	if err != nil {
		return false, err
	}

	masterAddr, _, err := net.SplitHostPort(masterHostPort)
	if err != nil {
		return false, err
	}

	for _, addr := range addrs {
		if addr.Value == masterAddr {
			return true, nil
		}
	}
	return false, nil
}

// SelectPeerAddress returns the address to use as the
// mongo replica set peer address by selecting it from the given addresses.
func SelectPeerAddress(addrs []network.Address) string {
	return network.SelectInternalAddress(addrs, false)
}

// SelectPeerHostPort returns the HostPort to use as the
// mongo replica set peer by selecting it from the given hostPorts.
func SelectPeerHostPort(hostPorts []network.HostPort) string {
	return network.SelectInternalHostPort(hostPorts, false)
}

// GenerateSharedSecret generates a pseudo-random shared secret (keyfile)
// for use with Mongo replica sets.
func GenerateSharedSecret() (string, error) {
	// "A key’s length must be between 6 and 1024 characters and may
	// only contain characters in the base64 set."
	//   -- http://docs.mongodb.org/manual/tutorial/generate-key-file/
	buf := make([]byte, base64.StdEncoding.DecodedLen(1024))
	if _, err := rand.Read(buf); err != nil {
		return "", fmt.Errorf("cannot read random secret: %v", err)
	}
	return base64.StdEncoding.EncodeToString(buf), nil
}

// Path returns the executable path to be used to run mongod on this
// machine. If the juju-bundled version of mongo exists, it will return that
// path, otherwise it will return the command to run mongod from the path.
func Path() (string, error) {
	if _, err := os.Stat(JujuMongodPath); err == nil {
		return JujuMongodPath, nil
	}

	path, err := exec.LookPath("mongod")
	if err != nil {
		logger.Infof("could not find %v or mongod in $PATH", JujuMongodPath)
		return "", err
	}
	return path, nil
}

// EnsureServerParams is a parameter struct for EnsureServer.
type EnsureServerParams struct {
	// APIPort is the port to connect to the api server.
	APIPort int

	// StatePort is the port to connect to the mongo server.
	StatePort int

	// Cert is the certificate.
	Cert string

	// PrivateKey is the certificate's private key.
	PrivateKey string

	// CAPrivateKey is the CA certificate's private key.
	CAPrivateKey string

	// SharedSecret is a secret shared between mongo servers.
	SharedSecret string

	// SystemIdentity is the identity of the system.
	SystemIdentity string

	// DataDir is the machine agent data directory.
	DataDir string

	// Namespace is the machine agent's namespace, which is used to
	// generate a unique service name for Mongo.
	Namespace string

	// OplogSize is the size of the Mongo oplog.
	// If this is zero, then EnsureServer will
	// calculate a default size according to the
	// algorithm defined in Mongo.
	OplogSize int

	// SetNumaControlPolicy preference - whether the user
	// wants to set the numa control policy when starting mongo.
	SetNumaControlPolicy bool
}

// EnsureServer ensures that the MongoDB server is installed,
// configured, and ready to run.
//
// This method will remove old versions of the mongo init service as necessary
// before installing the new version.
//
// The namespace is a unique identifier to prevent multiple instances of mongo
// on this machine from colliding. This should be empty unless using
// the local provider.
func EnsureServer(args EnsureServerParams) error {
	logger.Infof(
		"Ensuring mongo server is running; data directory %s; port %d",
		args.DataDir, args.StatePort,
	)

	dbDir := filepath.Join(args.DataDir, "db")
	if err := os.MkdirAll(dbDir, 0700); err != nil {
		return fmt.Errorf("cannot create mongo database directory: %v", err)
	}

	oplogSizeMB := args.OplogSize
	if oplogSizeMB == 0 {
		var err error
		if oplogSizeMB, err = defaultOplogSize(dbDir); err != nil {
			return err
		}
	}

	if err := installMongod(args.SetNumaControlPolicy); err != nil {
		// This isn't treated as fatal because the Juju MongoDB
		// package is likely to be already installed anyway. There
		// could just be a temporary issue with apt-get/yum/whatever
		// and we don't want this to stop jujud from starting.
		// (LP #1441904)
		logger.Errorf("cannot install/upgrade mongod (will proceed anyway): %v", err)
	}
	mongoPath, err := Path()
	if err != nil {
		return err
	}
	logVersion(mongoPath)

	svcConf := newConf(args.DataDir, dbDir, mongoPath, args.StatePort, oplogSizeMB, args.SetNumaControlPolicy)
	svc, err := newService(ServiceName(args.Namespace), svcConf)
	if err != nil {
		return err
	}
	installed, err := svc.Installed()
	if err != nil {
		return errors.Trace(err)
	}
	if installed {
		exists, err := svc.Exists()
		if err != nil {
			return errors.Trace(err)
		}
		if exists {
			logger.Debugf("mongo exists as expected")
			running, err := svc.Running()
			if err != nil {
				return errors.Trace(err)
			}
			if !running {
				return svc.Start()
			}
			return nil
		}
	}

	if err := UpdateSSLKey(args.DataDir, args.Cert, args.PrivateKey); err != nil {
		return err
	}

	err = utils.AtomicWriteFile(sharedSecretPath(args.DataDir), []byte(args.SharedSecret), 0600)
	if err != nil {
		return fmt.Errorf("cannot write mongod shared secret: %v", err)
	}

	// Disable the default mongodb installed by the mongodb-server package.
	// Only do this if the file doesn't exist already, so users can run
	// their own mongodb server if they wish to.
	if _, err := os.Stat(mongoConfigPath); os.IsNotExist(err) {
		err = utils.AtomicWriteFile(
			mongoConfigPath,
			[]byte("ENABLE_MONGODB=no"),
			0644,
		)
		if err != nil {
			return err
		}
	}

	if err := svc.Stop(); err != nil {
		return errors.Annotatef(err, "failed to stop mongo")
	}
	if err := makeJournalDirs(dbDir); err != nil {
		return fmt.Errorf("error creating journal directories: %v", err)
	}
	if err := preallocOplog(dbDir, oplogSizeMB); err != nil {
		return fmt.Errorf("error creating oplog files: %v", err)
	}
	if err := service.InstallAndStart(svc); err != nil {
		return errors.Trace(err)
	}
	return nil
}

// UpdateSSLKey writes a new SSL key used by mongo to validate connections from Juju state server(s)
func UpdateSSLKey(dataDir, cert, privateKey string) error {
	certKey := cert + "\n" + privateKey
	err := utils.AtomicWriteFile(sslKeyPath(dataDir), []byte(certKey), 0600)
	return errors.Annotate(err, "cannot write SSL key")
}

func makeJournalDirs(dataDir string) error {
	journalDir := path.Join(dataDir, "journal")
	if err := os.MkdirAll(journalDir, 0700); err != nil {
		logger.Errorf("failed to make mongo journal dir %s: %v", journalDir, err)
		return err
	}

	// Manually create the prealloc files, since otherwise they get
	// created as 100M files. We create three files of 1MB each.
	prefix := filepath.Join(journalDir, "prealloc.")
	preallocSize := 1024 * 1024
	return preallocFiles(prefix, preallocSize, preallocSize, preallocSize)
}

func logVersion(mongoPath string) {
	cmd := exec.Command(mongoPath, "--version")
	output, err := cmd.CombinedOutput()
	if err != nil {
		logger.Infof("failed to read the output from %s --version: %v", mongoPath, err)
		return
	}
	logger.Debugf("using mongod: %s --version: %q", mongoPath, output)
}

// getPackageManager is a helper function which returns the
// package manager implementation for the current system.
func getPackageManager() (manager.PackageManager, error) {
	return manager.NewPackageManager(version.Current.Series)
}

// getPackagingConfigurer is a helper function which returns the
// packaging configuration manager for the current system.
func getPackagingConfigurer() (config.PackagingConfigurer, error) {
	return config.NewPackagingConfigurer(version.Current.Series)
}

func installMongod(numaCtl bool) error {
	series := version.Current.Series

	pacconfer, err := getPackagingConfigurer()
	if err != nil {
		return err
	}

	pacman, err := getPackageManager()
	if err != nil {
		return err
	}

	// Only Quantal requires the PPA.
	if series == "quantal" {
		// install python-software-properties:
		if err := pacman.InstallPrerequisite(); err != nil {
			return err
		}
		if err := pacman.AddRepository("ppa:juju/stable"); err != nil {
			return err
		}
	}
	// CentOS requires "epel-release" for the epel repo mongodb-server is in.
	if series == "centos7" {
		// install epel-release
		if err := pacman.Install("epel-release"); err != nil {
			return err
		}
	}

	mongoPkg := packageForSeries(series)

	pkgs := []string{mongoPkg}
	if numaCtl {
		pkgs = []string{mongoPkg, numaCtlPkg}
		logger.Infof("installing %s and %s", mongoPkg, numaCtlPkg)
	} else {
		logger.Infof("installing %s", mongoPkg)
	}

	for i, _ := range pkgs {
		// apply release targeting if needed.
		if pacconfer.IsCloudArchivePackage(pkgs[i]) {
			pkgs[i] = strings.Join(pacconfer.ApplyCloudArchiveTarget(pkgs[i]), " ")
		}

		if err := pacman.Install(pkgs[i]); err != nil {
			return err
		}
	}

	// Work around SELinux on centos7
	if series == "centos7" {
		cmd := []string{"chcon", "-R", "-v", "-t", "mongod_var_lib_t", "/var/lib/juju/"}
		logger.Infof("running %s %v", cmd[0], cmd[1:])
		_, err = utils.RunCommand(cmd[0], cmd[1:]...)
		if err != nil {
			logger.Infof("chcon error %s", err)
			logger.Infof("chcon error %s", err.Error())
			return err
		}

		cmd = []string{"semanage", "port", "-a", "-t", "mongod_port_t", "-p", "tcp", "37017"}
		logger.Infof("running %s %v", cmd[0], cmd[1:])
		_, err = utils.RunCommand(cmd[0], cmd[1:]...)
		if err != nil {
			if !strings.Contains(err.Error(), "exit status 1") {
				return err
			}
		}
	}

	return nil
}

// packageForSeries returns the name of the mongo package for the series
// of the machine that it is going to be running on.
func packageForSeries(series string) string {
	switch series {
	case "precise", "quantal", "raring", "saucy", "centos7":
		return "mongodb-server"
	default:
		// trusty and onwards
		return "juju-mongodb"
	}
}

// noauthCommand returns an os/exec.Cmd that may be executed to
// run mongod without security.
func noauthCommand(dataDir string, port int) (*exec.Cmd, error) {
	sslKeyFile := path.Join(dataDir, "server.pem")
	dbDir := filepath.Join(dataDir, "db")
	mongoPath, err := Path()
	if err != nil {
		return nil, err
	}
	cmd := exec.Command(mongoPath,
		"--noauth",
		"--dbpath", dbDir,
		"--sslOnNormalPorts",
		"--sslPEMKeyFile", sslKeyFile,
		"--sslPEMKeyPassword", "ignored",
		"--bind_ip", "127.0.0.1",
		"--port", fmt.Sprint(port),
		"--noprealloc",
		"--syslog",
		"--smallfiles",
		"--journal",
	)
	return cmd, nil
}
<commit_msg>mongo: make it clear mongo installation depends on operatingsystem<commit_after>// Copyright 2014 Canonical Ltd.
// Licensed under the AGPLv3, see LICENCE file for details.

package mongo

import (
	"crypto/rand"
	"encoding/base64"
	"fmt"
	"net"
	"os"
	"os/exec"
	"path"
	"path/filepath"
	"strings"

	"github.com/juju/errors"
	"github.com/juju/loggo"
	"github.com/juju/replicaset"
	"github.com/juju/utils"
	"github.com/juju/utils/packaging/config"
	"github.com/juju/utils/packaging/manager"
	"gopkg.in/mgo.v2"

	"github.com/juju/juju/network"
	"github.com/juju/juju/service"
	"github.com/juju/juju/version"
)

var (
	logger          = loggo.GetLogger("juju.mongo")
	mongoConfigPath = "/etc/default/mongodb"

	// JujuMongodPath holds the default path to the juju-specific
	// mongod.
	JujuMongodPath = "/usr/lib/juju/bin/mongod"

	// This is NUMACTL package name for apt-get
	numaCtlPkg = "numactl"
)

// WithAddresses represents an entity that has a set of
// addresses. e.g. a state Machine object
type WithAddresses interface {
	Addresses() []network.Address
}

// IsMaster returns a boolean that represents whether the given
// machine's peer address is the primary mongo host for the replicaset
func IsMaster(session *mgo.Session, obj WithAddresses) (bool, error) {
	addrs := obj.Addresses()

	masterHostPort, err := replicaset.MasterHostPort(session)

	// If the replica set has not been configured, then we
	// can have only one master and the caller must
	// be that master.
	if err == replicaset.ErrMasterNotConfigured {
		return true, nil
	}
	if err != nil {
		return false, err
	}

	masterAddr, _, err := net.SplitHostPort(masterHostPort)
	if err != nil {
		return false, err
	}

	for _, addr := range addrs {
		if addr.Value == masterAddr {
			return true, nil
		}
	}
	return false, nil
}

// SelectPeerAddress returns the address to use as the
// mongo replica set peer address by selecting it from the given addresses.
func SelectPeerAddress(addrs []network.Address) string {
	return network.SelectInternalAddress(addrs, false)
}

// SelectPeerHostPort returns the HostPort to use as the
// mongo replica set peer by selecting it from the given hostPorts.
func SelectPeerHostPort(hostPorts []network.HostPort) string {
	return network.SelectInternalHostPort(hostPorts, false)
}

// GenerateSharedSecret generates a pseudo-random shared secret (keyfile)
// for use with Mongo replica sets.
func GenerateSharedSecret() (string, error) {
	// "A key’s length must be between 6 and 1024 characters and may
	// only contain characters in the base64 set."
	//   -- http://docs.mongodb.org/manual/tutorial/generate-key-file/
	buf := make([]byte, base64.StdEncoding.DecodedLen(1024))
	if _, err := rand.Read(buf); err != nil {
		return "", fmt.Errorf("cannot read random secret: %v", err)
	}
	return base64.StdEncoding.EncodeToString(buf), nil
}

// Path returns the executable path to be used to run mongod on this
// machine. If the juju-bundled version of mongo exists, it will return that
// path, otherwise it will return the command to run mongod from the path.
func Path() (string, error) {
	if _, err := os.Stat(JujuMongodPath); err == nil {
		return JujuMongodPath, nil
	}

	path, err := exec.LookPath("mongod")
	if err != nil {
		logger.Infof("could not find %v or mongod in $PATH", JujuMongodPath)
		return "", err
	}
	return path, nil
}

// EnsureServerParams is a parameter struct for EnsureServer.
type EnsureServerParams struct {
	// APIPort is the port to connect to the api server.
	APIPort int

	// StatePort is the port to connect to the mongo server.
	StatePort int

	// Cert is the certificate.
	Cert string

	// PrivateKey is the certificate's private key.
	PrivateKey string

	// CAPrivateKey is the CA certificate's private key.
	CAPrivateKey string

	// SharedSecret is a secret shared between mongo servers.
	SharedSecret string

	// SystemIdentity is the identity of the system.
	SystemIdentity string

	// DataDir is the machine agent data directory.
	DataDir string

	// Namespace is the machine agent's namespace, which is used to
	// generate a unique service name for Mongo.
	Namespace string

	// OplogSize is the size of the Mongo oplog.
	// If this is zero, then EnsureServer will
	// calculate a default size according to the
	// algorithm defined in Mongo.
	OplogSize int

	// SetNumaControlPolicy preference - whether the user
	// wants to set the numa control policy when starting mongo.
	SetNumaControlPolicy bool
}

// EnsureServer ensures that the MongoDB server is installed,
// configured, and ready to run.
//
// This method will remove old versions of the mongo init service as necessary
// before installing the new version.
//
// The namespace is a unique identifier to prevent multiple instances of mongo
// on this machine from colliding. This should be empty unless using
// the local provider.
func EnsureServer(args EnsureServerParams) error {
	logger.Infof(
		"Ensuring mongo server is running; data directory %s; port %d",
		args.DataDir, args.StatePort,
	)

	dbDir := filepath.Join(args.DataDir, "db")
	if err := os.MkdirAll(dbDir, 0700); err != nil {
		return fmt.Errorf("cannot create mongo database directory: %v", err)
	}

	oplogSizeMB := args.OplogSize
	if oplogSizeMB == 0 {
		var err error
		if oplogSizeMB, err = defaultOplogSize(dbDir); err != nil {
			return err
		}
	}

	operatingsystem := version.Current.Series
	if err := installMongod(operatingsystem, args.SetNumaControlPolicy); err != nil {
		// This isn't treated as fatal because the Juju MongoDB
		// package is likely to be already installed anyway. There
		// could just be a temporary issue with apt-get/yum/whatever
		// and we don't want this to stop jujud from starting.
		// (LP #1441904)
		logger.Errorf("cannot install/upgrade mongod (will proceed anyway): %v", err)
	}
	mongoPath, err := Path()
	if err != nil {
		return err
	}
	logVersion(mongoPath)

	svcConf := newConf(args.DataDir, dbDir, mongoPath, args.StatePort, oplogSizeMB, args.SetNumaControlPolicy)
	svc, err := newService(ServiceName(args.Namespace), svcConf)
	if err != nil {
		return err
	}
	installed, err := svc.Installed()
	if err != nil {
		return errors.Trace(err)
	}
	if installed {
		exists, err := svc.Exists()
		if err != nil {
			return errors.Trace(err)
		}
		if exists {
			logger.Debugf("mongo exists as expected")
			running, err := svc.Running()
			if err != nil {
				return errors.Trace(err)
			}
			if !running {
				return svc.Start()
			}
			return nil
		}
	}

	if err := UpdateSSLKey(args.DataDir, args.Cert, args.PrivateKey); err != nil {
		return err
	}

	err = utils.AtomicWriteFile(sharedSecretPath(args.DataDir), []byte(args.SharedSecret), 0600)
	if err != nil {
		return fmt.Errorf("cannot write mongod shared secret: %v", err)
	}

	// Disable the default mongodb installed by the mongodb-server package.
	// Only do this if the file doesn't exist already, so users can run
	// their own mongodb server if they wish to.
	if _, err := os.Stat(mongoConfigPath); os.IsNotExist(err) {
		err = utils.AtomicWriteFile(
			mongoConfigPath,
			[]byte("ENABLE_MONGODB=no"),
			0644,
		)
		if err != nil {
			return err
		}
	}

	if err := svc.Stop(); err != nil {
		return errors.Annotatef(err, "failed to stop mongo")
	}
	if err := makeJournalDirs(dbDir); err != nil {
		return fmt.Errorf("error creating journal directories: %v", err)
	}
	if err := preallocOplog(dbDir, oplogSizeMB); err != nil {
		return fmt.Errorf("error creating oplog files: %v", err)
	}
	if err := service.InstallAndStart(svc); err != nil {
		return errors.Trace(err)
	}
	return nil
}

// UpdateSSLKey writes a new SSL key used by mongo to validate connections from Juju state server(s)
func UpdateSSLKey(dataDir, cert, privateKey string) error {
	certKey := cert + "\n" + privateKey
	err := utils.AtomicWriteFile(sslKeyPath(dataDir), []byte(certKey), 0600)
	return errors.Annotate(err, "cannot write SSL key")
}

func makeJournalDirs(dataDir string) error {
	journalDir := path.Join(dataDir, "journal")
	if err := os.MkdirAll(journalDir, 0700); err != nil {
		logger.Errorf("failed to make mongo journal dir %s: %v", journalDir, err)
		return err
	}

	// Manually create the prealloc files, since otherwise they get
	// created as 100M files. We create three files of 1MB each.
	prefix := filepath.Join(journalDir, "prealloc.")
	preallocSize := 1024 * 1024
	return preallocFiles(prefix, preallocSize, preallocSize, preallocSize)
}

func logVersion(mongoPath string) {
	cmd := exec.Command(mongoPath, "--version")
	output, err := cmd.CombinedOutput()
	if err != nil {
		logger.Infof("failed to read the output from %s --version: %v", mongoPath, err)
		return
	}
	logger.Debugf("using mongod: %s --version: %q", mongoPath, output)
}

func installMongod(operatingsystem string, numaCtl bool) error {
	// fetch the packaging configuration manager for the current operating system.
	pacconfer, err := config.NewPackagingConfigurer(operatingsystem)
	if err != nil {
		return err
	}

	// fetch the package manager implementation for the current operating system.
	pacman, err := manager.NewPackageManager(operatingsystem)
	if err != nil {
		return err
	}

	// Only Quantal requires the PPA.
	if operatingsystem == "quantal" {
		// install python-software-properties:
		if err := pacman.InstallPrerequisite(); err != nil {
			return err
		}
		if err := pacman.AddRepository("ppa:juju/stable"); err != nil {
			return err
		}
	}
	// CentOS requires "epel-release" for the epel repo mongodb-server is in.
	if operatingsystem == "centos7" {
		// install epel-release
		if err := pacman.Install("epel-release"); err != nil {
			return err
		}
	}

	mongoPkg := packageForSeries(operatingsystem)

	pkgs := []string{mongoPkg}
	if numaCtl {
		pkgs = []string{mongoPkg, numaCtlPkg}
		logger.Infof("installing %s and %s", mongoPkg, numaCtlPkg)
	} else {
		logger.Infof("installing %s", mongoPkg)
	}

	for i, _ := range pkgs {
		// apply release targeting if needed.
		if pacconfer.IsCloudArchivePackage(pkgs[i]) {
			pkgs[i] = strings.Join(pacconfer.ApplyCloudArchiveTarget(pkgs[i]), " ")
		}

		if err := pacman.Install(pkgs[i]); err != nil {
			return err
		}
	}

	// Work around SELinux on centos7
	if operatingsystem == "centos7" {
		cmd := []string{"chcon", "-R", "-v", "-t", "mongod_var_lib_t", "/var/lib/juju/"}
		logger.Infof("running %s %v", cmd[0], cmd[1:])
		_, err = utils.RunCommand(cmd[0], cmd[1:]...)
		if err != nil {
			logger.Infof("chcon error %s", err)
			logger.Infof("chcon error %s", err.Error())
			return err
		}

		cmd = []string{"semanage", "port", "-a", "-t", "mongod_port_t", "-p", "tcp", "37017"}
		logger.Infof("running %s %v", cmd[0], cmd[1:])
		_, err = utils.RunCommand(cmd[0], cmd[1:]...)
		if err != nil {
			if !strings.Contains(err.Error(), "exit status 1") {
				return err
			}
		}
	}

	return nil
}

// packageForSeries returns the name of the mongo package for the series
// of the machine that it is going to be running on.
func packageForSeries(series string) string {
	switch series {
	case "precise", "quantal", "raring", "saucy", "centos7":
		return "mongodb-server"
	default:
		// trusty and onwards
		return "juju-mongodb"
	}
}

// noauthCommand returns an os/exec.Cmd that may be executed to
// run mongod without security.
func noauthCommand(dataDir string, port int) (*exec.Cmd, error) {
	sslKeyFile := path.Join(dataDir, "server.pem")
	dbDir := filepath.Join(dataDir, "db")
	mongoPath, err := Path()
	if err != nil {
		return nil, err
	}
	cmd := exec.Command(mongoPath,
		"--noauth",
		"--dbpath", dbDir,
		"--sslOnNormalPorts",
		"--sslPEMKeyFile", sslKeyFile,
		"--sslPEMKeyPassword", "ignored",
		"--bind_ip", "127.0.0.1",
		"--port", fmt.Sprint(port),
		"--noprealloc",
		"--syslog",
		"--smallfiles",
		"--journal",
	)
	return cmd, nil
}
<|endoftext|>


<commit_before>package monit

import (
	"fmt"
	"log"
	"os"
	"runtime/debug"
	"sync"
	"time"

	"github.com/tolexo/aero/conf"
)

var (
	logFp      *os.File
	fileErr    error
	logger     *log.Logger
	prevDay    int
	currentDay int
	isLog      bool
	syncOnce   sync.Once
)

func createPanicLog(sTime time.Time, panicMsg interface{}) {
	path := conf.String("logs.panic_log", "panic_log")
	path = fmt.Sprintf("%s_%d-%d-%d.log", path, sTime.Day(), sTime.Month(), sTime.Year())
	if logFp, fileErr = os.OpenFile(path, os.O_RDWR|os.O_CREATE|os.O_APPEND, 0666); fileErr != nil {
		fmt.Println("Could not create the panic log file")
		panic(panicMsg)
	}
}

func logPanic(panicMsg interface{}, serviceID string, requestURI string) {
	logger = log.New(logFp, "panic", log.Lshortfile)
	logFp.WriteString("\n")
	logFp.WriteString("serviceID: " + serviceID + " requestURI: " + requestURI + "\n")
	logger.Print(string(debug.Stack()))
	logger.Panic(panicMsg)
}

func PanicLogger(panicMsg interface{}, serviceID string, requestURI string) {

	syncOnce.Do(func() {
		isLog = conf.Bool("monitor.panic_log", false)
	})

	if isLog == true {
		sTime := time.Now().UTC()
		currentDay = sTime.Day()
		if currentDay != prevDay {
			if prevDay != 0 && logFp != nil {
				logFp.Close()
			}
			createPanicLog(sTime, panicMsg)
			prevDay = currentDay
		}
		if logFp != nil {
			logPanic(panicMsg, serviceID, requestURI)
		} else {
			fmt.Println("lost the pointer to the log file, creating again")
			createPanicLog(sTime, panicMsg)
			logPanic(panicMsg)
		}
	} else {
		panic(panicMsg)
	}
}
<commit_msg>TTA-226A code tweak<commit_after>package monit

import (
	"fmt"
	"log"
	"os"
	"runtime/debug"
	"sync"
	"time"

	"github.com/tolexo/aero/conf"
)

var (
	logFp      *os.File
	fileErr    error
	logger     *log.Logger
	prevDay    int
	currentDay int
	isLog      bool
	syncOnce   sync.Once
)

func createPanicLog(sTime time.Time, panicMsg interface{}) {
	path := conf.String("logs.panic_log", "panic_log")
	path = fmt.Sprintf("%s_%d-%d-%d.log", path, sTime.Day(), sTime.Month(), sTime.Year())
	if logFp, fileErr = os.OpenFile(path, os.O_RDWR|os.O_CREATE|os.O_APPEND, 0666); fileErr != nil {
		fmt.Println("Could not create the panic log file")
		panic(panicMsg)
	}
}

func logPanic(panicMsg interface{}, serviceID string, requestURI string) {
	logger = log.New(logFp, "panic", log.Lshortfile)
	logFp.WriteString("\n")
	logFp.WriteString("serviceID: " + serviceID + " requestURI: " + requestURI + "\n")
	logger.Print(string(debug.Stack()))
	logger.Panic(panicMsg)
}

func PanicLogger(panicMsg interface{}, serviceID string, requestURI string) {

	syncOnce.Do(func() {
		isLog = conf.Bool("monitor.panic_log", false)
	})

	if isLog == true {
		sTime := time.Now().UTC()
		currentDay = sTime.Day()
		if currentDay != prevDay {
			if prevDay != 0 && logFp != nil {
				logFp.Close()
			}
			createPanicLog(sTime, panicMsg)
			prevDay = currentDay
		}
		if logFp != nil {
			logPanic(panicMsg, serviceID, requestURI)
		} else {
			fmt.Println("lost the pointer to the log file, creating again")
			createPanicLog(sTime, panicMsg)
			logPanic(panicMsg, serviceID, requestURI)
		}
	} else {
		panic(panicMsg)
	}
}
<|endoftext|>


<commit_before>package monit

import (
	"fmt"
	"log"
	"os"
	"runtime/debug"
	"sync"
	"time"

	"github.com/tolexo/aero/conf"
)

var (
	logFp      *os.File
	fileErr    error
	logger     *log.Logger
	prevDay    int
	currentDay int
	isLog      bool
	syncOnce   sync.Once
)

func createPanicLog(sTime time.Time, panicMsg interface{}) {
	path := conf.String("logs.panic_log", "panic_log")
	path = fmt.Sprintf("%s_%d-%d-%d.log", path, sTime.Day(), sTime.Month(), sTime.Year())
	if logFp, fileErr = os.OpenFile(path, os.O_RDWR|os.O_CREATE|os.O_APPEND, 0666); fileErr != nil {
		fmt.Println("Could not create the panic log file")
		panic(panicMsg)
	}
}

func logPanic(panicMsg interface{}, serviceID string, requestURI string, curTime time.Time) {
	logger = log.New(logFp, "", log.Lshortfile)
	logFp.WriteString("\n")
	logger.Print("Time: ", curTime, "  serviceID: ", serviceID, "  requestURI: ", requestURI)
	logger.Print(string(debug.Stack()))
	logger.Panic(panicMsg)
}

func PanicLogger(panicMsg interface{}, serviceID string, requestURI string, curTime time.Time) {

	syncOnce.Do(func() {
		isLog = conf.Bool("monitor.panic_log", false)
	})

	if isLog == true {
		sTime := time.Now().UTC()
		currentDay = sTime.Day()
		if currentDay != prevDay {
			if prevDay != 0 && logFp != nil {
				logFp.Close()
			}
			createPanicLog(sTime, panicMsg)
			prevDay = currentDay
		}
		if logFp != nil {
			logPanic(panicMsg, serviceID, requestURI, curTime)
		} else {
			fmt.Println("lost the pointer to the log file, creating again")
			createPanicLog(sTime, panicMsg)
			logPanic(panicMsg, serviceID, requestURI, curTime)
		}
	} else {
		panic(panicMsg)
	}
}
<commit_msg>TTA-226 log format modified<commit_after>package monit

import (
	"fmt"
	"log"
	"os"
	"runtime/debug"
	"sync"
	"time"

	"github.com/tolexo/aero/conf"
)

var (
	logFp      *os.File
	fileErr    error
	logger     *log.Logger
	prevDay    int
	currentDay int
	isLog      bool
	syncOnce   sync.Once
)

func createPanicLog(sTime time.Time, panicMsg interface{}) {
	path := conf.String("logs.panic_log", "panic_log")
	path = fmt.Sprintf("%s_%d-%d-%d.log", path, sTime.Day(), sTime.Month(), sTime.Year())
	if logFp, fileErr = os.OpenFile(path, os.O_RDWR|os.O_CREATE|os.O_APPEND, 0666); fileErr != nil {
		fmt.Println("Could not create the panic log file")
		panic(panicMsg)
	}
}

func logPanic(panicMsg interface{}, serviceID string, requestURI string, curTime time.Time) {
	logger = log.New(logFp, "\n", log.Lshortfile)
	logger.Print("Time: ", curTime, "  serviceID: ", serviceID, "  requestURI: ", requestURI, "\n", string(debug.Stack()))
	logger.Panic(panicMsg)
}

func PanicLogger(panicMsg interface{}, serviceID string, requestURI string, curTime time.Time) {

	syncOnce.Do(func() {
		isLog = conf.Bool("monitor.panic_log", false)
	})

	if isLog == true {
		sTime := time.Now().UTC()
		currentDay = sTime.Day()
		if currentDay != prevDay {
			if prevDay != 0 && logFp != nil {
				logFp.Close()
			}
			createPanicLog(sTime, panicMsg)
			prevDay = currentDay
		}
		if logFp != nil {
			logPanic(panicMsg, serviceID, requestURI, curTime)
		} else {
			fmt.Println("lost the pointer to the log file, creating again")
			createPanicLog(sTime, panicMsg)
			logPanic(panicMsg, serviceID, requestURI, curTime)
		}
	} else {
		panic(panicMsg)
	}
}
<|endoftext|>


<commit_before>// Copyright 2019 Google LLC
//
// 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 test contains test utility code shared between IC and DAM.
package test

import (
	"fmt"
	"io"
	"net/http"
	"net/http/httptest"
	"regexp"
	"strings"
	"testing"

	"github.com/google/go-cmp/cmp" /* copybara-comment */
	"github.com/GoogleCloudPlatform/healthcare-federated-access-services/lib/persona" /* copybara-comment: persona */
	cpb "github.com/GoogleCloudPlatform/healthcare-federated-access-services/proto/common/v1" /* copybara-comment: go_proto */
	dampb "github.com/GoogleCloudPlatform/healthcare-federated-access-services/proto/dam/v1" /* copybara-comment: go_proto */
)

const (
	// TestClientID is the client ID for tests.
	TestClientID = "00000000-0000-0000-0000-000000000000"
	// TestClientSecret is the client secret for test client.
	TestClientSecret = "00000000-0000-0000-0000-000000000001"
	// TestIssuerURL is the URL of the fake OIDC Issuer service.
	TestIssuerURL = "https://example.org/oidc"
)

var (
	varRE = regexp.MustCompile(`(\$\((.*?)\))`)
)

// HandlerTest holds the test variables for a service handler test.
type HandlerTest struct {
	Name        string
	Method      string
	Path        string
	Input       string
	Params      string
	IsForm      bool
	Persona     string
	Scope       string
	LinkPersona string
	LinkScope   string
	Output      string
	CmpOptions  cmp.Options
	Status      int
}

type serviceHandler interface {
	ServeHTTP(w http.ResponseWriter, r *http.Request)
}

// HandlerTests run tests on a service handler.
func HandlerTests(t *testing.T, h serviceHandler, tests []HandlerTest, issuerURL string, cfg *dampb.DamConfig) {
	t.Helper()

	testOutput := make(map[string]string)
	for _, test := range tests {
		name := test.Name
		if len(name) == 0 {
			name = test.Method + " " + test.Path
		}
		pname := "dr_joe_elixir"
		if len(test.Persona) > 0 {
			pname = test.Persona
		}
		var p *cpb.TestPersona
		if cfg != nil {
			p = cfg.TestPersonas[pname]
		}
		acTok, _, err := persona.NewAccessToken(pname, issuerURL, TestClientID, test.Scope, p)
		if err != nil {
			t.Fatalf("persona.NewAccessToken(%q, %q, _, _) failed: %v", pname, issuerURL, err)
		}
		target := fmt.Sprintf("%s?client_id=%s&client_secret=%s", test.Path, TestClientID, TestClientSecret)
		if len(test.Params) > 0 {
			target += "&" + test.Params
		}
		var Input io.Reader
		varInput := false
		InputStr := test.Input
		if len(test.Input) > 0 {
			for match := varRE.FindStringSubmatch(InputStr); match != nil; match = varRE.FindStringSubmatch(InputStr) {
				varInput = true
				ref := match[2]
				if len(testOutput[ref]) == 0 {
					t.Fatalf("test %q uses %q, but %q must be tested separately earlier in the test list", name, ref, ref)
				}
				InputStr = strings.Replace(InputStr, match[1], testOutput[ref], -1)
			}
			Input = strings.NewReader(InputStr)
		}
		r := httptest.NewRequest(test.Method, target, Input)
		if test.IsForm {
			r.Header.Set("Content-Type", "application/x-www-form-urlencoded")
		}
		r.Header.Set("Authorization", "Bearer "+string(acTok))
		if len(test.LinkPersona) > 0 {
			linkTok, _, err := persona.NewAccessToken(test.LinkPersona, issuerURL, TestClientID, test.LinkScope, nil)
			if err != nil {
				t.Fatalf("persona.NewAccessToken(%q, %q, _, _) failed: %v", test.LinkPersona, issuerURL, err)
			}
			r.Header.Set("X-Link-Authorization", "Bearer "+string(linkTok))
		}
		w := httptest.NewRecorder()
		h.ServeHTTP(w, r)
		Output := w.Body.String()
		testOutput[name] = Output
		hasError := false

		if w.Code != test.Status {
			hasError = true
			t.Errorf("test %q returned wrong status code: got %d want %d", name, w.Code, test.Status)
		}
		switch {
		case strings.HasPrefix(test.Output, "^"):
			re, err := regexp.Compile(test.Output)
			if err != nil {
				t.Fatalf("test %q cannot compile regexp Output %q: %v", name, test.Output, err)
			}
			if !re.Match([]byte(Output)) {
				hasError = true
				t.Errorf("test %q returned unexpected body, got:\n%s\n\nwant regexp match of:\n%s", name, Output, test.Output)
			}
		case strings.HasPrefix(test.Output, "*") && strings.HasSuffix(test.Output, "*"):
			q := regexp.QuoteMeta(test.Output[1 : len(test.Output)-1])
			q = strings.Replace(q, "\\*", ".*", -1)
			re, err := regexp.Compile(q)
			if err != nil {
				t.Fatalf("test %q cannot compile regexp Output %q with regex %q: %v", name, test.Output, q, err)
			}
			if !re.Match([]byte(Output)) {
				hasError = true
				t.Errorf("test %q returned unexpected body, got:\n%s\n\nwant substring match of:\n%s\n\nregex:\n%s", name, Output, test.Output, q)
			}
		default:
			if diff := cmp.Diff(test.Output, Output, test.CmpOptions); diff != "" {
				hasError = true
				t.Errorf("test %q returned mismatching body (-want +got):\n%s", name, diff)
			}
		}
		if hasError && varInput {
			t.Logf("test %q Input value was: %s", name, InputStr)
		}
	}
}
<commit_msg>Use sub tests for HandlerTests<commit_after>// Copyright 2019 Google LLC
//
// 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 test contains test utility code shared between IC and DAM.
package test

import (
	"fmt"
	"io"
	"net/http"
	"net/http/httptest"
	"regexp"
	"strings"
	"testing"

	"github.com/google/go-cmp/cmp" /* copybara-comment */
	"github.com/GoogleCloudPlatform/healthcare-federated-access-services/lib/persona" /* copybara-comment: persona */

	cpb "github.com/GoogleCloudPlatform/healthcare-federated-access-services/proto/common/v1" /* copybara-comment: go_proto */
	dampb "github.com/GoogleCloudPlatform/healthcare-federated-access-services/proto/dam/v1" /* copybara-comment: go_proto */
)

const (
	// TestClientID is the client ID for tests.
	TestClientID = "00000000-0000-0000-0000-000000000000"
	// TestClientSecret is the client secret for test client.
	TestClientSecret = "00000000-0000-0000-0000-000000000001"
	// TestIssuerURL is the URL of the fake OIDC Issuer service.
	TestIssuerURL = "https://example.org/oidc"
)

var (
	varRE = regexp.MustCompile(`(\$\((.*?)\))`)
)

// HandlerTest holds the test variables for a service handler test.
type HandlerTest struct {
	Name        string
	Method      string
	Path        string
	Input       string
	Params      string
	IsForm      bool
	Persona     string
	Scope       string
	LinkPersona string
	LinkScope   string
	Output      string
	CmpOptions  cmp.Options
	Status      int
}

type serviceHandler interface {
	ServeHTTP(w http.ResponseWriter, r *http.Request)
}

// HandlerTests run tests on a service handler.
func HandlerTests(t *testing.T, h serviceHandler, tests []HandlerTest, issuerURL string, cfg *dampb.DamConfig) {
	t.Helper()

	testOutput := make(map[string]string)
	for _, test := range tests {
		t.Run(test.Name, func(t *testing.T) {
			name := test.Name
			if len(name) == 0 {
				name = test.Method + " " + test.Path
			}
			pname := "dr_joe_elixir"
			if len(test.Persona) > 0 {
				pname = test.Persona
			}
			var p *cpb.TestPersona
			if cfg != nil {
				p = cfg.TestPersonas[pname]
			}
			acTok, _, err := persona.NewAccessToken(pname, issuerURL, TestClientID, test.Scope, p)
			if err != nil {
				t.Fatalf("persona.NewAccessToken(%q, %q, _, _) failed: %v", pname, issuerURL, err)
			}
			target := fmt.Sprintf("%s?client_id=%s&client_secret=%s", test.Path, TestClientID, TestClientSecret)
			if len(test.Params) > 0 {
				target += "&" + test.Params
			}
			var Input io.Reader
			varInput := false
			InputStr := test.Input
			if len(test.Input) > 0 {
				for match := varRE.FindStringSubmatch(InputStr); match != nil; match = varRE.FindStringSubmatch(InputStr) {
					varInput = true
					ref := match[2]
					if len(testOutput[ref]) == 0 {
						t.Fatalf("test %q uses %q, but %q must be tested separately earlier in the test list", name, ref, ref)
					}
					InputStr = strings.Replace(InputStr, match[1], testOutput[ref], -1)
				}
				Input = strings.NewReader(InputStr)
			}
			r := httptest.NewRequest(test.Method, target, Input)
			if test.IsForm {
				r.Header.Set("Content-Type", "application/x-www-form-urlencoded")
			}
			r.Header.Set("Authorization", "Bearer "+string(acTok))
			if len(test.LinkPersona) > 0 {
				linkTok, _, err := persona.NewAccessToken(test.LinkPersona, issuerURL, TestClientID, test.LinkScope, nil)
				if err != nil {
					t.Fatalf("persona.NewAccessToken(%q, %q, _, _) failed: %v", test.LinkPersona, issuerURL, err)
				}
				r.Header.Set("X-Link-Authorization", "Bearer "+string(linkTok))
			}
			w := httptest.NewRecorder()
			h.ServeHTTP(w, r)
			Output := w.Body.String()
			testOutput[name] = Output
			hasError := false

			if w.Code != test.Status {
				hasError = true
				t.Errorf("test %q returned wrong status code: got %d want %d", name, w.Code, test.Status)
			}
			switch {
			case strings.HasPrefix(test.Output, "^"):
				re, err := regexp.Compile(test.Output)
				if err != nil {
					t.Fatalf("test %q cannot compile regexp Output %q: %v", name, test.Output, err)
				}
				if !re.Match([]byte(Output)) {
					hasError = true
					t.Errorf("test %q returned unexpected body, got:\n%s\n\nwant regexp match of:\n%s", name, Output, test.Output)
				}
			case strings.HasPrefix(test.Output, "*") && strings.HasSuffix(test.Output, "*"):
				q := regexp.QuoteMeta(test.Output[1 : len(test.Output)-1])
				q = strings.Replace(q, "\\*", ".*", -1)
				re, err := regexp.Compile(q)
				if err != nil {
					t.Fatalf("test %q cannot compile regexp Output %q with regex %q: %v", name, test.Output, q, err)
				}
				if !re.Match([]byte(Output)) {
					hasError = true
					t.Errorf("test %q returned unexpected body, got:\n%s\n\nwant substring match of:\n%s\n\nregex:\n%s", name, Output, test.Output, q)
				}
			default:
				if diff := cmp.Diff(test.Output, Output, test.CmpOptions); diff != "" {
					hasError = true
					t.Errorf("test %q returned mismatching body (-want +got):\n%s", name, diff)
				}
			}
			if hasError && varInput {
				t.Logf("test %q Input value was: %s", name, InputStr)
			}
		})
	}
}
<|endoftext|>


<commit_before>package main

import (
	"bitbucket.org/anacrolix/dms/dlna/dms"
	"github.com/anacrolix/go-gtk/gtk"
	"github.com/mattn/go-gtk/gdk"
	"io"
	"log"
	"os"
	"runtime"
)

func main() {
	runtime.LockOSThread()
	gtk.Init(&os.Args)
	gdk.ThreadsInit()

	window := gtk.Window(gtk.GTK_WINDOW_TOPLEVEL)
	window.SetTitle("DMS GUI")
	window.Connect("destroy", gtk.MainQuit)

	vbox := gtk.VBox(false, 0)
	window.Add(vbox)

	hbox := gtk.HBox(false, 0)
	vbox.PackStart(hbox, false, true, 0)

	hbox.PackStart(gtk.Label("Shared directory: "), false, true, 0)

	dialog := gtk.FileChooserDialog(
		"Select directory to share", window, gtk.GTK_FILE_CHOOSER_ACTION_SELECT_FOLDER,
		gtk.GTK_STOCK_CANCEL, gtk.GTK_RESPONSE_CANCEL,
		gtk.GTK_STOCK_OK, gtk.GTK_RESPONSE_ACCEPT)

	button := gtk.FileChooserButtonWithDialog(dialog)
	hbox.Add(button)

	logView := gtk.TextView()
	logView.SetEditable(false)
	logView.ModifyFontEasy("monospace")
	logView.SetWrapMode(gtk.GTK_WRAP_WORD_CHAR)
	logViewScroller := gtk.ScrolledWindow(nil, nil)
	logViewScroller.Add(logView)
	logViewScroller.SetPolicy(gtk.GTK_POLICY_AUTOMATIC, gtk.GTK_POLICY_ALWAYS)
	vbox.PackEnd(logViewScroller, true, true, 0)

	getPath := func() string {
		return button.GetFilename()
	}

	appendToLog := func(text string) {
		var endIter gtk.GtkTextIter
		gdk.ThreadsEnter()
		logBuffer := logView.GetBuffer()
		logBuffer.GetEndIter(&endIter)
		logBuffer.Insert(&endIter, text)
		logView.ScrollToIter(&endIter, 0, false, 0, 0)
		gdk.ThreadsLeave()
	}

	window.ShowAll()
	if dialog.Run() != gtk.GTK_RESPONSE_ACCEPT {
		return
	}
	go func() {
		logReader, logWriter := io.Pipe()
		go func() {
			var buf [128]byte
			for {
				n, err := logReader.Read(buf[:])
				appendToLog(string(buf[:n]))
				if err != nil {
					panic(err)
				}
			}
		}()
		dmsLogger := log.New(logWriter, "", 0)
		log.SetOutput(logWriter)
		dmsServer, err := dms.New(getPath(), dmsLogger)
		if err != nil {
			log.Fatalln(err)
		}
		defer dmsServer.Close()
		runtime.LockOSThread()
		gdk.ThreadsEnter()
		button.Connect("selection-changed", func() {
			dmsServer.SetRootPath(getPath())
		})
		gdk.ThreadsLeave()
		runtime.UnlockOSThread()
		dmsServer.Serve()
	}()
	gtk.Main()
	runtime.UnlockOSThread()
}
<commit_msg>mattn/go-gtk has matured such that anacrolix/go-gtk can be deleted<commit_after>package main

import (
	"bitbucket.org/anacrolix/dms/dlna/dms"
	"github.com/mattn/go-gtk/gdk"
	"github.com/mattn/go-gtk/gtk"
	"log"
	"os"
	"runtime"
)

func main() {
	runtime.LockOSThread()
	gtk.Init(&os.Args)
	gdk.ThreadsInit()

	window := gtk.NewWindow(gtk.WINDOW_TOPLEVEL)
	window.SetTitle("DMS GUI")
	window.Connect("destroy", gtk.MainQuit)

	vbox := gtk.NewVBox(false, 0)
	window.Add(vbox)

	hbox := gtk.NewHBox(false, 0)
	vbox.PackStart(hbox, false, true, 0)

	hbox.PackStart(gtk.NewLabel("Shared directory: "), false, true, 0)

	dialog := gtk.NewFileChooserDialog(
		"Select directory to share", window, gtk.FILE_CHOOSER_ACTION_SELECT_FOLDER,
		gtk.STOCK_CANCEL, gtk.RESPONSE_CANCEL,
		gtk.STOCK_OK, gtk.RESPONSE_ACCEPT)

	button := gtk.NewFileChooserButtonWithDialog(dialog)
	hbox.Add(button)

	logView := gtk.NewTextView()
	logView.SetEditable(false)
	logView.ModifyFontEasy("monospace")
	logView.SetWrapMode(gtk.WRAP_WORD_CHAR)
	logViewScroller := gtk.NewScrolledWindow(nil, nil)
	logViewScroller.Add(logView)
	logViewScroller.SetPolicy(gtk.POLICY_AUTOMATIC, gtk.POLICY_ALWAYS)
	vbox.PackEnd(logViewScroller, true, true, 0)

	getPath := func() string {
		return button.GetFilename()
	}

	appendToLog := func(text string) {
		var endIter gtk.GtkTextIter
		gdk.ThreadsEnter()
		logBuffer := logView.GetBuffer()
		logBuffer.GetEndIter(&endIter)
		logBuffer.Insert(&endIter, text)
		logView.ScrollToIter(&endIter, 0, false, 0, 0)
		gdk.ThreadsLeave()
	}

	window.ShowAll()
	if dialog.Run() != gtk.RESPONSE_ACCEPT {
		return
	}
	go func() {
		logReader, logWriter := io.Pipe()
		go func() {
			var buf [128]byte
			for {
				n, err := logReader.Read(buf[:])
				appendToLog(string(buf[:n]))
				if err != nil {
					panic(err)
				}
			}
		}()
		dmsLogger := log.New(logWriter, "", 0)
		log.SetOutput(logWriter)
		dmsServer, err := dms.New(getPath(), dmsLogger)
		if err != nil {
			log.Fatalln(err)
		}
		defer dmsServer.Close()
		runtime.LockOSThread()
		gdk.ThreadsEnter()
		button.Connect("selection-changed", func() {
			dmsServer.SetRootPath(getPath())
		})
		gdk.ThreadsLeave()
		runtime.UnlockOSThread()
		dmsServer.Serve()
	}()
	gtk.Main()
	runtime.UnlockOSThread()
}
<|endoftext|>


<commit_before>package mstate

import (
	"fmt"
	"labix.org/v2/mgo"
	"labix.org/v2/mgo/bson"
)

// Life represents the lifecycle state of the entities
// Relation, Unit, Service and Machine.
type Life int8

const (
	Alive Life = iota
	Dying
	Dead
	nLife
)

var lifeStrings = [nLife]string{
	Alive: "alive",
	Dying: "dying",
	Dead:  "dead",
}

func (l Life) String() string {
	return lifeStrings[l]
}

// Living describes state entities with a lifecycle.
type Living interface {
	Life() Life
	Kill() error
	Die() error
	Refresh() error
}

// ensureLife changes the lifecycle state of the entity with
// the id in the collection.
func ensureLife(id interface{}, coll *mgo.Collection, descr string, life Life) error {
	if life == Alive {
		panic("cannot set life to alive")
	}
	sel := bson.D{
		{"_id", id},
		// $lte is used so that we don't overwrite a previous
		// change we don't know about.
		{"life", bson.D{{"$lte", life}}},
	}
	change := bson.D{{"$set", bson.D{{"life", life}}}}
	err := coll.Update(sel, change)
	if err == mgo.ErrNotFound {
		return nil
	}
	if err != nil {
		return fmt.Errorf("cannot set life to %s for %s %q: %v", life, descr, id, err)
	}
	return nil
}
<commit_msg>mstate: added transaction to life<commit_after>package mstate

import (
	"fmt"
	"labix.org/v2/mgo"
	"labix.org/v2/mgo/bson"
	"labix.org/v2/mgo/txn"
)

// Life represents the lifecycle state of the entities
// Relation, Unit, Service and Machine.
type Life int8

const (
	Alive Life = iota
	Dying
	Dead
	nLife
)

var lifeStrings = [nLife]string{
	Alive: "alive",
	Dying: "dying",
	Dead:  "dead",
}

func (l Life) String() string {
	return lifeStrings[l]
}

// Living describes state entities with a lifecycle.
type Living interface {
	Life() Life
	Kill() error
	Die() error
	Refresh() error
}

// ensureLife changes the lifecycle state of the entity with
// the id in the collection.
func ensureLife(st *State, id interface{}, coll *mgo.Collection, descr string, life Life) error {
	if life == Alive {
		panic("cannot set life to alive")
	}
	sel := bson.D{
		{"_id", id},
		// $lte is used so that we don't overwrite a previous
		// change we don't know about.
		{"life", bson.D{{"$lte", life}}},
	}
	change := bson.D{{"$set", bson.D{{"life", life}}}}
	op := []txn.Operation{{
		Collection: coll.Name,
		DocId:      id,
		Assert:     sel,
		Change:     change,
	}}
	err = st.runner.Run(op, "", nil)
	if err != nil {
		return fmt.Errorf("cannot set life to %s for %s %q: %v", life, descr, id, err)
	}
	return nil
}
<|endoftext|>


<commit_before>// Copyright 2015 trivago GmbH
//
// 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 main

import (
	"github.com/trivago/gollum/core"
	"github.com/trivago/gollum/core/log"
	"github.com/trivago/gollum/shared"
	"log"
	"os"
	"os/signal"
	"sync"
	"syscall"
	"time"
)

const (
	metricMsgSec    = "MessagesPerSec"
	metricMsgSecAvg = "MessagesPerSecAvg"
	metricCons      = "Consumers"
	metricProds     = "Producers"
	metricStreams   = "ManagedStreams"
	metricMessages  = "Messages"
)

type multiplexerState byte

const (
	multiplexerStateConfigure      = multiplexerState(iota)
	multiplexerStateStartProducers = multiplexerState(iota)
	multiplexerStateStartConsumers = multiplexerState(iota)
	multiplexerStateRunning        = multiplexerState(iota)
	multiplexerStateShutdown       = multiplexerState(iota)
	multiplexerStateStopConsumers  = multiplexerState(iota)
	multiplexerStateStopProducers  = multiplexerState(iota)
	multiplexerStateStopped        = multiplexerState(iota)
)

type multiplexer struct {
	consumers      []core.Consumer
	producers      []core.Producer
	streams        map[core.MessageStreamID][]core.Distributor
	consumerWorker *sync.WaitGroup
	producerWorker *sync.WaitGroup
	state          multiplexerState
	signal         chan os.Signal
	profile        bool
}

// Create a new multiplexer based on a given config file.
func newMultiplexer(conf *core.Config, profile bool) multiplexer {
	// Configure the multiplexer, create a byte pool and assign it to the log

	logConsumer := core.LogConsumer{}
	logConsumer.Configure(core.PluginConfig{})

	shared.Metric.New(metricMsgSec)
	shared.Metric.New(metricMsgSecAvg)
	shared.Metric.New(metricCons)
	shared.Metric.New(metricProds)
	shared.Metric.New(metricStreams)
	shared.Metric.New(metricMessages)

	plex := multiplexer{
		consumers:      []core.Consumer{&logConsumer},
		streams:        make(map[core.MessageStreamID][]core.Distributor),
		consumerWorker: new(sync.WaitGroup),
		producerWorker: new(sync.WaitGroup),
		profile:        profile,
		state:          multiplexerStateConfigure,
	}

	// Initialize the plugins based on the config

	for _, config := range conf.Plugins {
		if !config.Enable {
			continue // ### continue, disabled ###
		}

		// Try to instantiate and configure the plugin

		plugin, err := core.NewPlugin(config)
		if err != nil {
			if plugin == nil {
				Log.Error.Panic(err.Error())
			} else {
				Log.Error.Print("Failed to configure ", config.TypeName, ": ", err)
				continue // ### continue ###
			}
		}

		// Register consumer plugins
		if consumer, isConsumer := plugin.(core.Consumer); isConsumer {
			plex.consumers = append(plex.consumers, consumer)
			shared.Metric.Inc(metricCons)

			for i := 1; i < config.Instances; i++ {
				clone, _ := core.NewPlugin(config)
				plex.consumers = append(plex.consumers, clone.(core.Consumer))
				shared.Metric.Inc(metricCons)
			}
		}

		// Register producer plugins
		if producer, isProducer := plugin.(core.Producer); isProducer {
			plex.producers = append(plex.producers, producer)
			shared.Metric.Inc(metricProds)

			for i := 1; i < config.Instances; i++ {
				clone, _ := core.NewPlugin(config)
				plex.producers = append(plex.producers, clone.(core.Producer))
				shared.Metric.Inc(metricCons)
			}
		}

		// Register dsitributor plugins
		if _, isDistributor := plugin.(core.Distributor); isDistributor {
			for _, stream := range config.Stream {
				streamID := core.GetStreamID(stream)

				// New instance per stream
				distributor, _ := core.NewPlugin(config)

				if stream, isMapped := plex.streams[streamID]; isMapped {
					plex.streams[streamID] = append(stream, distributor.(core.Distributor))
				} else {
					plex.streams[streamID] = []core.Distributor{distributor.(core.Distributor)}
				}
			}
		}
	}

	// Analyze all producers and add them to the corresponding distributors
	// Wildcard streams will be handled after this so we have a fully configured
	// map to add to.

	var wildcardProducers []core.Producer

	for _, prod := range plex.producers {
		// Check all streams for each producer
		for _, streamID := range prod.Streams() {
			if streamID == core.WildcardStreamID {
				wildcardProducers = append(wildcardProducers, prod)
			}

			if distList, isMapped := plex.streams[streamID]; isMapped {
				// Add to all distributors for this stream
				for _, dist := range distList {
					dist.AddProducer(prod)
				}
			} else {
				// No distributor for this stream: Create broadcast
				// distributor as a default.
				defaultDistPlugin, _ := shared.RuntimeType.New("distributor.Broadcast")
				defaultDist := defaultDistPlugin.(core.Distributor)
				defaultDist.AddProducer(prod)
				plex.streams[streamID] = []core.Distributor{defaultDist}
			}
		}
	}

	// Append wildcard producers to all streams

	for _, prod := range wildcardProducers {
		for streamID, distList := range plex.streams {
			switch streamID {
			case core.WildcardStreamID:
			case core.DroppedStreamID:
			case core.LogInternalStreamID:
			default:
				for _, dist := range distList {
					dist.AddProducer(prod)
				}
			}
		}
	}

	return plex
}

// Shutdown all consumers and producers in a clean way.
// The internal log is flushed after the consumers have been shut down so that
// consumer related messages are still in the log.
// Producers are flushed after flushing the log, so producer related shutdown
// messages will be posted to stdout
func (plex *multiplexer) shutdown() {
	// Make Ctrl+C possible during shutdown sequence
	if plex.signal != nil {
		signal.Stop(plex.signal)
	}

	Log.Note.Print("Filthy little hobbites. They stole it from us. (shutdown)")
	stateAtShutdown := plex.state

	// Shutdown consumers
	plex.state = multiplexerStateStopConsumers
	if stateAtShutdown >= multiplexerStateStartConsumers {
		for _, consumer := range plex.consumers {
			consumer.Control() <- core.ConsumerControlStop
		}

		// Make sure all remaining messages are flushed BEFORE waiting for all
		// consumers to stop. This is necessary as consumers might be waiting in
		// a push to channel.
		Log.Note.Print("It's the only way. Go in, or go back. (flushing)")

		for _, consumer := range plex.consumers {
		flushing:
			for {
				select {
				case message := <-consumer.Messages():
					plex.distribute(message)
				default:
					break flushing
				}
			}
		}

		plex.consumerWorker.Wait()
	}

	// Make sure remaining warning / errors are written to stderr
	Log.SetWriter(os.Stdout)

	// Shutdown producers
	plex.state = multiplexerStateStopProducers
	if stateAtShutdown >= multiplexerStateStartProducers {
		for _, producer := range plex.producers {
			producer.Control() <- core.ProducerControlStop
		}
		plex.producerWorker.Wait()
	}

	// Done
	plex.state = multiplexerStateStopped
}

func (plex *multiplexer) distribute(msg core.Message) {
	for _, streamID := range msg.Streams {
		msg.CurrentStream = streamID

		distList, isMapped := plex.streams[streamID]
		if !isMapped {
			distList = plex.streams[core.WildcardStreamID]
		}

		for _, distributor := range distList {
			distributor.Distribute(msg)
		}
	}
}

func (plex *multiplexer) handleSignals() {
	plex.signal = make(chan os.Signal, 1)
	signal.Notify(plex.signal, os.Interrupt, syscall.SIGHUP)

	for {
		sig := <-plex.signal
		switch sig {
		case syscall.SIGINT:
			Log.Note.Print("Master betrayed us. Wicked. Tricksy, False. (signal)")
			plex.state = multiplexerStateShutdown
			return // ### return, exit requested ###

		case syscall.SIGHUP:
			for _, consumer := range plex.consumers {
				consumer.Control() <- core.ConsumerControlRoll
			}
			for _, producer := range plex.producers {
				producer.Control() <- core.ProducerControlRoll
			}
		}
	}
}

func (plex *multiplexer) handlePanics() {
	if r := recover(); r != nil {
		log.Println(r)
	}

	plex.shutdown()
}

// Run the multiplexer.
// Fetch messags from the consumers and pass them to all producers.
func (plex multiplexer) run() {
	if len(plex.consumers) == 0 {
		Log.Error.Print("No consumers configured.")
		return // ### return, nothing to do ###
	}

	if len(plex.producers) == 0 {
		Log.Error.Print("No producers configured.")
		return // ### return, nothing to do ###
	}

	defer plex.handlePanics()

	// Launch producers
	plex.state = multiplexerStateStartProducers
	for _, producer := range plex.producers {
		producer := producer
		go func() {
			defer shared.RecoverShutdown()
			producer.Produce(plex.producerWorker)
		}()
	}

	// If there are intenal log listeners switch to stream mode
	if _, enableQueue := plex.streams[core.LogInternalStreamID]; enableQueue {
		Log.SetWriter(plex.consumers[0].(*core.LogConsumer))
	}

	// Launch consumers
	plex.state = multiplexerStateStartConsumers
	for _, consumer := range plex.consumers {
		consumer := consumer
		go func() {
			defer shared.RecoverShutdown()
			consumer.Consume(plex.consumerWorker)
		}()
	}

	// Main loop
	go plex.handleSignals()
	Log.Note.Print("We be nice to them, if they be nice to us. (startup)")

	measure := time.Now()
	messageCount := 0
	plex.state = multiplexerStateRunning

	for plex.state != multiplexerStateShutdown {
		// Go over all consumers in round-robin fashion
		// Don't block here, too as a consumer might not contain new messages

		messageCountBefore := messageCount

		for _, consumer := range plex.consumers {
			select {
			default:
			case message := <-consumer.Messages():
				plex.distribute(message)
				messageCount++
			}
		}

		// Sleep if there is nothing to do

		if messageCount-messageCountBefore == 0 {
			time.Sleep(time.Duration(100) * time.Millisecond)
		}

		// Profiling information

		duration := time.Since(measure)
		if messageCount >= 100000 || duration.Seconds() > 5 {
			// Local values
			value := float64(messageCount) / duration.Seconds()
			shared.Metric.SetF(metricMsgSec, value)
			shared.Metric.AddI(metricMessages, messageCount)

			if plex.profile {
				Log.Note.Printf("Processed %.2f msg/sec", value)
			}

			// Global values
			timeSinceStart := time.Since(shared.ProcessStartTime)
			if totalMessages, err := shared.Metric.Get(metricMessages); err == nil {
				value = float64(totalMessages) / timeSinceStart.Seconds()
				shared.Metric.SetF(metricMsgSecAvg, value)
			}

			// Prepare next run
			measure = time.Now()
			messageCount = 0
		}
	}
}
<commit_msg>removed idle code<commit_after>// Copyright 2015 trivago GmbH
//
// 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 main

import (
	"github.com/trivago/gollum/core"
	"github.com/trivago/gollum/core/log"
	"github.com/trivago/gollum/shared"
	"log"
	"os"
	"os/signal"
	"sync"
	"syscall"
	"time"
)

const (
	metricMsgSec    = "MessagesPerSec"
	metricMsgSecAvg = "MessagesPerSecAvg"
	metricCons      = "Consumers"
	metricProds     = "Producers"
	metricStreams   = "ManagedStreams"
	metricMessages  = "Messages"
)

type multiplexerState byte

const (
	multiplexerStateConfigure      = multiplexerState(iota)
	multiplexerStateStartProducers = multiplexerState(iota)
	multiplexerStateStartConsumers = multiplexerState(iota)
	multiplexerStateRunning        = multiplexerState(iota)
	multiplexerStateShutdown       = multiplexerState(iota)
	multiplexerStateStopConsumers  = multiplexerState(iota)
	multiplexerStateStopProducers  = multiplexerState(iota)
	multiplexerStateStopped        = multiplexerState(iota)
)

type multiplexer struct {
	consumers      []core.Consumer
	producers      []core.Producer
	streams        map[core.MessageStreamID][]core.Distributor
	consumerWorker *sync.WaitGroup
	producerWorker *sync.WaitGroup
	state          multiplexerState
	signal         chan os.Signal
	profile        bool
}

// Create a new multiplexer based on a given config file.
func newMultiplexer(conf *core.Config, profile bool) multiplexer {
	// Configure the multiplexer, create a byte pool and assign it to the log

	logConsumer := core.LogConsumer{}
	logConsumer.Configure(core.PluginConfig{})

	shared.Metric.New(metricMsgSec)
	shared.Metric.New(metricMsgSecAvg)
	shared.Metric.New(metricCons)
	shared.Metric.New(metricProds)
	shared.Metric.New(metricStreams)
	shared.Metric.New(metricMessages)

	plex := multiplexer{
		consumers:      []core.Consumer{&logConsumer},
		streams:        make(map[core.MessageStreamID][]core.Distributor),
		consumerWorker: new(sync.WaitGroup),
		producerWorker: new(sync.WaitGroup),
		profile:        profile,
		state:          multiplexerStateConfigure,
	}

	// Initialize the plugins based on the config

	for _, config := range conf.Plugins {
		if !config.Enable {
			continue // ### continue, disabled ###
		}

		// Try to instantiate and configure the plugin

		plugin, err := core.NewPlugin(config)
		if err != nil {
			if plugin == nil {
				Log.Error.Panic(err.Error())
			} else {
				Log.Error.Print("Failed to configure ", config.TypeName, ": ", err)
				continue // ### continue ###
			}
		}

		// Register consumer plugins
		if consumer, isConsumer := plugin.(core.Consumer); isConsumer {
			plex.consumers = append(plex.consumers, consumer)
			shared.Metric.Inc(metricCons)

			for i := 1; i < config.Instances; i++ {
				clone, _ := core.NewPlugin(config)
				plex.consumers = append(plex.consumers, clone.(core.Consumer))
				shared.Metric.Inc(metricCons)
			}
		}

		// Register producer plugins
		if producer, isProducer := plugin.(core.Producer); isProducer {
			plex.producers = append(plex.producers, producer)
			shared.Metric.Inc(metricProds)

			for i := 1; i < config.Instances; i++ {
				clone, _ := core.NewPlugin(config)
				plex.producers = append(plex.producers, clone.(core.Producer))
				shared.Metric.Inc(metricCons)
			}
		}

		// Register dsitributor plugins
		if _, isDistributor := plugin.(core.Distributor); isDistributor {
			for _, stream := range config.Stream {
				streamID := core.GetStreamID(stream)

				// New instance per stream
				distributor, _ := core.NewPlugin(config)

				if stream, isMapped := plex.streams[streamID]; isMapped {
					plex.streams[streamID] = append(stream, distributor.(core.Distributor))
				} else {
					plex.streams[streamID] = []core.Distributor{distributor.(core.Distributor)}
				}
			}
		}
	}

	// Analyze all producers and add them to the corresponding distributors
	// Wildcard streams will be handled after this so we have a fully configured
	// map to add to.

	var wildcardProducers []core.Producer

	for _, prod := range plex.producers {
		// Check all streams for each producer
		for _, streamID := range prod.Streams() {
			if streamID == core.WildcardStreamID {
				wildcardProducers = append(wildcardProducers, prod)
			}

			if distList, isMapped := plex.streams[streamID]; isMapped {
				// Add to all distributors for this stream
				for _, dist := range distList {
					dist.AddProducer(prod)
				}
			} else {
				// No distributor for this stream: Create broadcast
				// distributor as a default.
				defaultDistPlugin, _ := shared.RuntimeType.New("distributor.Broadcast")
				defaultDist := defaultDistPlugin.(core.Distributor)
				defaultDist.AddProducer(prod)
				plex.streams[streamID] = []core.Distributor{defaultDist}
			}
		}
	}

	// Append wildcard producers to all streams

	for _, prod := range wildcardProducers {
		for streamID, distList := range plex.streams {
			switch streamID {
			case core.WildcardStreamID:
			case core.DroppedStreamID:
			case core.LogInternalStreamID:
			default:
				for _, dist := range distList {
					dist.AddProducer(prod)
				}
			}
		}
	}

	return plex
}

// Shutdown all consumers and producers in a clean way.
// The internal log is flushed after the consumers have been shut down so that
// consumer related messages are still in the log.
// Producers are flushed after flushing the log, so producer related shutdown
// messages will be posted to stdout
func (plex *multiplexer) shutdown() {
	// Make Ctrl+C possible during shutdown sequence
	if plex.signal != nil {
		signal.Stop(plex.signal)
	}

	Log.Note.Print("Filthy little hobbites. They stole it from us. (shutdown)")
	stateAtShutdown := plex.state

	// Shutdown consumers
	plex.state = multiplexerStateStopConsumers
	if stateAtShutdown >= multiplexerStateStartConsumers {
		for _, consumer := range plex.consumers {
			consumer.Control() <- core.ConsumerControlStop
		}

		// Make sure all remaining messages are flushed BEFORE waiting for all
		// consumers to stop. This is necessary as consumers might be waiting in
		// a push to channel.
		Log.Note.Print("It's the only way. Go in, or go back. (flushing)")

		for _, consumer := range plex.consumers {
		flushing:
			for {
				select {
				case message := <-consumer.Messages():
					plex.distribute(message)
				default:
					break flushing
				}
			}
		}

		plex.consumerWorker.Wait()
	}

	// Make sure remaining warning / errors are written to stderr
	Log.SetWriter(os.Stdout)

	// Shutdown producers
	plex.state = multiplexerStateStopProducers
	if stateAtShutdown >= multiplexerStateStartProducers {
		for _, producer := range plex.producers {
			producer.Control() <- core.ProducerControlStop
		}
		plex.producerWorker.Wait()
	}

	// Done
	plex.state = multiplexerStateStopped
}

func (plex *multiplexer) distribute(msg core.Message) {
	for _, streamID := range msg.Streams {
		msg.CurrentStream = streamID

		distList, isMapped := plex.streams[streamID]
		if !isMapped {
			distList = plex.streams[core.WildcardStreamID]
		}

		for _, distributor := range distList {
			distributor.Distribute(msg)
		}
	}
}

func (plex *multiplexer) handleSignals() {
	plex.signal = make(chan os.Signal, 1)
	signal.Notify(plex.signal, os.Interrupt, syscall.SIGHUP)

	for {
		sig := <-plex.signal
		switch sig {
		case syscall.SIGINT:
			Log.Note.Print("Master betrayed us. Wicked. Tricksy, False. (signal)")
			plex.state = multiplexerStateShutdown
			return // ### return, exit requested ###

		case syscall.SIGHUP:
			for _, consumer := range plex.consumers {
				consumer.Control() <- core.ConsumerControlRoll
			}
			for _, producer := range plex.producers {
				producer.Control() <- core.ProducerControlRoll
			}
		}
	}
}

func (plex *multiplexer) handlePanics() {
	if r := recover(); r != nil {
		log.Println(r)
	}

	plex.shutdown()
}

// Run the multiplexer.
// Fetch messags from the consumers and pass them to all producers.
func (plex multiplexer) run() {
	if len(plex.consumers) == 0 {
		Log.Error.Print("No consumers configured.")
		return // ### return, nothing to do ###
	}

	if len(plex.producers) == 0 {
		Log.Error.Print("No producers configured.")
		return // ### return, nothing to do ###
	}

	defer plex.handlePanics()

	// Launch producers
	plex.state = multiplexerStateStartProducers
	for _, producer := range plex.producers {
		producer := producer
		go func() {
			defer shared.RecoverShutdown()
			producer.Produce(plex.producerWorker)
		}()
	}

	// If there are intenal log listeners switch to stream mode
	if _, enableQueue := plex.streams[core.LogInternalStreamID]; enableQueue {
		Log.SetWriter(plex.consumers[0].(*core.LogConsumer))
	}

	// Launch consumers
	plex.state = multiplexerStateStartConsumers
	for _, consumer := range plex.consumers {
		consumer := consumer
		go func() {
			defer shared.RecoverShutdown()
			consumer.Consume(plex.consumerWorker)
		}()
	}

	// Main loop
	go plex.handleSignals()
	Log.Note.Print("We be nice to them, if they be nice to us. (startup)")

	measure := time.Now()
	messageCount := 0
	plex.state = multiplexerStateRunning

	for plex.state != multiplexerStateShutdown {
		// Go over all consumers in round-robin fashion
		// Don't block here, too as a consumer might not contain new messages

		for _, consumer := range plex.consumers {
			select {
			default:
			case message := <-consumer.Messages():
				plex.distribute(message)
				messageCount++
			}
		}

		// Profiling information and idle state

		duration := time.Since(measure)
		if messageCount >= 100000 || duration.Seconds() > 5 {
			// Local values
			value := float64(messageCount) / duration.Seconds()
			shared.Metric.SetF(metricMsgSec, value)
			shared.Metric.AddI(metricMessages, messageCount)

			if plex.profile {
				Log.Note.Printf("Processed %.2f msg/sec", value)
			}

			// Global values
			timeSinceStart := time.Since(shared.ProcessStartTime)
			if totalMessages, err := shared.Metric.Get(metricMessages); err == nil {
				value = float64(totalMessages) / timeSinceStart.Seconds()
				shared.Metric.SetF(metricMsgSecAvg, value)
			}

			// Prepare next run
			measure = time.Now()
			messageCount = 0
		}
	}
}
<|endoftext|>


<commit_before>// Copyright 2013 Canonical Ltd.
// Licensed under the AGPLv3, see LICENCE file for details.

package environs_test

import (
	. "launchpad.net/gocheck"
	"launchpad.net/juju-core/environs"
)

type EmptyStorageSuite struct{}

var _ = Suite(&EmptyStorageSuite{})

func (s *EmptyStorageSuite) TestGet(c *C) {
	f, err := environs.EmptyStorage.Get("anything")
	c.Assert(f, IsNil)
	c.Assert(err, ErrorMatches, `file "anything" not found`)
}

func (s *EmptyStorageSuite) TestURL(c *C) {
	url, err := environs.EmptyStorage.URL("anything")
	c.Assert(url, Equals, "")
	c.Assert(err, ErrorMatches, `file "anything" not found`)
}

func (s *EmptyStorageSuite) TestList(c *C) {
	names, err := environs.EmptyStorage.List("anything")
	c.Assert(names, IsNil)
	c.Assert(err, IsNil)
}
<commit_msg>Start a VerifyStorage test.<commit_after>// Copyright 2013 Canonical Ltd.
// Licensed under the AGPLv3, see LICENCE file for details.

package environs_test

import (
	"fmt"
	"io"
	. "launchpad.net/gocheck"
	"launchpad.net/juju-core/environs"
)

type EmptyStorageSuite struct{}

var _ = Suite(&EmptyStorageSuite{})

func (s *EmptyStorageSuite) TestGet(c *C) {
	f, err := environs.EmptyStorage.Get("anything")
	c.Assert(f, IsNil)
	c.Assert(err, ErrorMatches, `file "anything" not found`)
}

func (s *EmptyStorageSuite) TestURL(c *C) {
	url, err := environs.EmptyStorage.URL("anything")
	c.Assert(url, Equals, "")
	c.Assert(err, ErrorMatches, `file "anything" not found`)
}

func (s *EmptyStorageSuite) TestList(c *C) {
	names, err := environs.EmptyStorage.List("anything")
	c.Assert(names, IsNil)
	c.Assert(err, IsNil)
}

type MockStorage struct {
	StorageRequests []string
}

func (ms *MockStorage) Put(filename string, reader io.Reader, length int64) error {
	var content []byte
	_, _ = reader.Read(content)
	log_message := fmt.Sprintf(
		"Put('%s', '%s', %d)", filename, content, length)
	ms.StorageRequests = append(ms.StorageRequests, log_message)
	return nil
}

func (ms *MockStorage) Get(name string) (io.ReadCloser, error) {
	return nil, nil
	//return nil, &NotFoundError{fmt.Errorf("file %q not found", name)}
}

func (ms *MockStorage) List(prefix string) ([]string, error) {
	return nil, nil
}

func (ms *MockStorage) Remove(file string) error {
	return nil
}

func (ms *MockStorage) URL(name string) (string, error) {
	return "", fmt.Errorf("file %q not found", name)
}

type VerifyStorageSuite struct{}

var _ = Suite(&VerifyStorageSuite{})

func (s *VerifyStorageSuite) TestVerifyStorage(c *C) {
	storage := MockStorage{}
	error := environs.VerifyStorage(&storage)
	c.Assert(error, IsNil)
	c.Assert(storage.StorageRequests[0], Equals,
		"Put('bootstrap-verify', '%s', 28)")
}
<|endoftext|>