<commit_before>package minion

import (
	"bytes"
	"encoding/json"
	"log"
	"os"
	"os/exec"
	"os/signal"
	"path/filepath"
	"strconv"
	"strings"
	"time"

	"github.com/dnaeon/gru/classifier"
	"github.com/dnaeon/gru/task"
	"github.com/dnaeon/gru/utils"

	"code.google.com/p/go-uuid/uuid"
	etcdclient "github.com/coreos/etcd/client"
	"golang.org/x/net/context"
)

// Minions keyspace in etcd
const EtcdMinionSpace = "/gru/minion"

// Etcd Minion
type etcdMinion struct {
	// Name of this minion
	name string

	// Minion root node in etcd
	rootDir string

	// Minion queue node in etcd
	queueDir string

	// Log directory to keep previously executed tasks
	logDir string

	// Root node for classifiers in etcd
	classifierDir string

	// Minion unique identifier
	id uuid.UUID

	// KeysAPI client to etcd
	kapi etcdclient.KeysAPI
}

// Creates a new etcd minion
func NewEtcdMinion(name string, cfg etcdclient.Config) Minion {
	c, err := etcdclient.New(cfg)
	if err != nil {
		log.Fatal(err)
	}

	kapi := etcdclient.NewKeysAPI(c)
	id := utils.GenerateUUID(name)
	rootDir := filepath.Join(EtcdMinionSpace, id.String())
	queueDir := filepath.Join(rootDir, "queue")
	classifierDir := filepath.Join(rootDir, "classifier")
	logDir := filepath.Join(rootDir, "log")

	m := &etcdMinion{
		name:          name,
		rootDir:       rootDir,
		queueDir:      queueDir,
		classifierDir: classifierDir,
		logDir:        logDir,
		id:            id,
		kapi:          kapi,
	}

	return m
}

// Set the human-readable name of the minion in etcd
func (m *etcdMinion) setName() error {
	nameKey := filepath.Join(m.rootDir, "name")
	opts := &etcdclient.SetOptions{
		PrevExist: etcdclient.PrevIgnore,
	}

	_, err := m.kapi.Set(context.Background(), nameKey, m.Name(), opts)

	return err
}

// Set the time the minion was last seen in seconds since the Epoch
func (m *etcdMinion) setLastseen(s int64) error {
	lastseenKey := filepath.Join(m.rootDir, "lastseen")
	lastseenValue := strconv.FormatInt(s, 10)
	opts := &etcdclient.SetOptions{
		PrevExist: etcdclient.PrevIgnore,
	}

	_, err := m.kapi.Set(context.Background(), lastseenKey, lastseenValue, opts)

	return err
}

// Checks for any pending tasks in queue
func (m *etcdMinion) checkQueue(c chan<- *task.Task) error {
	opts := &etcdclient.GetOptions{
		Recursive: true,
		Sort:      true,
	}

	// Get backlog tasks if any
	resp, err := m.kapi.Get(context.Background(), m.queueDir, opts)
	if err != nil {
		return nil
	}

	backlog := resp.Node.Nodes
	if len(backlog) == 0 {
		// No backlog tasks found
		return nil
	}

	log.Printf("Found %d tasks in queue", len(backlog))
	for _, node := range backlog {
		task, err := EtcdUnmarshalTask(node)
		m.kapi.Delete(context.Background(), node.Key, nil)

		if err != nil {
			continue
		}

		c <- task
	}

	return nil
}

// Runs periodic jobs such as refreshing classifiers and
// updating the lastseen time
func (m *etcdMinion) periodicRunner(ticker *time.Ticker) error {
	for {
		// Update minion classifiers
		m.Classify()

		// Update lastseen time
		now := time.Now().Unix()
		err := m.setLastseen(now)
		if err != nil {
			log.Printf("Failed to update lastseen time: %s\n", err)
		}

		<-ticker.C
	}

	return nil
}

// Processes new tasks
func (m *etcdMinion) processTask(t *task.Task) error {
	var buf bytes.Buffer

	// Update state of task that we are now processing it
	t.State = task.TaskStateProcessing
	m.SaveTaskResult(t)

	cmd := exec.Command(t.Command, t.Args...)
	cmd.Stdout = &buf
	cmd.Stderr = &buf

	log.Printf("Processing task %s\n", t.TaskID)

	cmdError := cmd.Run()
	t.TimeProcessed = time.Now().Unix()
	t.Result = buf.String()

	if cmdError != nil {
		log.Printf("Failed to process task %s\n", t.TaskID)
		t.Error = cmdError.Error()
		t.State = task.TaskStateFailed
	} else {
		log.Printf("Finished processing task %s\n", t.TaskID)
		t.State = task.TaskStateSuccess
	}

	m.SaveTaskResult(t)

	return cmdError
}

// Saves the task result
func (m *etcdMinion) SaveTaskResult(t *task.Task) error {
	// Task key in etcd
	taskKey := filepath.Join(m.logDir, t.TaskID.String())

	// Serialize task to JSON
	data, err := json.Marshal(t)
	if err != nil {
		log.Printf("Failed to serialize task %s: %s\n", t.TaskID, err)
		return err
	}

	// Save the task result in etcd
	opts := &etcdclient.SetOptions{
		PrevExist: etcdclient.PrevIgnore,
	}

	_, err = m.kapi.Set(context.Background(), taskKey, string(data), opts)
	if err != nil {
		log.Printf("Failed to save task %s: %s\n", t.TaskID, err)
	}

	return err
}

// Unmarshals task from etcd
func EtcdUnmarshalTask(node *etcdclient.Node) (*task.Task, error) {
	task := new(task.Task)
	err := json.Unmarshal([]byte(node.Value), &task)

	return task, err
}

// Returns the minion unique identifier
func (m *etcdMinion) ID() uuid.UUID {
	return m.id
}

// Returns the assigned name of the minion
func (m *etcdMinion) Name() string {
	return m.name
}

// Classifies the minion
func (m *etcdMinion) Classify() error {
	// Classifiers in etcd expire after an hour
	opts := &etcdclient.SetOptions{
		PrevExist: etcdclient.PrevIgnore,
		TTL:       time.Hour,
	}

	// Update classifiers
	for key, _ := range classifier.Registry {
		klassifier, err := classifier.Get(key)

		if err != nil {
			continue
		}

		// Serialize classifier to JSON and save it in etcd
		data, err := json.Marshal(klassifier)
		if err != nil {
			log.Printf("Failed to serialize classifier: %s\n", key)
			continue
		}

		// Classifier key in etcd
		klassifierKey := filepath.Join(m.classifierDir, key)
		_, err = m.kapi.Set(context.Background(), klassifierKey, string(data), opts)

		if err != nil {
			log.Printf("Failed to set classifier %s: %s\n", key, err)
		}
	}

	return nil
}

// Monitors etcd for new tasks
func (m *etcdMinion) TaskListener(c chan<- *task.Task) error {
	watcherOpts := &etcdclient.WatcherOptions{
		Recursive: true,
	}
	watcher := m.kapi.Watcher(m.queueDir, watcherOpts)

	for {
		resp, err := watcher.Next(context.Background())
		if err != nil {
			log.Printf("Failed to read task: %s\n", err)
			continue
		}

		// Ignore "delete" events when removing a task from the queue
		action := strings.ToLower(resp.Action)
		if strings.EqualFold(action, "delete") {
			continue
		}

		// Remove task from the queue
		t, err := EtcdUnmarshalTask(resp.Node)
		m.kapi.Delete(context.Background(), resp.Node.Key, nil)

		if err != nil {
			log.Printf("Invalid task %s: %s\n", resp.Node.Key, err)
			continue
		}

		// Update task state and save it
		t.State = task.TaskStateQueued
		t.TimeReceived = time.Now().Unix()
		m.SaveTaskResult(t)

		log.Printf("Received task %s\n", t.TaskID)

		c <- t
	}

	return nil
}

// Processes new tasks
func (m *etcdMinion) TaskRunner(c <-chan *task.Task) error {
	for {
		task := <-c

		if task.IsConcurrent {
			go m.processTask(task)
		} else {
			m.processTask(task)
		}
	}

	return nil
}

// Main entry point of the minion
func (m *etcdMinion) Serve() error {
	// Channel on which we send the quit signal
	quit := make(chan os.Signal, 1)
	signal.Notify(quit, os.Interrupt)

	m.setName()
	log.Printf("Minion %s is ready to serve", m.ID())

	// Run periodic tasks every fifteen minutes
	ticker := time.NewTicker(time.Minute * 15)
	go m.periodicRunner(ticker)

	// Check for pending tasks in queue first
	tasks := make(chan *task.Task)
	go m.TaskRunner(tasks)
	m.checkQueue(tasks)

	go m.TaskListener(tasks)

	// Block until a stop signal is received
	s := <-quit
	log.Printf("Received %s signal, shutting down", s)
	close(quit)
	close(tasks)

	return nil
}
<commit_msg>Comments update<commit_after>package minion

import (
	"bytes"
	"encoding/json"
	"log"
	"os"
	"os/exec"
	"os/signal"
	"path/filepath"
	"strconv"
	"strings"
	"time"

	"github.com/dnaeon/gru/classifier"
	"github.com/dnaeon/gru/task"
	"github.com/dnaeon/gru/utils"

	"code.google.com/p/go-uuid/uuid"
	etcdclient "github.com/coreos/etcd/client"
	"golang.org/x/net/context"
)

// Minions keyspace in etcd
const EtcdMinionSpace = "/gru/minion"

// Etcd Minion
type etcdMinion struct {
	// Name of this minion
	name string

	// Minion root directory in etcd
	rootDir string

	// Minion queue directory in etcd
	queueDir string

	// Log directory to keep previously executed tasks
	logDir string

	// Classifier directory in etcd
	classifierDir string

	// Minion unique identifier
	id uuid.UUID

	// KeysAPI client to etcd
	kapi etcdclient.KeysAPI
}

// Creates a new etcd minion
func NewEtcdMinion(name string, cfg etcdclient.Config) Minion {
	c, err := etcdclient.New(cfg)
	if err != nil {
		log.Fatal(err)
	}

	kapi := etcdclient.NewKeysAPI(c)
	id := utils.GenerateUUID(name)
	rootDir := filepath.Join(EtcdMinionSpace, id.String())
	queueDir := filepath.Join(rootDir, "queue")
	classifierDir := filepath.Join(rootDir, "classifier")
	logDir := filepath.Join(rootDir, "log")

	m := &etcdMinion{
		name:          name,
		rootDir:       rootDir,
		queueDir:      queueDir,
		classifierDir: classifierDir,
		logDir:        logDir,
		id:            id,
		kapi:          kapi,
	}

	return m
}

// Set the human-readable name of the minion in etcd
func (m *etcdMinion) setName() error {
	nameKey := filepath.Join(m.rootDir, "name")
	opts := &etcdclient.SetOptions{
		PrevExist: etcdclient.PrevIgnore,
	}

	_, err := m.kapi.Set(context.Background(), nameKey, m.Name(), opts)

	return err
}

// Set the time the minion was last seen in seconds since the Epoch
func (m *etcdMinion) setLastseen(s int64) error {
	lastseenKey := filepath.Join(m.rootDir, "lastseen")
	lastseenValue := strconv.FormatInt(s, 10)
	opts := &etcdclient.SetOptions{
		PrevExist: etcdclient.PrevIgnore,
	}

	_, err := m.kapi.Set(context.Background(), lastseenKey, lastseenValue, opts)

	return err
}

// Checks for any pending tasks in queue
func (m *etcdMinion) checkQueue(c chan<- *task.Task) error {
	opts := &etcdclient.GetOptions{
		Recursive: true,
		Sort:      true,
	}

	// Get backlog tasks if any
	resp, err := m.kapi.Get(context.Background(), m.queueDir, opts)
	if err != nil {
		return nil
	}

	backlog := resp.Node.Nodes
	if len(backlog) == 0 {
		// No backlog tasks found
		return nil
	}

	log.Printf("Found %d tasks in queue", len(backlog))
	for _, node := range backlog {
		task, err := EtcdUnmarshalTask(node)
		m.kapi.Delete(context.Background(), node.Key, nil)

		if err != nil {
			continue
		}

		c <- task
	}

	return nil
}

// Runs periodic jobs such as refreshing classifiers and
// updating the lastseen time
func (m *etcdMinion) periodicRunner(ticker *time.Ticker) error {
	for {
		// Update minion classifiers
		m.Classify()

		// Update lastseen time
		now := time.Now().Unix()
		err := m.setLastseen(now)
		if err != nil {
			log.Printf("Failed to update lastseen time: %s\n", err)
		}

		<-ticker.C
	}

	return nil
}

// Processes new tasks
func (m *etcdMinion) processTask(t *task.Task) error {
	var buf bytes.Buffer

	// Update state of task that we are now processing it
	t.State = task.TaskStateProcessing
	m.SaveTaskResult(t)

	cmd := exec.Command(t.Command, t.Args...)
	cmd.Stdout = &buf
	cmd.Stderr = &buf

	log.Printf("Processing task %s\n", t.TaskID)

	cmdError := cmd.Run()
	t.TimeProcessed = time.Now().Unix()
	t.Result = buf.String()

	if cmdError != nil {
		log.Printf("Failed to process task %s\n", t.TaskID)
		t.Error = cmdError.Error()
		t.State = task.TaskStateFailed
	} else {
		log.Printf("Finished processing task %s\n", t.TaskID)
		t.State = task.TaskStateSuccess
	}

	m.SaveTaskResult(t)

	return cmdError
}

// Saves the task result
func (m *etcdMinion) SaveTaskResult(t *task.Task) error {
	// Task key in etcd
	taskKey := filepath.Join(m.logDir, t.TaskID.String())

	// Serialize task to JSON
	data, err := json.Marshal(t)
	if err != nil {
		log.Printf("Failed to serialize task %s: %s\n", t.TaskID, err)
		return err
	}

	// Save the task result in etcd
	opts := &etcdclient.SetOptions{
		PrevExist: etcdclient.PrevIgnore,
	}

	_, err = m.kapi.Set(context.Background(), taskKey, string(data), opts)
	if err != nil {
		log.Printf("Failed to save task %s: %s\n", t.TaskID, err)
	}

	return err
}

// Unmarshals task from etcd
func EtcdUnmarshalTask(node *etcdclient.Node) (*task.Task, error) {
	task := new(task.Task)
	err := json.Unmarshal([]byte(node.Value), &task)

	return task, err
}

// Returns the minion unique identifier
func (m *etcdMinion) ID() uuid.UUID {
	return m.id
}

// Returns the assigned name of the minion
func (m *etcdMinion) Name() string {
	return m.name
}

// Classifies the minion
func (m *etcdMinion) Classify() error {
	// Classifiers in etcd expire after an hour
	opts := &etcdclient.SetOptions{
		PrevExist: etcdclient.PrevIgnore,
		TTL:       time.Hour,
	}

	// Update classifiers
	for key, _ := range classifier.Registry {
		klassifier, err := classifier.Get(key)

		if err != nil {
			continue
		}

		// Serialize classifier to JSON and save it in etcd
		data, err := json.Marshal(klassifier)
		if err != nil {
			log.Printf("Failed to serialize classifier: %s\n", key)
			continue
		}

		// Classifier key in etcd
		klassifierKey := filepath.Join(m.classifierDir, key)
		_, err = m.kapi.Set(context.Background(), klassifierKey, string(data), opts)

		if err != nil {
			log.Printf("Failed to set classifier %s: %s\n", key, err)
		}
	}

	return nil
}

// Monitors etcd for new tasks
func (m *etcdMinion) TaskListener(c chan<- *task.Task) error {
	watcherOpts := &etcdclient.WatcherOptions{
		Recursive: true,
	}
	watcher := m.kapi.Watcher(m.queueDir, watcherOpts)

	for {
		resp, err := watcher.Next(context.Background())
		if err != nil {
			log.Printf("Failed to read task: %s\n", err)
			continue
		}

		// Ignore "delete" events when removing a task from the queue
		action := strings.ToLower(resp.Action)
		if strings.EqualFold(action, "delete") {
			continue
		}

		// Remove task from the queue
		t, err := EtcdUnmarshalTask(resp.Node)
		m.kapi.Delete(context.Background(), resp.Node.Key, nil)

		if err != nil {
			log.Printf("Invalid task %s: %s\n", resp.Node.Key, err)
			continue
		}

		// Update task state and save it
		t.State = task.TaskStateQueued
		t.TimeReceived = time.Now().Unix()
		m.SaveTaskResult(t)

		log.Printf("Received task %s\n", t.TaskID)

		c <- t
	}

	return nil
}

// Processes new tasks
func (m *etcdMinion) TaskRunner(c <-chan *task.Task) error {
	for {
		task := <-c

		if task.IsConcurrent {
			go m.processTask(task)
		} else {
			m.processTask(task)
		}
	}

	return nil
}

// Main entry point of the minion
func (m *etcdMinion) Serve() error {
	// Channel on which we send the quit signal
	quit := make(chan os.Signal, 1)
	signal.Notify(quit, os.Interrupt)

	m.setName()
	log.Printf("Minion %s is ready to serve", m.ID())

	// Run periodic tasks every fifteen minutes
	ticker := time.NewTicker(time.Minute * 15)
	go m.periodicRunner(ticker)

	// Check for pending tasks in queue first
	tasks := make(chan *task.Task)
	go m.TaskRunner(tasks)
	m.checkQueue(tasks)

	go m.TaskListener(tasks)

	// Block until a stop signal is received
	s := <-quit
	log.Printf("Received %s signal, shutting down", s)
	close(quit)
	close(tasks)

	return nil
}
<|endoftext|>


<commit_before>package minion

import (
	"bytes"
	"encoding/json"
	"log"
	"os/exec"
	"path/filepath"
	"strconv"
	"strings"
	"time"

	"github.com/dnaeon/gru/classifier"
	"github.com/dnaeon/gru/task"
	"github.com/dnaeon/gru/utils"

	"code.google.com/p/go-uuid/uuid"
	etcdclient "github.com/coreos/etcd/client"
	"golang.org/x/net/context"
)

// Minions keyspace in etcd
const EtcdMinionSpace = "/gru/minion"

// Etcd Minion
type etcdMinion struct {
	// Name of the minion
	name string

	// Minion root directory in etcd
	rootDir string

	// Minion queue directory in etcd
	queueDir string

	// Log directory of previously executed tasks
	logDir string

	// Classifier directory in etcd
	classifierDir string

	// Minion unique identifier
	id uuid.UUID

	// KeysAPI client to etcd
	kapi etcdclient.KeysAPI

	// Channel over which tasks are sent for processing
	taskQueue chan *task.Task
}

// Creates a new etcd minion
func NewEtcdMinion(name string, cfg etcdclient.Config) Minion {
	c, err := etcdclient.New(cfg)
	if err != nil {
		log.Fatal(err)
	}

	kapi := etcdclient.NewKeysAPI(c)
	id := utils.GenerateUUID(name)
	rootDir := filepath.Join(EtcdMinionSpace, id.String())
	queueDir := filepath.Join(rootDir, "queue")
	classifierDir := filepath.Join(rootDir, "classifier")
	logDir := filepath.Join(rootDir, "log")
	taskQueue := make(chan *task.Task)

	m := &etcdMinion{
		name:          name,
		rootDir:       rootDir,
		queueDir:      queueDir,
		classifierDir: classifierDir,
		logDir:        logDir,
		id:            id,
		kapi:          kapi,
		taskQueue:     taskQueue,
	}

	return m
}

// Set the human-readable name of the minion in etcd
func (m *etcdMinion) setName() error {
	nameKey := filepath.Join(m.rootDir, "name")
	opts := &etcdclient.SetOptions{
		PrevExist: etcdclient.PrevIgnore,
	}

	_, err := m.kapi.Set(context.Background(), nameKey, m.Name(), opts)
	if err != nil {
		log.Printf("Failed to set name of minion: %s\n", err)
	}

	return err
}

// Set the time the minion was last seen in seconds since the Epoch
func (m *etcdMinion) setLastseen(s int64) error {
	lastseenKey := filepath.Join(m.rootDir, "lastseen")
	lastseenValue := strconv.FormatInt(s, 10)
	opts := &etcdclient.SetOptions{
		PrevExist: etcdclient.PrevIgnore,
	}

	_, err := m.kapi.Set(context.Background(), lastseenKey, lastseenValue, opts)
	if err != nil {
		log.Printf("Failed to set lastseen time: %s\n", err)
	}

	return err
}

// Checks for any pending tasks and sends them
// for processing if there are any
func (m *etcdMinion) checkQueue() error {
	opts := &etcdclient.GetOptions{
		Recursive: true,
		Sort:      true,
	}

	// Get backlog tasks if any
	// If the directory key in etcd is missing that is okay, since
	// it means there are no pending tasks for processing
	resp, err := m.kapi.Get(context.Background(), m.queueDir, opts)
	if err != nil {
		if eerr, ok := err.(etcdclient.Error); !ok || eerr.Code != etcdclient.ErrorCodeKeyNotFound {
			return err
		}
	}

	backlog := resp.Node.Nodes
	if len(backlog) == 0 {
		// No backlog tasks found
		return nil
	}

	log.Printf("Found %d pending tasks in queue", len(backlog))
	for _, node := range backlog {
		t, err := EtcdUnmarshalTask(node)
		m.kapi.Delete(context.Background(), node.Key, nil)

		if err != nil {
			continue
		}

		m.taskQueue <- t
	}

	return nil
}

// Runs periodic jobs such as refreshing classifiers and
// updating the lastseen time every fifteen minutes
func (m *etcdMinion) periodicRunner() {
	schedule := time.Minute * 15
	ticker := time.NewTicker(schedule)
	log.Printf("Periodic scheduler set to run every %s\n", schedule)

	for now := range ticker.C {
		// Run any periodic jobs
		m.Classify()
		m.setLastseen(now.Unix())
	}
}

// Processes new tasks
func (m *etcdMinion) processTask(t *task.Task) error {
	var buf bytes.Buffer

	// Update state of task to indicate that we are now processing it
	t.State = task.TaskStateProcessing
	m.SaveTaskResult(t)

	cmd := exec.Command(t.Command, t.Args...)
	cmd.Stdout = &buf
	cmd.Stderr = &buf

	log.Printf("Processing task %s\n", t.TaskID)

	cmdError := cmd.Run()
	t.TimeProcessed = time.Now().Unix()
	t.Result = buf.String()

	if cmdError != nil {
		log.Printf("Failed to process task %s\n", t.TaskID)
		t.Error = cmdError.Error()
		t.State = task.TaskStateFailed
	} else {
		log.Printf("Finished processing task %s\n", t.TaskID)
		t.State = task.TaskStateSuccess
	}

	m.SaveTaskResult(t)

	return cmdError
}

// Saves a task in etcd
func (m *etcdMinion) SaveTaskResult(t *task.Task) error {
	taskKey := filepath.Join(m.logDir, t.TaskID.String())

	data, err := json.Marshal(t)
	if err != nil {
		log.Printf("Failed to serialize task %s: %s\n", t.TaskID, err)
		return err
	}

	opts := &etcdclient.SetOptions{
		PrevExist: etcdclient.PrevIgnore,
	}

	_, err = m.kapi.Set(context.Background(), taskKey, string(data), opts)
	if err != nil {
		log.Printf("Failed to save task %s: %s\n", t.TaskID, err)
	}

	return err
}

// Unmarshals task from etcd
func EtcdUnmarshalTask(node *etcdclient.Node) (*task.Task, error) {
	task := new(task.Task)
	err := json.Unmarshal([]byte(node.Value), &task)

	return task, err
}

// Returns the minion unique identifier
func (m *etcdMinion) ID() uuid.UUID {
	return m.id
}

// Returns the assigned name of the minion
func (m *etcdMinion) Name() string {
	return m.name
}

// Classifies the minion
func (m *etcdMinion) Classify() error {
	// Classifiers in etcd expire after an hour
	opts := &etcdclient.SetOptions{
		PrevExist: etcdclient.PrevIgnore,
		TTL:       time.Hour,
	}

	// Set/update classifiers in etcd
	for key, _ := range classifier.Registry {
		klassifier, err := classifier.Get(key)

		if err != nil {
			log.Printf("Failed to get classifier %s: %s\n", key, err)
			continue
		}

		// Serialize classifier to JSON and save it in etcd
		data, err := json.Marshal(klassifier)
		if err != nil {
			log.Printf("Failed to serialize classifier: %s\n", key)
			continue
		}

		// Classifier key in etcd
		klassifierKey := filepath.Join(m.classifierDir, key)
		_, err = m.kapi.Set(context.Background(), klassifierKey, string(data), opts)

		if err != nil {
			log.Printf("Failed to set classifier %s: %s\n", key, err)
		}
	}

	return nil
}

// Monitors etcd for new tasks
func (m *etcdMinion) TaskListener(c chan<- *task.Task) error {
	log.Printf("Task listener is watching %s\n", m.queueDir)

	watcherOpts := &etcdclient.WatcherOptions{
		Recursive: true,
	}
	watcher := m.kapi.Watcher(m.queueDir, watcherOpts)

	for {
		resp, err := watcher.Next(context.Background())
		if err != nil {
			log.Printf("Failed to receive task: %s\n", err)
			continue
		}

		// Ignore "delete" events when removing a task from the queue
		action := strings.ToLower(resp.Action)
		if strings.EqualFold(action, "delete") {
			continue
		}

		// Unmarshal and remove task from the queue
		t, err := EtcdUnmarshalTask(resp.Node)
		m.kapi.Delete(context.Background(), resp.Node.Key, nil)

		if err != nil {
			log.Printf("Received invalid task %s: %s\n", resp.Node.Key, err)
			continue
		}

		// Update task state and send it for processing
		log.Printf("Received task %s\n", t.TaskID)
		t.State = task.TaskStateQueued
		t.TimeReceived = time.Now().Unix()
		m.SaveTaskResult(t)
		c <- t
	}

	return nil
}

// Processes new tasks
func (m *etcdMinion) TaskRunner(c <-chan *task.Task) error {
	log.Println("Starting task runner")

	for t := range c {
		if t.IsConcurrent {
			go m.processTask(t)
		} else {
			m.processTask(t)
		}
	}

	return nil
}

// Main entry point of the minion
func (m *etcdMinion) Serve() error {
	err := m.setName()
	if err != nil {
		return err
	}

	// Start minion services
	go m.periodicRunner()
	go m.checkQueue()
	go m.TaskRunner(m.taskQueue)
	go m.TaskListener(m.taskQueue)

	log.Printf("Minion %s is ready to serve", m.ID())

	return nil
}

// Stops the minion and performs any cleanup tasks
func (m *etcdMinion) Stop() error {
	log.Println("Minion is shutting down")

	close(m.taskQueue)

	return nil
}
<commit_msg>Fix error checking when queue directory is missing<commit_after>package minion

import (
	"bytes"
	"encoding/json"
	"log"
	"os/exec"
	"path/filepath"
	"strconv"
	"strings"
	"time"

	"github.com/dnaeon/gru/classifier"
	"github.com/dnaeon/gru/task"
	"github.com/dnaeon/gru/utils"

	"code.google.com/p/go-uuid/uuid"
	etcdclient "github.com/coreos/etcd/client"
	"golang.org/x/net/context"
)

// Minions keyspace in etcd
const EtcdMinionSpace = "/gru/minion"

// Etcd Minion
type etcdMinion struct {
	// Name of the minion
	name string

	// Minion root directory in etcd
	rootDir string

	// Minion queue directory in etcd
	queueDir string

	// Log directory of previously executed tasks
	logDir string

	// Classifier directory in etcd
	classifierDir string

	// Minion unique identifier
	id uuid.UUID

	// KeysAPI client to etcd
	kapi etcdclient.KeysAPI

	// Channel over which tasks are sent for processing
	taskQueue chan *task.Task
}

// Creates a new etcd minion
func NewEtcdMinion(name string, cfg etcdclient.Config) Minion {
	c, err := etcdclient.New(cfg)
	if err != nil {
		log.Fatal(err)
	}

	kapi := etcdclient.NewKeysAPI(c)
	id := utils.GenerateUUID(name)
	rootDir := filepath.Join(EtcdMinionSpace, id.String())
	queueDir := filepath.Join(rootDir, "queue")
	classifierDir := filepath.Join(rootDir, "classifier")
	logDir := filepath.Join(rootDir, "log")
	taskQueue := make(chan *task.Task)

	m := &etcdMinion{
		name:          name,
		rootDir:       rootDir,
		queueDir:      queueDir,
		classifierDir: classifierDir,
		logDir:        logDir,
		id:            id,
		kapi:          kapi,
		taskQueue:     taskQueue,
	}

	return m
}

// Set the human-readable name of the minion in etcd
func (m *etcdMinion) setName() error {
	nameKey := filepath.Join(m.rootDir, "name")
	opts := &etcdclient.SetOptions{
		PrevExist: etcdclient.PrevIgnore,
	}

	_, err := m.kapi.Set(context.Background(), nameKey, m.Name(), opts)
	if err != nil {
		log.Printf("Failed to set name of minion: %s\n", err)
	}

	return err
}

// Set the time the minion was last seen in seconds since the Epoch
func (m *etcdMinion) setLastseen(s int64) error {
	lastseenKey := filepath.Join(m.rootDir, "lastseen")
	lastseenValue := strconv.FormatInt(s, 10)
	opts := &etcdclient.SetOptions{
		PrevExist: etcdclient.PrevIgnore,
	}

	_, err := m.kapi.Set(context.Background(), lastseenKey, lastseenValue, opts)
	if err != nil {
		log.Printf("Failed to set lastseen time: %s\n", err)
	}

	return err
}

// Checks for any pending tasks and sends them
// for processing if there are any
func (m *etcdMinion) checkQueue() error {
	opts := &etcdclient.GetOptions{
		Recursive: true,
		Sort:      true,
	}

	// Get backlog tasks if any
	// If the directory key in etcd is missing that is okay, since
	// it means there are no pending tasks for processing
	resp, err := m.kapi.Get(context.Background(), m.queueDir, opts)
	if err != nil {
		if eerr, ok := err.(etcdclient.Error); !ok || eerr.Code == etcdclient.ErrorCodeKeyNotFound {
			return err
		}
	}

	backlog := resp.Node.Nodes
	if len(backlog) == 0 {
		// No backlog tasks found
		return nil
	}

	log.Printf("Found %d pending tasks in queue", len(backlog))
	for _, node := range backlog {
		t, err := EtcdUnmarshalTask(node)
		m.kapi.Delete(context.Background(), node.Key, nil)

		if err != nil {
			continue
		}

		m.taskQueue <- t
	}

	return nil
}

// Runs periodic jobs such as refreshing classifiers and
// updating the lastseen time every fifteen minutes
func (m *etcdMinion) periodicRunner() {
	schedule := time.Minute * 15
	ticker := time.NewTicker(schedule)
	log.Printf("Periodic scheduler set to run every %s\n", schedule)

	for now := range ticker.C {
		// Run any periodic jobs
		m.Classify()
		m.setLastseen(now.Unix())
	}
}

// Processes new tasks
func (m *etcdMinion) processTask(t *task.Task) error {
	var buf bytes.Buffer

	// Update state of task to indicate that we are now processing it
	t.State = task.TaskStateProcessing
	m.SaveTaskResult(t)

	cmd := exec.Command(t.Command, t.Args...)
	cmd.Stdout = &buf
	cmd.Stderr = &buf

	log.Printf("Processing task %s\n", t.TaskID)

	cmdError := cmd.Run()
	t.TimeProcessed = time.Now().Unix()
	t.Result = buf.String()

	if cmdError != nil {
		log.Printf("Failed to process task %s\n", t.TaskID)
		t.Error = cmdError.Error()
		t.State = task.TaskStateFailed
	} else {
		log.Printf("Finished processing task %s\n", t.TaskID)
		t.State = task.TaskStateSuccess
	}

	m.SaveTaskResult(t)

	return cmdError
}

// Saves a task in etcd
func (m *etcdMinion) SaveTaskResult(t *task.Task) error {
	taskKey := filepath.Join(m.logDir, t.TaskID.String())

	data, err := json.Marshal(t)
	if err != nil {
		log.Printf("Failed to serialize task %s: %s\n", t.TaskID, err)
		return err
	}

	opts := &etcdclient.SetOptions{
		PrevExist: etcdclient.PrevIgnore,
	}

	_, err = m.kapi.Set(context.Background(), taskKey, string(data), opts)
	if err != nil {
		log.Printf("Failed to save task %s: %s\n", t.TaskID, err)
	}

	return err
}

// Unmarshals task from etcd
func EtcdUnmarshalTask(node *etcdclient.Node) (*task.Task, error) {
	task := new(task.Task)
	err := json.Unmarshal([]byte(node.Value), &task)

	return task, err
}

// Returns the minion unique identifier
func (m *etcdMinion) ID() uuid.UUID {
	return m.id
}

// Returns the assigned name of the minion
func (m *etcdMinion) Name() string {
	return m.name
}

// Classifies the minion
func (m *etcdMinion) Classify() error {
	// Classifiers in etcd expire after an hour
	opts := &etcdclient.SetOptions{
		PrevExist: etcdclient.PrevIgnore,
		TTL:       time.Hour,
	}

	// Set/update classifiers in etcd
	for key, _ := range classifier.Registry {
		klassifier, err := classifier.Get(key)

		if err != nil {
			log.Printf("Failed to get classifier %s: %s\n", key, err)
			continue
		}

		// Serialize classifier to JSON and save it in etcd
		data, err := json.Marshal(klassifier)
		if err != nil {
			log.Printf("Failed to serialize classifier: %s\n", key)
			continue
		}

		// Classifier key in etcd
		klassifierKey := filepath.Join(m.classifierDir, key)
		_, err = m.kapi.Set(context.Background(), klassifierKey, string(data), opts)

		if err != nil {
			log.Printf("Failed to set classifier %s: %s\n", key, err)
		}
	}

	return nil
}

// Monitors etcd for new tasks
func (m *etcdMinion) TaskListener(c chan<- *task.Task) error {
	log.Printf("Task listener is watching %s\n", m.queueDir)

	watcherOpts := &etcdclient.WatcherOptions{
		Recursive: true,
	}
	watcher := m.kapi.Watcher(m.queueDir, watcherOpts)

	for {
		resp, err := watcher.Next(context.Background())
		if err != nil {
			log.Printf("Failed to receive task: %s\n", err)
			continue
		}

		// Ignore "delete" events when removing a task from the queue
		action := strings.ToLower(resp.Action)
		if strings.EqualFold(action, "delete") {
			continue
		}

		// Unmarshal and remove task from the queue
		t, err := EtcdUnmarshalTask(resp.Node)
		m.kapi.Delete(context.Background(), resp.Node.Key, nil)

		if err != nil {
			log.Printf("Received invalid task %s: %s\n", resp.Node.Key, err)
			continue
		}

		// Update task state and send it for processing
		log.Printf("Received task %s\n", t.TaskID)
		t.State = task.TaskStateQueued
		t.TimeReceived = time.Now().Unix()
		m.SaveTaskResult(t)
		c <- t
	}

	return nil
}

// Processes new tasks
func (m *etcdMinion) TaskRunner(c <-chan *task.Task) error {
	log.Println("Starting task runner")

	for t := range c {
		if t.IsConcurrent {
			go m.processTask(t)
		} else {
			m.processTask(t)
		}
	}

	return nil
}

// Main entry point of the minion
func (m *etcdMinion) Serve() error {
	err := m.setName()
	if err != nil {
		return err
	}

	// Start minion services
	go m.periodicRunner()
	go m.checkQueue()
	go m.TaskRunner(m.taskQueue)
	go m.TaskListener(m.taskQueue)

	log.Printf("Minion %s is ready to serve", m.ID())

	return nil
}

// Stops the minion and performs any cleanup tasks
func (m *etcdMinion) Stop() error {
	log.Println("Minion is shutting down")

	close(m.taskQueue)

	return nil
}
<|endoftext|>


<commit_before>// sqrt_test.go
package simplemath

import "testing"

func TestSqrt1(t *testing.T) {
  v := Sqrt(16)
  if v != 4 {
    t.Errorf("Sqrt(16) failed. Got %v, expected 4.", v)
  }
}
<commit_msg>update<commit_after>// sqrt_test.go
package simplemath

import "testing"

func TestSqrt1(t *testing.T) {
  v := Sqrt(1677)
  if v != 4 {
    t.Errorf("Sqrt(16) failed. Got %v, expected 4.", v)
  }
}
<|endoftext|>


<commit_before>package objects

import (
	"fmt"
	"io"
	"io/ioutil"
	"strings"

	"github.com/rackspace/gophercloud"
	"github.com/rackspace/gophercloud/pagination"

	"github.com/mitchellh/mapstructure"
)

// Object is a structure that holds information related to a storage object.
type Object struct {
	Bytes        int    `json:"bytes" mapstructure:"bytes"`
	ContentType  string `json:"content_type" mapstructure:"content_type"`
	Hash         string `json:"hash" mapstructure:"hash"`
	LastModified string `json:"last_modified" mapstructure:"last_modified"`
	Name         string `json:"name" mapstructure:"name"`
}

// ObjectPage is a single page of objects that is returned from a call to the
// List function.
type ObjectPage struct {
	pagination.MarkerPageBase
}

// IsEmpty returns true if a ListResult contains no object names.
func (r ObjectPage) IsEmpty() (bool, error) {
	names, err := ExtractNames(r)
	if err != nil {
		return true, err
	}
	return len(names) == 0, nil
}

// LastMarker returns the last object name in a ListResult.
func (r ObjectPage) LastMarker() (string, error) {
	names, err := ExtractNames(r)
	if err != nil {
		return "", err
	}
	if len(names) == 0 {
		return "", nil
	}
	return names[len(names)-1], nil
}

// ExtractInfo is a function that takes a page of objects and returns their full information.
func ExtractInfo(page pagination.Page) ([]Object, error) {
	untyped := page.(ObjectPage).Body.([]interface{})
	results := make([]Object, len(untyped))
	for index, each := range untyped {
		object := each.(map[string]interface{})
		err := mapstructure.Decode(object, &results[index])
		if err != nil {
			return results, err
		}
	}
	return results, nil
}

// ExtractNames is a function that takes a page of objects and returns only their names.
func ExtractNames(page pagination.Page) ([]string, error) {
	casted := page.(ObjectPage)
	ct := casted.Header.Get("Content-Type")
	switch {
	case strings.HasPrefix(ct, "application/json"):
		parsed, err := ExtractInfo(page)
		if err != nil {
			return nil, err
		}

		names := make([]string, 0, len(parsed))
		for _, object := range parsed {
			names = append(names, object.Name)
		}

		return names, nil
	case strings.HasPrefix(ct, "text/plain"):
		names := make([]string, 0, 50)

		body := string(page.(ObjectPage).Body.([]uint8))
		for _, name := range strings.Split(body, "\n") {
			if len(name) > 0 {
				names = append(names, name)
			}
		}

		return names, nil
	case strings.HasPrefix(ct, "text/html"):
		return []string{}, nil
	default:
		return nil, fmt.Errorf("Cannot extract names from response with content-type: [%s]", ct)
	}
}

type commonResult struct {
	gophercloud.Result
}

// DownloadResult is a *http.Response that is returned from a call to the Download function.
type DownloadResult struct {
	commonResult
	Body io.Reader
}

// ExtractContent is a function that takes a DownloadResult's io.Reader body
// and reads all available data into a slice of bytes. Please be aware that due
// the nature of io.Reader is forward-only - meaning that it can only be read
// once and not rewound. You can recreate a reader from the output of this
// function by using bytes.NewReader(downloadBytes)
func (dr DownloadResult) ExtractContent() ([]byte, error) {
	if dr.Err != nil {
		return nil, dr.Err
	}
	body, err := ioutil.ReadAll(dr.Body)
	if err != nil {
		return nil, err
	}
	return body, nil
}

// GetResult is a *http.Response that is returned from a call to the Get function.
type GetResult struct {
	commonResult
}

// ExtractMetadata is a function that takes a GetResult (of type *http.Response)
// and returns the custom metadata associated with the object.
func (gr GetResult) ExtractMetadata() (map[string]string, error) {
	if gr.Err != nil {
		return nil, gr.Err
	}
	metadata := make(map[string]string)
	for k, v := range gr.Header {
		if strings.HasPrefix(k, "X-Object-Meta-") {
			key := strings.TrimPrefix(k, "X-Object-Meta-")
			metadata[key] = v[0]
		}
	}
	return metadata, nil
}

// CreateResult represents the result of a create operation.
type CreateResult struct {
	commonResult
}

// UpdateResult represents the result of an update operation.
type UpdateResult struct {
	commonResult
}

// DeleteResult represents the result of a delete operation.
type DeleteResult struct {
	commonResult
}

// CopyResult represents the result of a copy operation.
type CopyResult struct {
	commonResult
}
<commit_msg>Ensuring that reader is closed after being read<commit_after>package objects

import (
	"fmt"
	"io"
	"io/ioutil"
	"strings"

	"github.com/rackspace/gophercloud"
	"github.com/rackspace/gophercloud/pagination"

	"github.com/mitchellh/mapstructure"
)

// Object is a structure that holds information related to a storage object.
type Object struct {
	Bytes        int    `json:"bytes" mapstructure:"bytes"`
	ContentType  string `json:"content_type" mapstructure:"content_type"`
	Hash         string `json:"hash" mapstructure:"hash"`
	LastModified string `json:"last_modified" mapstructure:"last_modified"`
	Name         string `json:"name" mapstructure:"name"`
}

// ObjectPage is a single page of objects that is returned from a call to the
// List function.
type ObjectPage struct {
	pagination.MarkerPageBase
}

// IsEmpty returns true if a ListResult contains no object names.
func (r ObjectPage) IsEmpty() (bool, error) {
	names, err := ExtractNames(r)
	if err != nil {
		return true, err
	}
	return len(names) == 0, nil
}

// LastMarker returns the last object name in a ListResult.
func (r ObjectPage) LastMarker() (string, error) {
	names, err := ExtractNames(r)
	if err != nil {
		return "", err
	}
	if len(names) == 0 {
		return "", nil
	}
	return names[len(names)-1], nil
}

// ExtractInfo is a function that takes a page of objects and returns their full information.
func ExtractInfo(page pagination.Page) ([]Object, error) {
	untyped := page.(ObjectPage).Body.([]interface{})
	results := make([]Object, len(untyped))
	for index, each := range untyped {
		object := each.(map[string]interface{})
		err := mapstructure.Decode(object, &results[index])
		if err != nil {
			return results, err
		}
	}
	return results, nil
}

// ExtractNames is a function that takes a page of objects and returns only their names.
func ExtractNames(page pagination.Page) ([]string, error) {
	casted := page.(ObjectPage)
	ct := casted.Header.Get("Content-Type")
	switch {
	case strings.HasPrefix(ct, "application/json"):
		parsed, err := ExtractInfo(page)
		if err != nil {
			return nil, err
		}

		names := make([]string, 0, len(parsed))
		for _, object := range parsed {
			names = append(names, object.Name)
		}

		return names, nil
	case strings.HasPrefix(ct, "text/plain"):
		names := make([]string, 0, 50)

		body := string(page.(ObjectPage).Body.([]uint8))
		for _, name := range strings.Split(body, "\n") {
			if len(name) > 0 {
				names = append(names, name)
			}
		}

		return names, nil
	case strings.HasPrefix(ct, "text/html"):
		return []string{}, nil
	default:
		return nil, fmt.Errorf("Cannot extract names from response with content-type: [%s]", ct)
	}
}

type commonResult struct {
	gophercloud.Result
}

// DownloadResult is a *http.Response that is returned from a call to the Download function.
type DownloadResult struct {
	commonResult
	Body io.ReadCloser
}

// ExtractContent is a function that takes a DownloadResult's io.Reader body
// and reads all available data into a slice of bytes. Please be aware that due
// the nature of io.Reader is forward-only - meaning that it can only be read
// once and not rewound. You can recreate a reader from the output of this
// function by using bytes.NewReader(downloadBytes)
func (dr DownloadResult) ExtractContent() ([]byte, error) {
	if dr.Err != nil {
		return nil, dr.Err
	}
	body, err := ioutil.ReadAll(dr.Body)
	if err != nil {
		return nil, err
	}
	dr.Body.Close()
	return body, nil
}

// GetResult is a *http.Response that is returned from a call to the Get function.
type GetResult struct {
	commonResult
}

// ExtractMetadata is a function that takes a GetResult (of type *http.Response)
// and returns the custom metadata associated with the object.
func (gr GetResult) ExtractMetadata() (map[string]string, error) {
	if gr.Err != nil {
		return nil, gr.Err
	}
	metadata := make(map[string]string)
	for k, v := range gr.Header {
		if strings.HasPrefix(k, "X-Object-Meta-") {
			key := strings.TrimPrefix(k, "X-Object-Meta-")
			metadata[key] = v[0]
		}
	}
	return metadata, nil
}

// CreateResult represents the result of a create operation.
type CreateResult struct {
	commonResult
}

// UpdateResult represents the result of an update operation.
type UpdateResult struct {
	commonResult
}

// DeleteResult represents the result of a delete operation.
type DeleteResult struct {
	commonResult
}

// CopyResult represents the result of a copy operation.
type CopyResult struct {
	commonResult
}
<|endoftext|>


<commit_before>// Copyright 2016, RadiantBlue Technologies, Inc.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package logger

import (
	"encoding/json"
	"fmt"
	"log"
	"net/http"
	"sync"
	"time"

	"github.com/venicegeo/pz-gocommon/elasticsearch"
	"github.com/venicegeo/pz-gocommon/gocommon"
	syslogger "github.com/venicegeo/pz-gocommon/syslog"
)

type Service struct {
	sync.Mutex

	stats  Stats
	origin string

	logWriters   []syslogger.Writer
	auditWriters []syslogger.Writer

	esIndex elasticsearch.IIndex
}

func (service *Service) Init(sys *piazza.SystemConfig, logWriters []syslogger.Writer, auditWriters []syslogger.Writer, esi elasticsearch.IIndex) error {

	service.stats.CreatedOn = time.Now()

	service.logWriters = logWriters
	service.auditWriters = auditWriters

	service.esIndex = esi

	service.origin = string(sys.Name)

	return nil
}

func (service *Service) newInternalErrorResponse(err error) *piazza.JsonResponse {
	return &piazza.JsonResponse{
		StatusCode: http.StatusInternalServerError,
		Message:    err.Error(),
		Origin:     service.origin,
	}
}

func (service *Service) newBadRequestResponse(err error) *piazza.JsonResponse {
	return &piazza.JsonResponse{
		StatusCode: http.StatusBadRequest,
		Message:    err.Error(),
		Origin:     service.origin,
	}
}

func (service *Service) GetRoot() *piazza.JsonResponse {
	resp := &piazza.JsonResponse{
		StatusCode: 200,
		Data:       "Hi. I'm pz-logger.",
	}

	err := resp.SetType()
	if err != nil {
		return service.newInternalErrorResponse(err)
	}

	return resp
}

func (service *Service) GetStats() *piazza.JsonResponse {
	service.Lock()
	t := service.stats
	service.Unlock()

	resp := &piazza.JsonResponse{
		StatusCode: http.StatusOK,
		Data:       t,
	}

	err := resp.SetType()
	if err != nil {
		return service.newInternalErrorResponse(err)
	}

	return resp
}

func createQueryDslAsString(
	pagination *piazza.JsonPagination,
	params *piazza.HttpQueryParams) (string, error) {

	must := []map[string]interface{}{}

	service, err := params.GetAsString("service", "")
	if err != nil {
		return "", err
	}

	contains, err := params.GetAsString("contains", "")
	if err != nil {
		return "", err
	}

	before, err := params.GetBefore(time.Time{})
	if err != nil {
		return "", err
	}

	after, err := params.GetAfter(time.Time{})
	if err != nil {
		return "", err
	}

	if service != "" {
		must = append(must, map[string]interface{}{
			"match": map[string]interface{}{
				"service": service,
			},
		})
	}
	if contains != "" {
		must = append(must, map[string]interface{}{
			"multi_match": map[string]interface{}{
				"query":  contains,
				"fields": []string{"address", "message", "service", "severity"},
			},
		})
	}

	if !after.IsZero() || !before.IsZero() {
		rangeParams := map[string]time.Time{}

		if !after.IsZero() {
			rangeParams["gte"] = after
		}

		if !before.IsZero() {
			rangeParams["lte"] = before
		}

		must = append(must, map[string]interface{}{
			"range": map[string]interface{}{
				"timeStamp": rangeParams,
			},
		})
	}

	if len(must) == 0 {
		return "", nil
	}

	dsl := map[string]interface{}{
		"query": map[string]interface{}{
			"filtered": map[string]interface{}{
				"query": map[string]interface{}{
					"bool": map[string]interface{}{
						"must": must,
					},
				},
			},
		},
		"size": pagination.PerPage,
		"from": pagination.PerPage * pagination.Page,
	}

	dsl["sort"] = map[string]string{
		pagination.SortBy: string(pagination.Order),
	}

	output, err := json.Marshal(dsl)
	if err != nil {
		return "", err
	}

	return string(output), nil
}

func (service *Service) PostSyslog(mNew *syslogger.Message) *piazza.JsonResponse {
	err := mNew.Validate()
	if err != nil {
		return service.newBadRequestResponse(err)
	}

	err = service.postSyslog(mNew)
	if err != nil {
		return service.newInternalErrorResponse(err)
	}

	resp := &piazza.JsonResponse{
		StatusCode: http.StatusOK,
	}
	return resp
}

func (service *Service) postSyslog(mssg *syslogger.Message) error {

	var err error
	isAudit := mssg.AuditData != nil

	for _, writer := range service.logWriters {
		err = writer.Write(mssg)
		if err != nil {
			return fmt.Errorf("syslog.Service.postSyslog: %s", err.Error())
		}
	}

	if isAudit {
		for _, writer := range service.auditWriters {
			err = writer.Write(mssg)
			if err != nil {
				return fmt.Errorf("syslog.Service.postSyslog (audit): %s", err.Error())
			}
		}
	}

	service.Lock()
	service.stats.NumMessages++
	service.Unlock()

	return nil
}

func (service *Service) getMessageCommon(params *piazza.HttpQueryParams) (*elasticsearch.SearchResult, *piazza.JsonPagination, *piazza.JsonResponse) {
	pagination, err := piazza.NewJsonPagination(params)
	if err != nil {
		return nil, nil, service.newBadRequestResponse(err)
	}

	paginationCreatedOnToTimeStamp(pagination)

	dsl, err := createQueryDslAsString(pagination, params)
	if err != nil {
		return nil, pagination, service.newBadRequestResponse(err)
	}

	var searchResult *elasticsearch.SearchResult

	if dsl == "" {
		searchResult, err = service.esIndex.FilterByMatchAll(logSchema, pagination)
	} else {
		searchResult, err = service.esIndex.SearchByJSON(logSchema, dsl)
	}
	if err != nil {
		return nil, pagination, service.newInternalErrorResponse(err)
	}

	return searchResult, pagination, nil
}

func (service *Service) GetSyslog(params *piazza.HttpQueryParams) *piazza.JsonResponse {
	var err error

	format, err := params.GetAsString("format", "json")
	if err != nil {
		return service.newInternalErrorResponse(err)
	}

	searchResult, pagination, jErr := service.getMessageCommon(params)
	if jErr != nil {
		return jErr
	}

	lines, err := extractFromSearchResult(searchResult)
	if err != nil {
		return service.newInternalErrorResponse(err)
	}

	var data interface{} = lines

	if format == "string" {
		t := make([]string, len(lines))
		for i := 0; i < len(lines); i++ {
			t[i] = lines[i].String()
		}
		data = t
	}

	pagination.Count = int(searchResult.TotalHits())
	resp := &piazza.JsonResponse{
		StatusCode: http.StatusOK,
		Data:       data,
		Pagination: pagination,
	}

	err = resp.SetType()
	if err != nil {
		return service.newInternalErrorResponse(err)
	}

	return resp
}

func extractFromSearchResult(searchResult *elasticsearch.SearchResult) ([]syslogger.Message, error) {

	var lines = make([]syslogger.Message, 0, len(*searchResult.GetHits()))

	if searchResult == nil || searchResult.GetHits() == nil {
		return lines, nil
	}

	for _, hit := range *searchResult.GetHits() {
		if hit.Source == nil {
			log.Printf("null source hit")
			continue
		}

		msg := syslogger.Message{}
		err := json.Unmarshal(*hit.Source, &msg)
		if err != nil {
			log.Printf("UNABLE TO PARSE: %s", string(*hit.Source))
			continue
		}

		// just in case
		/*err = msg.Validate()
		if err != nil {
			log.Printf("UNABLE TO VALIDATE: %s", string(*hit.Source))
			continue
		}*/

		lines = append(lines, msg)
	}

	return lines, nil
}

func (service *Service) PostQuery(params *piazza.HttpQueryParams, jsnQuery string) *piazza.JsonResponse {

	pagination, err := piazza.NewJsonPagination(params)
	if err != nil {
		return service.newBadRequestResponse(err)
	}

	searchResult, err := service.esIndex.SearchByJSON(logSchema, jsnQuery)
	if err != nil {
		return service.newInternalErrorResponse(err)
	}

	lines, err := extractFromSearchResult(searchResult)
	if err != nil {
		return service.newInternalErrorResponse(err)
	}

	pagination.Count = int(searchResult.TotalHits())
	resp := &piazza.JsonResponse{
		StatusCode: http.StatusOK,
		Data:       lines,
		Pagination: pagination,
	}

	err = resp.SetType()
	if err != nil {
		return service.newInternalErrorResponse(err)
	}

	return resp
}
<commit_msg>update query for new Message format<commit_after>// Copyright 2016, RadiantBlue Technologies, Inc.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package logger

import (
	"encoding/json"
	"fmt"
	"log"
	"net/http"
	"sync"
	"time"

	"github.com/venicegeo/pz-gocommon/elasticsearch"
	"github.com/venicegeo/pz-gocommon/gocommon"
	syslogger "github.com/venicegeo/pz-gocommon/syslog"
)

type Service struct {
	sync.Mutex

	stats  Stats
	origin string

	logWriters   []syslogger.Writer
	auditWriters []syslogger.Writer

	esIndex elasticsearch.IIndex
}

func (service *Service) Init(sys *piazza.SystemConfig, logWriters []syslogger.Writer, auditWriters []syslogger.Writer, esi elasticsearch.IIndex) error {

	service.stats.CreatedOn = time.Now()

	service.logWriters = logWriters
	service.auditWriters = auditWriters

	service.esIndex = esi

	service.origin = string(sys.Name)

	return nil
}

func (service *Service) newInternalErrorResponse(err error) *piazza.JsonResponse {
	return &piazza.JsonResponse{
		StatusCode: http.StatusInternalServerError,
		Message:    err.Error(),
		Origin:     service.origin,
	}
}

func (service *Service) newBadRequestResponse(err error) *piazza.JsonResponse {
	return &piazza.JsonResponse{
		StatusCode: http.StatusBadRequest,
		Message:    err.Error(),
		Origin:     service.origin,
	}
}

func (service *Service) GetRoot() *piazza.JsonResponse {
	resp := &piazza.JsonResponse{
		StatusCode: 200,
		Data:       "Hi. I'm pz-logger.",
	}

	err := resp.SetType()
	if err != nil {
		return service.newInternalErrorResponse(err)
	}

	return resp
}

func (service *Service) GetStats() *piazza.JsonResponse {
	service.Lock()
	t := service.stats
	service.Unlock()

	resp := &piazza.JsonResponse{
		StatusCode: http.StatusOK,
		Data:       t,
	}

	err := resp.SetType()
	if err != nil {
		return service.newInternalErrorResponse(err)
	}

	return resp
}

func createQueryDslAsString(
	pagination *piazza.JsonPagination,
	params *piazza.HttpQueryParams) (string, error) {

	must := []map[string]interface{}{}

	service, err := params.GetAsString("service", "")
	if err != nil {
		return "", err
	}

	contains, err := params.GetAsString("contains", "")
	if err != nil {
		return "", err
	}

	before, err := params.GetBefore(time.Time{})
	if err != nil {
		return "", err
	}

	after, err := params.GetAfter(time.Time{})
	if err != nil {
		return "", err
	}

	if service != "" {
		must = append(must, map[string]interface{}{
			"match": map[string]interface{}{
				"application": service,
			},
		})
	}
	if contains != "" {
		must = append(must, map[string]interface{}{
			"multi_match": map[string]interface{}{
				"query":  contains,
				"fields": []string{"hostName", "application", "process", "messageId", "message"},
			},
		})
	}

	if !after.IsZero() || !before.IsZero() {
		rangeParams := map[string]time.Time{}

		if !after.IsZero() {
			rangeParams["gte"] = after
		}

		if !before.IsZero() {
			rangeParams["lte"] = before
		}

		must = append(must, map[string]interface{}{
			"range": map[string]interface{}{
				"timeStamp": rangeParams,
			},
		})
	}

	if len(must) == 0 {
		return "", nil
	}

	dsl := map[string]interface{}{
		"query": map[string]interface{}{
			"filtered": map[string]interface{}{
				"query": map[string]interface{}{
					"bool": map[string]interface{}{
						"must": must,
					},
				},
			},
		},
		"size": pagination.PerPage,
		"from": pagination.PerPage * pagination.Page,
	}

	dsl["sort"] = map[string]string{
		pagination.SortBy: string(pagination.Order),
	}

	output, err := json.Marshal(dsl)
	if err != nil {
		return "", err
	}

	return string(output), nil
}

func (service *Service) PostSyslog(mNew *syslogger.Message) *piazza.JsonResponse {
	err := mNew.Validate()
	if err != nil {
		return service.newBadRequestResponse(err)
	}

	err = service.postSyslog(mNew)
	if err != nil {
		return service.newInternalErrorResponse(err)
	}

	resp := &piazza.JsonResponse{
		StatusCode: http.StatusOK,
	}
	return resp
}

func (service *Service) postSyslog(mssg *syslogger.Message) error {

	var err error
	isAudit := mssg.AuditData != nil

	for _, writer := range service.logWriters {
		err = writer.Write(mssg)
		if err != nil {
			return fmt.Errorf("syslog.Service.postSyslog: %s", err.Error())
		}
	}

	if isAudit {
		for _, writer := range service.auditWriters {
			err = writer.Write(mssg)
			if err != nil {
				return fmt.Errorf("syslog.Service.postSyslog (audit): %s", err.Error())
			}
		}
	}

	service.Lock()
	service.stats.NumMessages++
	service.Unlock()

	return nil
}

func (service *Service) getMessageCommon(params *piazza.HttpQueryParams) (*elasticsearch.SearchResult, *piazza.JsonPagination, *piazza.JsonResponse) {
	pagination, err := piazza.NewJsonPagination(params)
	if err != nil {
		return nil, nil, service.newBadRequestResponse(err)
	}

	paginationCreatedOnToTimeStamp(pagination)

	dsl, err := createQueryDslAsString(pagination, params)
	if err != nil {
		return nil, pagination, service.newBadRequestResponse(err)
	}

	var searchResult *elasticsearch.SearchResult

	if dsl == "" {
		searchResult, err = service.esIndex.FilterByMatchAll(logSchema, pagination)
	} else {
		searchResult, err = service.esIndex.SearchByJSON(logSchema, dsl)
	}
	if err != nil {
		return nil, pagination, service.newInternalErrorResponse(err)
	}

	return searchResult, pagination, nil
}

func (service *Service) GetSyslog(params *piazza.HttpQueryParams) *piazza.JsonResponse {
	var err error

	format, err := params.GetAsString("format", "json")
	if err != nil {
		return service.newInternalErrorResponse(err)
	}

	searchResult, pagination, jErr := service.getMessageCommon(params)
	if jErr != nil {
		return jErr
	}

	lines, err := extractFromSearchResult(searchResult)
	if err != nil {
		return service.newInternalErrorResponse(err)
	}

	var data interface{} = lines

	if format == "string" {
		t := make([]string, len(lines))
		for i := 0; i < len(lines); i++ {
			t[i] = lines[i].String()
		}
		data = t
	}

	pagination.Count = int(searchResult.TotalHits())
	resp := &piazza.JsonResponse{
		StatusCode: http.StatusOK,
		Data:       data,
		Pagination: pagination,
	}

	err = resp.SetType()
	if err != nil {
		return service.newInternalErrorResponse(err)
	}

	return resp
}

func extractFromSearchResult(searchResult *elasticsearch.SearchResult) ([]syslogger.Message, error) {

	var lines = make([]syslogger.Message, 0, len(*searchResult.GetHits()))

	if searchResult == nil || searchResult.GetHits() == nil {
		return lines, nil
	}

	for _, hit := range *searchResult.GetHits() {
		if hit.Source == nil {
			log.Printf("null source hit")
			continue
		}

		msg := syslogger.Message{}
		err := json.Unmarshal(*hit.Source, &msg)
		if err != nil {
			log.Printf("UNABLE TO PARSE: %s", string(*hit.Source))
			continue
		}

		// just in case
		/*err = msg.Validate()
		if err != nil {
			log.Printf("UNABLE TO VALIDATE: %s", string(*hit.Source))
			continue
		}*/

		lines = append(lines, msg)
	}

	return lines, nil
}

func (service *Service) PostQuery(params *piazza.HttpQueryParams, jsnQuery string) *piazza.JsonResponse {

	pagination, err := piazza.NewJsonPagination(params)
	if err != nil {
		return service.newBadRequestResponse(err)
	}

	searchResult, err := service.esIndex.SearchByJSON(logSchema, jsnQuery)
	if err != nil {
		return service.newInternalErrorResponse(err)
	}

	lines, err := extractFromSearchResult(searchResult)
	if err != nil {
		return service.newInternalErrorResponse(err)
	}

	pagination.Count = int(searchResult.TotalHits())
	resp := &piazza.JsonResponse{
		StatusCode: http.StatusOK,
		Data:       lines,
		Pagination: pagination,
	}

	err = resp.SetType()
	if err != nil {
		return service.newInternalErrorResponse(err)
	}

	return resp
}
<|endoftext|>


<commit_before>package slack_incoming_webhooks

type Attachment struct {
	Fallback   string   `json:"fallback,omitempty"`
	Color      string   `json:"color,omitempty"`
	Pretext    string   `json:"pretext,omitempty"`
	AuthorName string   `json:"author_name,omitempty"`
	AuthorLink string   `json:"author_link,omitempty"`
	AuthorIcon string   `json:"author_icon,omitempty"`
	Title      string   `json:"title,omitempty"`
	TitleLink  string   `json:"title_link,omitempty"`
	Text       string   `json:"text,omitempty"`
	Fields     []*Field `json:"fields,omitempty"`
	ImageURL   string   `json:"image_url,omitempty"`
	ThumbURL   string   `json:"thumb_url,omitempty"`
}

func (a Attachment) IsEmpty() bool {
	return a.Fallback == "" &&
		a.Color == "" &&
		a.Pretext == "" &&
		a.AuthorName == "" &&
		a.AuthorLink == "" &&
		a.AuthorIcon == "" &&
		a.Title == "" &&
		a.TitleLink == "" &&
		a.Text == "" &&
		len(a.Fields) == 0 &&
		a.ImageURL == "" &&
		a.ThumbURL == ""
}

func (a *Attachment) AddField(f *Field) {
	if !f.IsEmpty() {
		a.Fields = append(a.Fields, f)
	}
}
<commit_msg>Enable markdown attachments<commit_after>package slack_incoming_webhooks

type Attachment struct {
	Fallback   string   `json:"fallback,omitempty"`
	Color      string   `json:"color,omitempty"`
	Pretext    string   `json:"pretext,omitempty"`
	AuthorName string   `json:"author_name,omitempty"`
	AuthorLink string   `json:"author_link,omitempty"`
	AuthorIcon string   `json:"author_icon,omitempty"`
	Title      string   `json:"title,omitempty"`
	TitleLink  string   `json:"title_link,omitempty"`
	Text       string   `json:"text,omitempty"`
	Fields     []*Field `json:"fields,omitempty"`
	ImageURL   string   `json:"image_url,omitempty"`
	ThumbURL   string   `json:"thumb_url,omitempty"`
	MarkdownIn []string `json:"mrkdwn_in,omitempty"`
}

func (a Attachment) IsEmpty() bool {
	return a.Fallback == "" &&
		a.Color == "" &&
		a.Pretext == "" &&
		a.AuthorName == "" &&
		a.AuthorLink == "" &&
		a.AuthorIcon == "" &&
		a.Title == "" &&
		a.TitleLink == "" &&
		a.Text == "" &&
		len(a.Fields) == 0 &&
		a.ImageURL == "" &&
		a.ThumbURL == "" &&
		len(a.MarkdownIn) == 0
}

func (a *Attachment) AddField(f *Field) {
	if !f.IsEmpty() {
		a.Fields = append(a.Fields, f)
	}
}
<|endoftext|>


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

package annotations_test

import (
	jc "github.com/juju/testing/checkers"
	gc "gopkg.in/check.v1"

	"github.com/juju/juju/api/annotations"
	basetesting "github.com/juju/juju/api/base/testing"
	"github.com/juju/juju/apiserver/params"
	jujutesting "github.com/juju/juju/juju/testing"
	coretesting "github.com/juju/juju/testing"
	"github.com/juju/juju/testing/factory"
)

type annotationsMockSuite struct {
	coretesting.BaseSuite
	annotationsClient *annotations.Client
}

var _ = gc.Suite(&annotationsMockSuite{})

func (s *annotationsMockSuite) TestSetEntitiesAnnotation(c *gc.C) {
	var called bool
	annts := map[string]string{"annotation": "test"}
	apiCaller := basetesting.APICallerFunc(
		func(
			objType string,
			version int,
			id, request string,
			a, result interface{}) error {
			called = true
			c.Check(objType, gc.Equals, "Annotations")
			c.Check(id, gc.Equals, "")
			c.Check(request, gc.Equals, "Set")

			args, ok := a.(params.AnnotationsSet)
			c.Assert(ok, jc.IsTrue)
			expected := params.AnnotationsSet{
				Annotations: []params.EntityAnnotations{
					{Entity: params.Entity{"charmA"}, Annotations: annts},
					{Entity: params.Entity{"serviceB"}, Annotations: annts},
				}}
			c.Assert(args, gc.DeepEquals, expected)
			return nil
		})
	annotationsClient := annotations.NewClient(apiCaller)
	err := annotationsClient.Set(
		constructTestEntityAnnotations(
			[]string{"charmA", "serviceB"},
			annts))
	c.Assert(err, jc.ErrorIsNil)
	c.Assert(called, jc.IsTrue)
}

func (s *annotationsMockSuite) TestGetEntitiesAnnotations(c *gc.C) {
	var called bool
	apiCaller := basetesting.APICallerFunc(
		func(
			objType string,
			version int,
			id, request string,
			a, response interface{}) error {
			called = true
			c.Check(objType, gc.Equals, "Annotations")
			c.Check(id, gc.Equals, "")
			c.Check(request, gc.Equals, "Get")
			args, ok := a.(params.Entities)
			c.Assert(ok, jc.IsTrue)
			c.Assert(args.Entities, gc.HasLen, 1)
			c.Assert(args.Entities[0], gc.DeepEquals, params.Entity{"charm"})

			result := response.(*params.AnnotationsGetResults)
			facadeAnnts := map[string]string{
				"annotations": "test",
			}
			entitiesAnnts := params.AnnotationsGetResult{
				Entity:      params.Entity{"charm"},
				Annotations: facadeAnnts,
			}
			result.Results = []params.AnnotationsGetResult{entitiesAnnts}
			return nil
		})
	annotationsClient := annotations.NewClient(apiCaller)
	found, err := annotationsClient.Get([]string{"charm"})
	c.Assert(err, jc.ErrorIsNil)
	c.Assert(called, jc.IsTrue)
	c.Assert(found, gc.HasLen, 1)
}
func constructTestEntityAnnotations(tags []string, pairs map[string]string) map[string]map[string]string {
	result := make(map[string]map[string]string)
	for _, tag := range tags {
		result[tag] = pairs
	}
	return result
}

type annotationsSuite struct {
	jujutesting.JujuConnSuite
	annotationsClient *annotations.Client
}

var _ = gc.Suite(&annotationsSuite{})

func (s *annotationsSuite) SetUpTest(c *gc.C) {
	s.JujuConnSuite.SetUpTest(c)
	s.annotationsClient = annotations.NewClient(s.APIState)
	c.Assert(s.annotationsClient, gc.NotNil)
}

func (s *annotationsSuite) TearDownTest(c *gc.C) {
	s.annotationsClient.ClientFacade.Close()
	s.JujuConnSuite.TearDownTest(c)
}

func (s *annotationsSuite) TestAnnotationFacadeCall(c *gc.C) {
	charm := s.Factory.MakeCharm(c, &factory.CharmParams{Name: "wordpress"})

	annts := map[string]string{"annotation": "test"}
	err := s.annotationsClient.Set(
		constructTestEntityAnnotations([]string{charm.Tag().String()},
			annts))
	c.Assert(err, jc.ErrorIsNil)

	found, err := s.annotationsClient.Get([]string{charm.Tag().String()})
	c.Assert(err, jc.ErrorIsNil)
	c.Assert(found, gc.HasLen, 1)
}
<commit_msg>Tests that different entity annotations are set correctly.<commit_after>// Copyright 2015 Canonical Ltd.
// Licensed under the AGPLv3, see LICENCE file for details.

package annotations_test

import (
	jc "github.com/juju/testing/checkers"
	gc "gopkg.in/check.v1"

	"github.com/juju/juju/api/annotations"
	basetesting "github.com/juju/juju/api/base/testing"
	"github.com/juju/juju/apiserver/params"
	jujutesting "github.com/juju/juju/juju/testing"
	coretesting "github.com/juju/juju/testing"
	"github.com/juju/juju/testing/factory"
)

type annotationsMockSuite struct {
	coretesting.BaseSuite
	annotationsClient *annotations.Client
}

var _ = gc.Suite(&annotationsMockSuite{})

func (s *annotationsMockSuite) TestSetEntitiesAnnotation(c *gc.C) {
	var called bool
	annts := map[string]string{"annotation1": "test"}
	annts2 := map[string]string{"annotation2": "test"}
	apiCaller := basetesting.APICallerFunc(
		func(
			objType string,
			version int,
			id, request string,
			a, result interface{}) error {
			called = true
			c.Check(objType, gc.Equals, "Annotations")
			c.Check(id, gc.Equals, "")
			c.Check(request, gc.Equals, "Set")

			args, ok := a.(params.AnnotationsSet)
			c.Assert(ok, jc.IsTrue)
			expected := params.AnnotationsSet{
				Annotations: []params.EntityAnnotations{
					{Entity: params.Entity{"charmA"}, Annotations: annts},
					{Entity: params.Entity{"serviceB"}, Annotations: annts2},
				}}
			c.Assert(args, gc.DeepEquals, expected)
			return nil
		})
	annotationsClient := annotations.NewClient(apiCaller)
	err := annotationsClient.Set(
		map[string]map[string]string{
			"charmA":   annts,
			"serviceB": annts2,
		})
	c.Assert(err, jc.ErrorIsNil)
	c.Assert(called, jc.IsTrue)
}

func (s *annotationsMockSuite) TestGetEntitiesAnnotations(c *gc.C) {
	var called bool
	apiCaller := basetesting.APICallerFunc(
		func(
			objType string,
			version int,
			id, request string,
			a, response interface{}) error {
			called = true
			c.Check(objType, gc.Equals, "Annotations")
			c.Check(id, gc.Equals, "")
			c.Check(request, gc.Equals, "Get")
			args, ok := a.(params.Entities)
			c.Assert(ok, jc.IsTrue)
			c.Assert(args.Entities, gc.HasLen, 1)
			c.Assert(args.Entities[0], gc.DeepEquals, params.Entity{"charm"})

			result := response.(*params.AnnotationsGetResults)
			facadeAnnts := map[string]string{
				"annotations": "test",
			}
			entitiesAnnts := params.AnnotationsGetResult{
				Entity:      params.Entity{"charm"},
				Annotations: facadeAnnts,
			}
			result.Results = []params.AnnotationsGetResult{entitiesAnnts}
			return nil
		})
	annotationsClient := annotations.NewClient(apiCaller)
	found, err := annotationsClient.Get([]string{"charm"})
	c.Assert(err, jc.ErrorIsNil)
	c.Assert(called, jc.IsTrue)
	c.Assert(found, gc.HasLen, 1)
}

type annotationsSuite struct {
	jujutesting.JujuConnSuite
	annotationsClient *annotations.Client
}

var _ = gc.Suite(&annotationsSuite{})

func (s *annotationsSuite) SetUpTest(c *gc.C) {
	s.JujuConnSuite.SetUpTest(c)
	s.annotationsClient = annotations.NewClient(s.APIState)
	c.Assert(s.annotationsClient, gc.NotNil)
}

func (s *annotationsSuite) TearDownTest(c *gc.C) {
	s.annotationsClient.ClientFacade.Close()
	s.JujuConnSuite.TearDownTest(c)
}

func (s *annotationsSuite) TestAnnotationFacadeCall(c *gc.C) {
	charm := s.Factory.MakeCharm(c, &factory.CharmParams{Name: "wordpress"})

	annts := map[string]string{"annotation": "test"}
	err := s.annotationsClient.Set(
		map[string]map[string]string{
			charm.Tag().String(): annts,
		})
	c.Assert(err, jc.ErrorIsNil)

	found, err := s.annotationsClient.Get([]string{charm.Tag().String()})
	c.Assert(err, jc.ErrorIsNil)
	c.Assert(found, gc.HasLen, 1)
}
<|endoftext|>


<commit_before>//go:build linux && cgo && !agent
// +build linux,cgo,!agent

package db

import "fmt"

// Code generation directives.
//
//go:generate -command mapper lxd-generate db mapper -t devices.mapper.go
//go:generate mapper reset
//
//go:generate mapper stmt -p db -e device objects
//go:generate mapper stmt -p db -e device create struct=Device
//go:generate mapper stmt -p db -e device delete
//
//go:generate mapper method -p db -e device GetMany
//go:generate mapper method -p db -e device Create struct=Device
//go:generate mapper method -p db -e device Update struct=Device
//go:generate mapper method -p db -e device DeleteMany

// DeviceType represents the types of supported devices.
type DeviceType int

// Device is a reference struct representing another entity's device.
type Device struct {
	ID          int
	ReferenceID int
	Name        string
	Type        DeviceType
	Config      map[string]string
}

// Supported device types.
const (
	TypeNone        = DeviceType(0)
	TypeNIC         = DeviceType(1)
	TypeDisk        = DeviceType(2)
	TypeUnixChar    = DeviceType(3)
	TypeUnixBlock   = DeviceType(4)
	TypeUSB         = DeviceType(5)
	TypeGPU         = DeviceType(6)
	TypeInfiniband  = DeviceType(7)
	TypeProxy       = DeviceType(8)
	TypeUnixHotplug = DeviceType(9)
	TypeTPM         = DeviceType(10)
	TypePCI         = DeviceType(11)
)

func (t DeviceType) String() string {
	switch t {
	case TypeNone:
		return "none"
	case TypeNIC:
		return "nic"
	case TypeDisk:
		return "disk"
	case TypeUnixChar:
		return "unix-char"
	case TypeUnixBlock:
		return "unix-block"
	case TypeUSB:
		return "usb"
	case TypeGPU:
		return "gpu"
	case TypeInfiniband:
		return "infiniband"
	case TypeProxy:
		return "proxy"
	case TypeUnixHotplug:
		return "unix-hotplug"
	case TypeTPM:
		return "tpm"
	case TypePCI:
		return "pci"
	}

	return ""
}

// NewDeviceType determines the device type from the given string, if supported.
func NewDeviceType(t string) (DeviceType, error) {
	switch t {
	case "none":
		return TypeNone, nil
	case "nic":
		return TypeNIC, nil
	case "disk":
		return TypeDisk, nil
	case "unix-char":
		return TypeUnixChar, nil
	case "unix-block":
		return TypeUnixBlock, nil
	case "usb":
		return TypeUSB, nil
	case "gpu":
		return TypeGPU, nil
	case "infiniband":
		return TypeInfiniband, nil
	case "proxy":
		return TypeProxy, nil
	case "unix-hotplug":
		return TypeUnixHotplug, nil
	case "tpm":
		return TypeTPM, nil
	case "pci":
		return TypePCI, nil
	default:
		return -1, fmt.Errorf("Invalid device type %s", t)
	}
}
<commit_msg>lxd/db/devices: Devices to/from API helpers<commit_after>//go:build linux && cgo && !agent
// +build linux,cgo,!agent

package db

import (
	"fmt"
)

// Code generation directives.
//
//go:generate -command mapper lxd-generate db mapper -t devices.mapper.go
//go:generate mapper reset
//
//go:generate mapper stmt -p db -e device objects
//go:generate mapper stmt -p db -e device create struct=Device
//go:generate mapper stmt -p db -e device delete
//
//go:generate mapper method -p db -e device GetMany
//go:generate mapper method -p db -e device Create struct=Device
//go:generate mapper method -p db -e device Update struct=Device
//go:generate mapper method -p db -e device DeleteMany

// DeviceType represents the types of supported devices.
type DeviceType int

// Device is a reference struct representing another entity's device.
type Device struct {
	ID          int
	ReferenceID int
	Name        string
	Type        DeviceType
	Config      map[string]string
}

// Supported device types.
const (
	TypeNone        = DeviceType(0)
	TypeNIC         = DeviceType(1)
	TypeDisk        = DeviceType(2)
	TypeUnixChar    = DeviceType(3)
	TypeUnixBlock   = DeviceType(4)
	TypeUSB         = DeviceType(5)
	TypeGPU         = DeviceType(6)
	TypeInfiniband  = DeviceType(7)
	TypeProxy       = DeviceType(8)
	TypeUnixHotplug = DeviceType(9)
	TypeTPM         = DeviceType(10)
	TypePCI         = DeviceType(11)
)

func (t DeviceType) String() string {
	switch t {
	case TypeNone:
		return "none"
	case TypeNIC:
		return "nic"
	case TypeDisk:
		return "disk"
	case TypeUnixChar:
		return "unix-char"
	case TypeUnixBlock:
		return "unix-block"
	case TypeUSB:
		return "usb"
	case TypeGPU:
		return "gpu"
	case TypeInfiniband:
		return "infiniband"
	case TypeProxy:
		return "proxy"
	case TypeUnixHotplug:
		return "unix-hotplug"
	case TypeTPM:
		return "tpm"
	case TypePCI:
		return "pci"
	}

	return ""
}

// NewDeviceType determines the device type from the given string, if supported.
func NewDeviceType(t string) (DeviceType, error) {
	switch t {
	case "none":
		return TypeNone, nil
	case "nic":
		return TypeNIC, nil
	case "disk":
		return TypeDisk, nil
	case "unix-char":
		return TypeUnixChar, nil
	case "unix-block":
		return TypeUnixBlock, nil
	case "usb":
		return TypeUSB, nil
	case "gpu":
		return TypeGPU, nil
	case "infiniband":
		return TypeInfiniband, nil
	case "proxy":
		return TypeProxy, nil
	case "unix-hotplug":
		return TypeUnixHotplug, nil
	case "tpm":
		return TypeTPM, nil
	case "pci":
		return TypePCI, nil
	default:
		return -1, fmt.Errorf("Invalid device type %s", t)
	}
}

// DevicesToAPI takes a map of devices and converts them to API format
func DevicesToAPI(devices map[string]Device) map[string]map[string]string {
	config := map[string]map[string]string{}
	for _, d := range devices {
		if d.Config == nil {
			d.Config = map[string]string{}
		}

		config[d.Name] = d.Config
		config[d.Name]["type"] = d.Type.String()
	}

	return config
}

// APIToDevices takes an API format devices map and converts it to a map of db.Device.
func APIToDevices(apiDevices map[string]map[string]string) (map[string]Device, error) {
	devices := map[string]Device{}
	for name, config := range apiDevices {
		newType, err := NewDeviceType(config["type"])
		if err != nil {
			return nil, err
		}

		device := Device{
			Name:   name,
			Type:   newType,
			Config: config,
		}

		devices[name] = device
	}

	return devices, nil
}
<|endoftext|>


<commit_before>package main

import (
	"flag"
	"log"
	"github.com/hanwen/go-fuse/termite"
	"io/ioutil"
	"strings"
)

func main() {
	cachedir := flag.String("cachedir", "/tmp/fsserver-cache", "content cache")
	workers := flag.String("workers", "localhost:1235", "comma separated list of worker addresses")
	socket := flag.String("socket", ".termite-socket", "socket to listen for commands")
	exclude := flag.String("exclude", "/sys,/proc,/dev", "prefixes to not export.")
	secretFile := flag.String("secret", "/tmp/secret.txt", "file containing password.")
	jobs := flag.Int("jobs", 1, "number of jobs to run")

	flag.Parse()
	secret, err := ioutil.ReadFile(*secretFile)
	if err != nil {
		log.Fatal("ReadFile", err)
	}

	workerList := strings.Split(*workers, ",", -1)
	excludeList := strings.Split(*exclude, ",", -1)
	c := termite.NewDiskFileCache(*cachedir)
	master := termite.NewMaster(
		c, workerList, secret, excludeList, *jobs)
	master.Start(*socket)
}
<commit_msg>Do not export /selinux.<commit_after>package main

import (
	"flag"
	"log"
	"github.com/hanwen/go-fuse/termite"
	"io/ioutil"
	"strings"
)

func main() {
	cachedir := flag.String("cachedir", "/tmp/fsserver-cache", "content cache")
	workers := flag.String("workers", "localhost:1235", "comma separated list of worker addresses")
	socket := flag.String("socket", ".termite-socket", "socket to listen for commands")
	exclude := flag.String("exclude", "/sys,/proc,/dev,/selinux", "prefixes to not export.")
	secretFile := flag.String("secret", "/tmp/secret.txt", "file containing password.")
	jobs := flag.Int("jobs", 1, "number of jobs to run")

	flag.Parse()
	secret, err := ioutil.ReadFile(*secretFile)
	if err != nil {
		log.Fatal("ReadFile", err)
	}

	workerList := strings.Split(*workers, ",", -1)
	excludeList := strings.Split(*exclude, ",", -1)
	c := termite.NewDiskFileCache(*cachedir)
	master := termite.NewMaster(
		c, workerList, secret, excludeList, *jobs)
	master.Start(*socket)
}
<|endoftext|>


<commit_before>package client

import (
	"bytes"
	"encoding/json"
	"io/ioutil"
	"net/http"

	"github.com/danielkrainas/tinkersnest/api/v1"
)

type BlogAPI interface {
	SearchPosts() ([]*v1.Post, error)
	CreatePost(create *v1.CreatePostRequest) (*v1.Post, error)
}

type blogAPI struct {
	*Client
}

func (c *Client) Blog() BlogAPI {
	return &blogAPI{c}
}

func (api *blogAPI) SearchPosts() ([]*v1.Post, error) {
	url, err := api.urls().BuildBlog()
	if err != nil {
		return nil, err
	}

	r, err := http.NewRequest(http.MethodGet, url, nil)
	if err != nil {
		return nil, err
	}

	resp, err := api.do(r)
	if err != nil {
		return nil, err
	}

	defer resp.Body.Close()

	body, err := ioutil.ReadAll(resp.Body)
	if err != nil {
		return nil, err
	}

	p := make([]*v1.Post, 0)
	if err = json.Unmarshal(body, &p); err != nil {
		return nil, err
	}

	return p, nil
}

func (api *blogAPI) CreatePost(create *v1.CreatePostRequest) (*v1.Post, error) {
	body, err := json.Marshal(&create)
	if err != nil {
		return nil, err
	}

	url, err := api.urls().BuildBlog()
	if err != nil {
		return nil, err
	}

	r, err := http.NewRequest(http.MethodPost, url, bytes.NewBuffer(body))
	if err != nil {
		return nil, err
	}

	resp, err := api.do(r)
	if err != nil {
		return nil, err
	}

	defer resp.Body.Close()
	body, err = ioutil.ReadAll(resp.Body)
	if err != nil {
		return nil, err
	}

	p := &v1.Post{}
	if err = json.Unmarshal(body, p); err != nil {
		return nil, err
	}

	return p, nil
}
<commit_msg>just use post data for create<commit_after>package client

import (
	"bytes"
	"encoding/json"
	"io/ioutil"
	"net/http"

	"github.com/danielkrainas/tinkersnest/api/v1"
)

type BlogAPI interface {
	SearchPosts() ([]*v1.Post, error)
	CreatePost(post *v1.Post) (*v1.Post, error)
}

type blogAPI struct {
	*Client
}

func (c *Client) Blog() BlogAPI {
	return &blogAPI{c}
}

func (api *blogAPI) SearchPosts() ([]*v1.Post, error) {
	url, err := api.urls().BuildBlog()
	if err != nil {
		return nil, err
	}

	r, err := http.NewRequest(http.MethodGet, url, nil)
	if err != nil {
		return nil, err
	}

	resp, err := api.do(r)
	if err != nil {
		return nil, err
	}

	defer resp.Body.Close()

	body, err := ioutil.ReadAll(resp.Body)
	if err != nil {
		return nil, err
	}

	p := make([]*v1.Post, 0)
	if err = json.Unmarshal(body, &p); err != nil {
		return nil, err
	}

	return p, nil
}

func (api *blogAPI) CreatePost(post *v1.Post) (*v1.Post, error) {
	body, err := json.Marshal(&post)
	if err != nil {
		return nil, err
	}

	url, err := api.urls().BuildBlog()
	if err != nil {
		return nil, err
	}

	r, err := http.NewRequest(http.MethodPost, url, bytes.NewBuffer(body))
	if err != nil {
		return nil, err
	}

	resp, err := api.do(r)
	if err != nil {
		return nil, err
	}

	defer resp.Body.Close()
	body, err = ioutil.ReadAll(resp.Body)
	if err != nil {
		return nil, err
	}

	p := &v1.Post{}
	if err = json.Unmarshal(body, p); err != nil {
		return nil, err
	}

	return p, nil
}
<|endoftext|>


<commit_before>package main

import (
	"fmt"
	"log"

	"github.com/grokify/gotilla/image/imageutil"
)

func main() {
	urls := []string{
		"https://example/img1.jpg",
		"https://example/img2.jpg",
		"https://example/img3.jpg",
		"https://example/img4.jpg"}

	img12, err := imageutil.MergeXSameYRead(urls[0], urls[1], true)
	if err != nil {
		log.Fatal(err)
	}

	outfile12 := "_img12.jpg"
	err = imageutil.WriteFileJPEG(outfile12, img12, imageutil.JPEGQualityMax)
	if err != nil {
		log.Fatal(err)
	}
	fmt.Printf("WROTE [%v]\n", outfile12)

	img34, err := imageutil.MergeXSameYRead(urls[2], urls[3], true)
	if err != nil {
		log.Fatal(err)
	}

	outfile34 := "_img34.jpg"
	err = imageutil.WriteFileJPEG(outfile34, img34, imageutil.JPEGQualityMax)
	if err != nil {
		log.Fatal(err)
	}
	fmt.Printf("WROTE [%v]\n", outfile34)

	img4 := imageutil.MergeYSameX(img12, img34, true)
	outfile := "_four.jpg"
	err = imageutil.WriteFileJPEG(outfile, img4, imageutil.JPEGQualityMax)
	if err != nil {
		log.Fatal(err)
	}
	fmt.Printf("WROTE [%v]\n", outfile)

	outfile4 := "_img4.jpg"
	img4a, err := imageutil.Merge4Read(urls[0], urls[1], urls[2], urls[3], true)
	err = imageutil.WriteFileJPEG(outfile4, img4a, imageutil.JPEGQualityMax)
	if err != nil {
		log.Fatal(err)
	}
	fmt.Printf("WROTE [%v]\n", outfile4)

	fmt.Println("DONE")
}
<commit_msg>enhance: imageutil: update `merge_four` example to use generic code<commit_after>package main

import (
	"fmt"
	"image"
	"log"

	"github.com/grokify/gotilla/image/imageutil"
)

func main() {
	urls := []string{
		"https://example.com/img1.jpg",
		"https://example.com/img2.jpg",
		"https://example.com/img3.jpg",
		"https://example.com/img4.jpg"}

	writeManual := true
	writeManualIntermediate := true

	img12, err := imageutil.MergeXSameYRead([]string{urls[0], urls[1]}, true)
	if err != nil {
		log.Fatal(err)
	}
	if writeManual && writeManualIntermediate {
		outfile12 := "_img12.jpg"
		err = imageutil.WriteFileJPEG(outfile12, img12, imageutil.JPEGQualityMax)
		if err != nil {
			log.Fatal(err)
		}
		fmt.Printf("WROTE [%v]\n", outfile12)
	}
	img34, err := imageutil.MergeXSameYRead([]string{urls[2], urls[3]}, true)
	if err != nil {
		log.Fatal(err)
	}

	if writeManual && writeManualIntermediate {
		outfile34 := "_img34.jpg"
		err = imageutil.WriteFileJPEG(outfile34, img34, imageutil.JPEGQualityMax)
		if err != nil {
			log.Fatal(err)
		}
		fmt.Printf("WROTE [%v]\n", outfile34)
	}
	if writeManual {
		img4 := imageutil.MergeYSameX([]image.Image{img12, img34}, true)
		outfile := "_img4_manual.jpg"
		err = imageutil.WriteFileJPEG(outfile, img4, imageutil.JPEGQualityMax)
		if err != nil {
			log.Fatal(err)
		}
		fmt.Printf("WROTE [%v]\n", outfile)
	}

	outfile4 := "_img4_auto.jpg"

	img4a, err := imageutil.MatrixMergeRead(
		[][]string{
			[]string{urls[0]},
			[]string{urls[0], urls[1]},
			[]string{urls[2], urls[3], urls[2]},
		},
		true, true)

	err = imageutil.WriteFileJPEG(outfile4, img4a, imageutil.JPEGQualityMax)
	if err != nil {
		log.Fatal(err)
	}
	fmt.Printf("WROTE [%v]\n", outfile4)

	fmt.Println("DONE")
}
<|endoftext|>


<commit_before>package main

import (
	"go/ast"
	"go/parser"
	"go/token"
	"strings"
)

type mImportDecl struct {
	Name string `json:"name"`
	Path string `json:"path"`
}

type mImportDeclArg struct {
	Name string `json:"name"`
	Path string `json:"path"`
	Add  bool   `json:"add"`
}

type mImports struct {
	Fn        string
	Src       string
	Toggle    []mImportDeclArg
	TabWidth  int
	TabIndent bool
}

func (m *mImports) Call() (interface{}, string) {
	lineRef := 0
	src := ""

	fset, af, err := parseAstFile(m.Fn, m.Src, parser.ImportsOnly|parser.ParseComments)
	if err == nil {
		// we neither return, nor attempt the whole source because it likely contains
		// syntax errors after the imports... as a result we need to tell the client
		// which part of the original source we edited so they are able to patch their
		// copy of the source...

		// find the last line number before we edit anything so the client can use it as
		// as reference as where to patch
		lineRef = fset.Position(af.Name.End()).Line
		if l := len(af.Decls); l > 0 {
			lineRef = fset.Position(af.Decls[l-1].End()).Line
		}

		// trim trailing comments so they don't make the line ref incorrect
		for i, c := range af.Comments {
			cp := fset.Position(c.Pos())
			if cp.Line > lineRef {
				af.Comments = af.Comments[:i]
				break
			}
		}

		af = imp(fset, af, m.Toggle)
		src, err = printSrc(fset, af, m.TabIndent, m.TabWidth)
	}

	res := M{
		"src":     src,
		"lineRef": lineRef,
	}
	return res, errStr(err)
}

func init() {
	registry.Register("imports", func(_ *Broker) Caller {
		return &mImports{
			Toggle:    []mImportDeclArg{},
			TabWidth:  8,
			TabIndent: true,
		}
	})
}

func unquote(s string) string {
	return strings.Trim(s, "\"`")
}

func quote(s string) string {
	return `"` + unquote(s) + `"`
}

func imp(fset *token.FileSet, af *ast.File, toggle []mImportDeclArg) *ast.File {
	add := map[mImportDecl]bool{}
	del := map[mImportDecl]bool{}
	for _, sda := range toggle {
		sd := mImportDecl{
			Path: sda.Path,
			Name: sda.Name,
		}
		if sda.Add {
			add[sd] = true
		} else {
			del[sd] = true
		}
	}

	var firstDecl *ast.GenDecl
	imports := map[mImportDecl]bool{}
	for _, decl := range af.Decls {
		if gdecl, ok := decl.(*ast.GenDecl); ok && len(gdecl.Specs) > 0 {
			hasC := false
			sj := 0
			for _, spec := range gdecl.Specs {
				if ispec, ok := spec.(*ast.ImportSpec); ok {
					sd := mImportDecl{
						Path: unquote(ispec.Path.Value),
					}
					if ispec.Name != nil {
						sd.Name = ispec.Name.String()
					}

					if sd.Path == "C" {
						hasC = true
					} else if del[sd] {
						if sj > 0 {
							if lspec, ok := gdecl.Specs[sj-1].(*ast.ImportSpec); ok {
								lspec.EndPos = ispec.Pos()
							}
						}
						continue
					} else {
						imports[sd] = true
					}
				}

				gdecl.Specs[sj] = spec
				sj += 1
			}
			gdecl.Specs = gdecl.Specs[:sj]

			if !hasC && firstDecl == nil {
				firstDecl = gdecl
			}
		}
	}

	if len(add) > 0 {
		if firstDecl == nil {
			firstDecl = &ast.GenDecl{
				Tok:    token.IMPORT,
				Lparen: 1,
			}
			af.Decls = append(af.Decls, firstDecl)
		} else if firstDecl.Lparen == token.NoPos {
			firstDecl.Lparen = 1
		}

		addSpecs := make([]ast.Spec, 0, len(firstDecl.Specs)+len(add))
		for sd, _ := range add {
			if !imports[sd] {
				ispec := &ast.ImportSpec{
					Path: &ast.BasicLit{
						Value: quote(sd.Path),
						Kind:  token.STRING,
					},
				}
				if sd.Name != "" {
					ispec.Name = &ast.Ident{
						Name: sd.Name,
					}
				}
				addSpecs = append(addSpecs, ispec)
				imports[sd] = true
			}
		}
		firstDecl.Specs = append(addSpecs, firstDecl.Specs...)
	}

	dj := 0
	for _, decl := range af.Decls {
		if gdecl, ok := decl.(*ast.GenDecl); ok {
			if len(gdecl.Specs) == 0 {
				continue
			}
		}
		af.Decls[dj] = decl
		dj += 1
	}
	af.Decls = af.Decls[:dj]

	return af
}
<commit_msg>* add autoInstall support<commit_after>package main

import (
	"go/ast"
	"go/parser"
	"go/token"
	"strings"
)

type mImportDecl struct {
	Name string `json:"name"`
	Path string `json:"path"`
}

type mImportDeclArg struct {
	Name string `json:"name"`
	Path string `json:"path"`
	Add  bool   `json:"add"`
}

type mImports struct {
	Fn        string
	Src       string
	Toggle    []mImportDeclArg
	TabWidth  int
	TabIndent bool
	Env       map[string]string
	Autoinst  bool
}

func (m *mImports) Call() (interface{}, string) {
	lineRef := 0
	src := ""

	fset, af, err := parseAstFile(m.Fn, m.Src, parser.ImportsOnly|parser.ParseComments)
	if err == nil {
		// we neither return, nor attempt the whole source because it likely contains
		// syntax errors after the imports... as a result we need to tell the client
		// which part of the original source we edited so they are able to patch their
		// copy of the source...

		// find the last line number before we edit anything so the client can use it as
		// as reference as where to patch
		lineRef = fset.Position(af.Name.End()).Line
		if l := len(af.Decls); l > 0 {
			lineRef = fset.Position(af.Decls[l-1].End()).Line
		}

		// trim trailing comments so they don't make the line ref incorrect
		for i, c := range af.Comments {
			cp := fset.Position(c.Pos())
			if cp.Line > lineRef {
				af.Comments = af.Comments[:i]
				break
			}
		}

		af = imp(fset, af, m.Toggle)
		src, err = printSrc(fset, af, m.TabIndent, m.TabWidth)
	}

	if m.Autoinst {
		autoInstall(AutoInstOptions{
			Env:         m.Env,
			ImportPaths: fileImportPaths(af),
		})
	}

	res := M{
		"src":     src,
		"lineRef": lineRef,
	}
	return res, errStr(err)
}

func init() {
	registry.Register("imports", func(_ *Broker) Caller {
		return &mImports{
			Toggle:    []mImportDeclArg{},
			TabWidth:  8,
			TabIndent: true,
		}
	})
}

func unquote(s string) string {
	return strings.Trim(s, "\"`")
}

func quote(s string) string {
	return `"` + unquote(s) + `"`
}

func imp(fset *token.FileSet, af *ast.File, toggle []mImportDeclArg) *ast.File {
	add := map[mImportDecl]bool{}
	del := map[mImportDecl]bool{}
	for _, sda := range toggle {
		sd := mImportDecl{
			Path: sda.Path,
			Name: sda.Name,
		}
		if sda.Add {
			add[sd] = true
		} else {
			del[sd] = true
		}
	}

	var firstDecl *ast.GenDecl
	imports := map[mImportDecl]bool{}
	for _, decl := range af.Decls {
		if gdecl, ok := decl.(*ast.GenDecl); ok && len(gdecl.Specs) > 0 {
			hasC := false
			sj := 0
			for _, spec := range gdecl.Specs {
				if ispec, ok := spec.(*ast.ImportSpec); ok {
					sd := mImportDecl{
						Path: unquote(ispec.Path.Value),
					}
					if ispec.Name != nil {
						sd.Name = ispec.Name.String()
					}

					if sd.Path == "C" {
						hasC = true
					} else if del[sd] {
						if sj > 0 {
							if lspec, ok := gdecl.Specs[sj-1].(*ast.ImportSpec); ok {
								lspec.EndPos = ispec.Pos()
							}
						}
						continue
					} else {
						imports[sd] = true
					}
				}

				gdecl.Specs[sj] = spec
				sj += 1
			}
			gdecl.Specs = gdecl.Specs[:sj]

			if !hasC && firstDecl == nil {
				firstDecl = gdecl
			}
		}
	}

	if len(add) > 0 {
		if firstDecl == nil {
			firstDecl = &ast.GenDecl{
				Tok:    token.IMPORT,
				Lparen: 1,
			}
			af.Decls = append(af.Decls, firstDecl)
		} else if firstDecl.Lparen == token.NoPos {
			firstDecl.Lparen = 1
		}

		addSpecs := make([]ast.Spec, 0, len(firstDecl.Specs)+len(add))
		for sd, _ := range add {
			if !imports[sd] {
				ispec := &ast.ImportSpec{
					Path: &ast.BasicLit{
						Value: quote(sd.Path),
						Kind:  token.STRING,
					},
				}
				if sd.Name != "" {
					ispec.Name = &ast.Ident{
						Name: sd.Name,
					}
				}
				addSpecs = append(addSpecs, ispec)
				imports[sd] = true
			}
		}
		firstDecl.Specs = append(addSpecs, firstDecl.Specs...)
	}

	dj := 0
	for _, decl := range af.Decls {
		if gdecl, ok := decl.(*ast.GenDecl); ok {
			if len(gdecl.Specs) == 0 {
				continue
			}
		}
		af.Decls[dj] = decl
		dj += 1
	}
	af.Decls = af.Decls[:dj]

	return af
}
<|endoftext|>


<commit_before>package protocol

import (
	"fmt"
	"io"
	"sync"

	context "github.com/jbenet/go-ipfs/Godeps/_workspace/src/code.google.com/p/go.net/context"

	inet "github.com/jbenet/go-ipfs/p2p/net"
	eventlog "github.com/jbenet/go-ipfs/thirdparty/eventlog"
	lgbl "github.com/jbenet/go-ipfs/util/eventlog/loggables"
)

var log = eventlog.Logger("net/mux")

type streamHandlerMap map[ID]inet.StreamHandler

// Mux provides simple stream multixplexing.
// It helps you precisely when:
//  * You have many streams
//  * You have function handlers
//
// It contains the handlers for each protocol accepted.
// It dispatches handlers for streams opened by remote peers.
type Mux struct {
	// defaultHandler handles unknown protocols. Callers modify at your own risk.
	defaultHandler inet.StreamHandler

	lock     sync.RWMutex
	handlers streamHandlerMap
}

func NewMux() *Mux {
	return &Mux{
		handlers: streamHandlerMap{},
	}
}

// Protocols returns the list of protocols this muxer has handlers for
func (m *Mux) Protocols() []ID {
	m.lock.RLock()
	l := make([]ID, 0, len(m.handlers))
	for p := range m.handlers {
		l = append(l, p)
	}
	m.lock.RUnlock()
	return l
}

// readHeader reads the stream and returns the next Handler function
// according to the muxer encoding.
func (m *Mux) readHeader(s io.Reader) (ID, inet.StreamHandler, error) {
	// log.Error("ReadProtocolHeader")
	p, err := ReadHeader(s)
	if err != nil {
		return "", nil, err
	}

	// log.Debug("readHeader got:", p)
	m.lock.RLock()
	h, found := m.handlers[p]
	m.lock.RUnlock()

	switch {
	case !found && m.defaultHandler != nil:
		return p, m.defaultHandler, nil
	case !found && m.defaultHandler == nil:
		return p, nil, fmt.Errorf("%s no handler with name: %s (%d)", m, p, len(p))
	default:
		return p, h, nil
	}
}

// String returns the muxer's printing representation
func (m *Mux) String() string {
	m.lock.RLock()
	defer m.lock.RUnlock()
	return fmt.Sprintf("<Muxer %p %d>", m, len(m.handlers))
}

func (m *Mux) SetDefaultHandler(h inet.StreamHandler) {
	m.lock.Lock()
	m.defaultHandler = h
	m.lock.Unlock()
}

// SetHandler sets the protocol handler on the Network's Muxer.
// This operation is threadsafe.
func (m *Mux) SetHandler(p ID, h inet.StreamHandler) {
	log.Debugf("%s setting handler for protocol: %s (%d)", m, p, len(p))
	m.lock.Lock()
	m.handlers[p] = h
	m.lock.Unlock()
}

// Handle reads the next name off the Stream, and calls a handler function
// This is done in its own goroutine, to avoid blocking the caller.
func (m *Mux) Handle(s inet.Stream) {
	go m.HandleSync(s)
}

// HandleSync reads the next name off the Stream, and calls a handler function
// This is done synchronously. The handler function will return before
// HandleSync returns.
func (m *Mux) HandleSync(s inet.Stream) {
	ctx := context.Background()

	name, handler, err := m.readHeader(s)
	if err != nil {
		err = fmt.Errorf("protocol mux error: %s", err)
		log.Error(err)
		log.Event(ctx, "muxError", lgbl.Error(err))
		s.Close()
		return
	}

	log.Infof("muxer handle protocol: %s", name)
	handler(s)
}

// ReadLengthPrefix reads the name from Reader with a length-byte-prefix.
func ReadLengthPrefix(r io.Reader) (string, error) {
	// c-string identifier
	// the first byte is our length
	l := make([]byte, 1)
	if _, err := io.ReadFull(r, l); err != nil {
		return "", err
	}
	length := int(l[0])

	// the next are our identifier
	name := make([]byte, length)
	if _, err := io.ReadFull(r, name); err != nil {
		return "", err
	}

	return string(name), nil
}
<commit_msg>p2p/proto/mux: make log more useful<commit_after>package protocol

import (
	"fmt"
	"io"
	"sync"

	context "github.com/jbenet/go-ipfs/Godeps/_workspace/src/code.google.com/p/go.net/context"

	inet "github.com/jbenet/go-ipfs/p2p/net"
	eventlog "github.com/jbenet/go-ipfs/thirdparty/eventlog"
	lgbl "github.com/jbenet/go-ipfs/util/eventlog/loggables"
)

var log = eventlog.Logger("net/mux")

type streamHandlerMap map[ID]inet.StreamHandler

// Mux provides simple stream multixplexing.
// It helps you precisely when:
//  * You have many streams
//  * You have function handlers
//
// It contains the handlers for each protocol accepted.
// It dispatches handlers for streams opened by remote peers.
type Mux struct {
	// defaultHandler handles unknown protocols. Callers modify at your own risk.
	defaultHandler inet.StreamHandler

	lock     sync.RWMutex
	handlers streamHandlerMap
}

func NewMux() *Mux {
	return &Mux{
		handlers: streamHandlerMap{},
	}
}

// Protocols returns the list of protocols this muxer has handlers for
func (m *Mux) Protocols() []ID {
	m.lock.RLock()
	l := make([]ID, 0, len(m.handlers))
	for p := range m.handlers {
		l = append(l, p)
	}
	m.lock.RUnlock()
	return l
}

// readHeader reads the stream and returns the next Handler function
// according to the muxer encoding.
func (m *Mux) readHeader(s io.Reader) (ID, inet.StreamHandler, error) {
	// log.Error("ReadProtocolHeader")
	p, err := ReadHeader(s)
	if err != nil {
		return "", nil, err
	}

	// log.Debug("readHeader got:", p)
	m.lock.RLock()
	h, found := m.handlers[p]
	m.lock.RUnlock()

	switch {
	case !found && m.defaultHandler != nil:
		return p, m.defaultHandler, nil
	case !found && m.defaultHandler == nil:
		return p, nil, fmt.Errorf("%s no handler with name: %s (%d)", m, p, len(p))
	default:
		return p, h, nil
	}
}

// String returns the muxer's printing representation
func (m *Mux) String() string {
	m.lock.RLock()
	defer m.lock.RUnlock()
	return fmt.Sprintf("<Muxer %p %d>", m, len(m.handlers))
}

func (m *Mux) SetDefaultHandler(h inet.StreamHandler) {
	m.lock.Lock()
	m.defaultHandler = h
	m.lock.Unlock()
}

// SetHandler sets the protocol handler on the Network's Muxer.
// This operation is threadsafe.
func (m *Mux) SetHandler(p ID, h inet.StreamHandler) {
	log.Debugf("%s setting handler for protocol: %s (%d)", m, p, len(p))
	m.lock.Lock()
	m.handlers[p] = h
	m.lock.Unlock()
}

// Handle reads the next name off the Stream, and calls a handler function
// This is done in its own goroutine, to avoid blocking the caller.
func (m *Mux) Handle(s inet.Stream) {
	go m.HandleSync(s)
}

// HandleSync reads the next name off the Stream, and calls a handler function
// This is done synchronously. The handler function will return before
// HandleSync returns.
func (m *Mux) HandleSync(s inet.Stream) {
	ctx := context.Background()

	name, handler, err := m.readHeader(s)
	if err != nil {
		err = fmt.Errorf("protocol mux error: %s", err)
		log.Error(err)
		log.Event(ctx, "muxError", lgbl.Error(err))
		s.Close()
		return
	}

	log.Infof("muxer handle protocol %s: %s", s.Conn().RemotePeer(), name)
	handler(s)
}

// ReadLengthPrefix reads the name from Reader with a length-byte-prefix.
func ReadLengthPrefix(r io.Reader) (string, error) {
	// c-string identifier
	// the first byte is our length
	l := make([]byte, 1)
	if _, err := io.ReadFull(r, l); err != nil {
		return "", err
	}
	length := int(l[0])

	// the next are our identifier
	name := make([]byte, length)
	if _, err := io.ReadFull(r, name); err != nil {
		return "", err
	}

	return string(name), nil
}
<|endoftext|>