<commit_before>package notification

import (
	"encoding/json"
	"errors"
	"github.com/koding/logging"
	"github.com/koding/rabbitmq"
	"github.com/koding/worker"
	"github.com/streadway/amqp"
	mongomodels "koding/db/models"
	"koding/db/mongodb/modelhelper"
	"labix.org/v2/mgo"
	"socialapi/models"
	"socialapi/workers/cache"
)

type Action func(*NotificationWorkerController, []byte) error

type NotificationWorkerController struct {
	routes          map[string]Action
	log             logging.Logger
	rmqConn         *amqp.Connection
	notifierRmqConn *amqp.Connection
	cacheEnabled    bool
}

type NotificationEvent struct {
	RoutingKey string              `json:"routingKey"`
	Event      string              `json:"event"`
	Content    NotificationContent `json:"contents"`
}

type NotificationContent struct {
	TypeConstant string `json:"type"`
	TargetId     int64  `json:"targetId"`
	ActorId      int64  `json:"actorId"`
}

func (n *NotificationWorkerController) DefaultErrHandler(delivery amqp.Delivery, err error) {
	n.log.Error("an error occured: %s", err)

	delivery.Ack(false)
}

func NewNotificationWorkerController(rmq *rabbitmq.RabbitMQ, log logging.Logger, cacheEnabled bool) (*NotificationWorkerController, error) {
	rmqConn, err := rmq.Connect("NewNotificationWorkerController")
	if err != nil {
		return nil, err
	}

	notifierRmqConn, err := rmq.Connect("NotifierWorkerController")
	if err != nil {
		return nil, err
	}

	nwc := &NotificationWorkerController{
		log:             log,
		rmqConn:         rmqConn.Conn(),
		notifierRmqConn: notifierRmqConn.Conn(),
		cacheEnabled:    cacheEnabled,
	}

	routes := map[string]Action{
		"api.message_reply_created": (*NotificationWorkerController).CreateReplyNotification,
		"api.interaction_created":   (*NotificationWorkerController).CreateInteractionNotification,
		"api.notification_created":  (*NotificationWorkerController).NotifyUser,
		"api.notification_updated":  (*NotificationWorkerController).NotifyUser,
	}

	nwc.routes = routes

	models.Log = log

	return nwc, nil
}

// copy/paste
func (n *NotificationWorkerController) HandleEvent(event string, data []byte) error {
	n.log.Debug("New Event Received %s", event)
	handler, ok := n.routes[event]
	if !ok {
		return worker.HandlerNotFoundErr
	}

	return handler(n, data)
}

func (n *NotificationWorkerController) CreateReplyNotification(data []byte) error {
	mr, err := mapMessageToMessageReply(data)
	if err != nil {
		return err
	}

	// fetch replier
	cm := models.NewChannelMessage()
	if err := cm.ById(mr.ReplyId); err != nil {
		return err
	}

	rn := models.NewReplyNotification()
	rn.TargetId = mr.MessageId
	rn.NotifierId = cm.AccountId

	nc, err := models.CreateNotificationContent(rn)
	if err != nil {
		return err
	}

	cm = models.NewChannelMessage()
	// notify message owner
	if err = cm.ById(mr.MessageId); err != nil {
		return err
	}

	n.createNotification(nc.Id, cm.AccountId)

	notifiedUsers, err := rn.GetNotifiedUsers(nc.Id)
	if err != nil {
		return err
	}

	notifierSubscribed := false
	for _, recipient := range notifiedUsers {
		if recipient == rn.NotifierId {
			notifierSubscribed = true
		}
		n.createNotification(nc.Id, recipient)
	}

	// if not subcribed, create a notification for new subscriber
	if !notifierSubscribed {
		n.createNotification(nc.Id, rn.NotifierId)
	}

	return nil
}

func (n *NotificationWorkerController) createNotification(contentId, notifierId int64) {
	notification := models.NewNotification()
	notification.NotificationContentId = contentId
	notification.AccountId = notifierId
	if err := notification.Create(); err != nil {
		n.log.Error("An error occurred while notifying user %d: %s", notifierId, err.Error())
	}
}

func (n *NotificationWorkerController) CreateInteractionNotification(data []byte) error {
	i, err := mapMessageToInteraction(data)
	if err != nil {
		return err
	}

	cm := models.NewChannelMessage()
	if err := cm.ById(i.MessageId); err != nil {
		return err
	}

	// user likes her own message, so we bypass notification
	if cm.AccountId == i.AccountId {
		return nil
	}

	in := models.NewInteractionNotification(i.TypeConstant)
	in.TargetId = i.MessageId
	in.NotifierId = i.AccountId
	nc, err := models.CreateNotificationContent(in)
	if err != nil {
		return err
	}

	recipients, err := in.GetNotifiedUsers(nc.Id)
	if err != nil {
		return err
	}

	if len(recipients) == 0 {
		return errors.New("interaction notifiee not found")
	}

	notification := models.NewNotification()
	notification.NotificationContentId = nc.Id
	notification.AccountId = recipients[0]
	if err = notification.Create(); err != nil {
		n.log.Error("An error occurred while notifying user %d: %s", recipients[0], err.Error())
	}

	return nil
}

func (n *NotificationWorkerController) NotifyUser(data []byte) error {
	channel, err := n.notifierRmqConn.Channel()
	if err != nil {
		return errors.New("channel connection error")
	}
	defer channel.Close()

	notification, err := mapMessageToNotification(data)
	if err != nil {
		return err
	}

	// fetch notification content and get event type
	nc, err := notification.FetchContent()
	if err != nil {
		return err
	}

	nt, err := models.CreateNotificationContentType(nc.TypeConstant)
	if err != nil {
		return err
	}

	nt.SetTargetId(nc.TargetId)
	// ac, err := nt.FetchActors()
	// if err != nil {
	// 	return err
	// }

	// var actorId int64
	// if len(ac.LatestActors) > 0 {
	// 	actorId = ac.LatestActors[0]
	// }

	go func() {
		if n.cacheEnabled {
			notificationCache := cache.NewNotificationCache()
			if err := notificationCache.UpdateCache(notification, nc); err != nil {
				n.log.Error("an error occurred %s", err)
			}
		}
	}()

	accountId := notification.AccountId
	oldAccount, err := fetchNotifierOldAccount(accountId)
	if err != nil {
		return err
	}

	// fetch user profile name from bongo as routing key
	ne := &NotificationEvent{}
	ne.Event = nc.GetEventType()
	ne.Content = NotificationContent{
		// ActorId:      actorId,
		TargetId:     nc.TargetId,
		TypeConstant: nc.TypeConstant,
	}

	notificationMessage, err := json.Marshal(ne)
	if err != nil {
		return err
	}

	routingKey := oldAccount.Profile.Nickname

	return channel.Publish(
		"notification",
		routingKey,
		false,
		false,
		amqp.Publishing{Body: notificationMessage},
	)
}

// fetchNotifierOldAccount fetches mongo account of a given new account id.
// this function must be used under another file for further use
func fetchNotifierOldAccount(accountId int64) (*mongomodels.Account, error) {
	newAccount := models.NewAccount()
	if err := newAccount.ById(accountId); err != nil {
		return nil, err
	}

	account, err := modelhelper.GetAccountById(newAccount.OldId)
	if err != nil {
		if err == mgo.ErrNotFound {
			return nil, errors.New("old account not found")
		}

		return nil, err
	}

	return account, nil
}

// copy/pasted from realtime package
func mapMessageToMessageReply(data []byte) (*models.MessageReply, error) {
	mr := models.NewMessageReply()
	if err := json.Unmarshal(data, mr); err != nil {
		return nil, err
	}

	return mr, nil
}

// copy/pasted from realtime package
func mapMessageToInteraction(data []byte) (*models.Interaction, error) {
	i := models.NewInteraction()
	if err := json.Unmarshal(data, i); err != nil {
		return nil, err
	}

	return i, nil
}

func mapMessageToNotification(data []byte) (*models.Notification, error) {
	n := models.NewNotification()
	if err := json.Unmarshal(data, n); err != nil {
		return nil, err
	}

	return n, nil
}
<commit_msg>Social: SubscribedAt time is injected in new subscription creation<commit_after>package notification

import (
	"encoding/json"
	"errors"
	// "fmt"
	"github.com/koding/logging"
	"github.com/koding/rabbitmq"
	"github.com/koding/worker"
	"github.com/streadway/amqp"
	mongomodels "koding/db/models"
	"koding/db/mongodb/modelhelper"
	"labix.org/v2/mgo"
	"socialapi/models"
	"socialapi/workers/cache"
	"time"
)

type Action func(*NotificationWorkerController, []byte) error

type NotificationWorkerController struct {
	routes          map[string]Action
	log             logging.Logger
	rmqConn         *amqp.Connection
	notifierRmqConn *amqp.Connection
	cacheEnabled    bool
}

type NotificationEvent struct {
	RoutingKey string              `json:"routingKey"`
	Event      string              `json:"event"`
	Content    NotificationContent `json:"contents"`
}

type NotificationContent struct {
	TypeConstant string `json:"type"`
	TargetId     int64  `json:"targetId"`
	ActorId      int64  `json:"actorId"`
}

func (n *NotificationWorkerController) DefaultErrHandler(delivery amqp.Delivery, err error) {
	n.log.Error("an error occured: %s", err)

	delivery.Ack(false)
}

func NewNotificationWorkerController(rmq *rabbitmq.RabbitMQ, log logging.Logger, cacheEnabled bool) (*NotificationWorkerController, error) {
	rmqConn, err := rmq.Connect("NewNotificationWorkerController")
	if err != nil {
		return nil, err
	}

	notifierRmqConn, err := rmq.Connect("NotifierWorkerController")
	if err != nil {
		return nil, err
	}

	nwc := &NotificationWorkerController{
		log:             log,
		rmqConn:         rmqConn.Conn(),
		notifierRmqConn: notifierRmqConn.Conn(),
		cacheEnabled:    cacheEnabled,
	}

	routes := map[string]Action{
		"api.message_reply_created": (*NotificationWorkerController).CreateReplyNotification,
		"api.interaction_created":   (*NotificationWorkerController).CreateInteractionNotification,
		"api.notification_created":  (*NotificationWorkerController).NotifyUser,
		"api.notification_updated":  (*NotificationWorkerController).NotifyUser,
	}

	nwc.routes = routes

	models.Log = log

	return nwc, nil
}

// copy/paste
func (n *NotificationWorkerController) HandleEvent(event string, data []byte) error {
	n.log.Debug("New Event Received %s", event)
	handler, ok := n.routes[event]
	if !ok {
		return worker.HandlerNotFoundErr
	}

	return handler(n, data)
}

func (n *NotificationWorkerController) CreateReplyNotification(data []byte) error {
	mr, err := mapMessageToMessageReply(data)
	if err != nil {
		return err
	}

	// fetch replier
	cm := models.NewChannelMessage()
	if err := cm.ById(mr.ReplyId); err != nil {
		return err
	}

	rn := models.NewReplyNotification()
	rn.TargetId = mr.MessageId
	rn.NotifierId = cm.AccountId
	subscribedAt := time.Now()

	nc, err := models.CreateNotificationContent(rn)
	if err != nil {
		return err
	}

	cm = models.NewChannelMessage()
	// notify message owner
	if err = cm.ById(mr.MessageId); err != nil {
		return err
	}

	// if it is not notifier's own message then add owner to subscribers
	if cm.AccountId != rn.NotifierId {
		n.createNotification(nc.Id, cm.AccountId, subscribedAt)
	}

	notifiedUsers, err := rn.GetNotifiedUsers(nc.Id)
	if err != nil {
		return err
	}

	notifierSubscribed := false
	for _, recipient := range notifiedUsers {
		if recipient == rn.NotifierId {
			notifierSubscribed = true
		}
		n.createNotification(nc.Id, recipient, subscribedAt)
	}

	// if not subcribed, subscribe the actor to message
	if !notifierSubscribed {
		n.createNotification(nc.Id, rn.NotifierId, subscribedAt)
	}

	return nil
}

func (n *NotificationWorkerController) createNotification(contentId, notifierId int64, subscribedAt time.Time) {
	notification := models.NewNotification()
	notification.NotificationContentId = contentId
	notification.AccountId = notifierId
	notification.SubscribedAt = subscribedAt
	if err := notification.Create(); err != nil {
		n.log.Error("An error occurred while notifying user %d: %s", notifierId, err.Error())
	}
}

func (n *NotificationWorkerController) CreateInteractionNotification(data []byte) error {
	i, err := mapMessageToInteraction(data)
	if err != nil {
		return err
	}

	cm := models.NewChannelMessage()
	if err := cm.ById(i.MessageId); err != nil {
		return err
	}

	// user likes her own message, so we bypass notification
	if cm.AccountId == i.AccountId {
		return nil
	}

	in := models.NewInteractionNotification(i.TypeConstant)
	in.TargetId = i.MessageId
	in.NotifierId = i.AccountId
	nc, err := models.CreateNotificationContent(in)
	if err != nil {
		return err
	}

	notification := models.NewNotification()
	notification.NotificationContentId = nc.Id
	notification.AccountId = cm.AccountId // notify message owner
	notification.ActivatedAt = time.Now() // enables notification immediately
	if err = notification.Create(); err != nil {
		n.log.Error("An error occurred while notifying user %d: %s", cm.AccountId, err.Error())
	}

	return nil
}

// TODO how this is handled, it looks like it needs refactoring
func (n *NotificationWorkerController) NotifyUser(data []byte) error {
	channel, err := n.notifierRmqConn.Channel()
	if err != nil {
		return errors.New("channel connection error")
	}
	defer channel.Close()

	notification, err := mapMessageToNotification(data)
	if err != nil {
		return err
	}

	// fetch notification content and get event type
	nc, err := notification.FetchContent()
	if err != nil {
		return err
	}

	nI, err := models.CreateNotificationContentType(nc.TypeConstant)
	if err != nil {
		return err
	}

	nI.SetTargetId(nc.TargetId)
	// ac, err := nt.FetchActors()
	// if err != nil {
	// 	return err
	// }

	// var actorId int64
	// if len(ac.LatestActors) > 0 {
	// 	actorId = ac.LatestActors[0]
	// }

	go func() {
		if n.cacheEnabled {
			notificationCache := cache.NewNotificationCache()
			if err := notificationCache.UpdateCache(notification, nc); err != nil {
				n.log.Error("an error occurred %s", err)
			}
		}
	}()

	accountId := notification.AccountId
	oldAccount, err := fetchNotifierOldAccount(accountId)
	if err != nil {
		return err
	}

	// fetch user profile name from bongo as routing key
	ne := &NotificationEvent{}
	ne.Event = nc.GetEventType()
	ne.Content = NotificationContent{
		// ActorId:      actorId,
		TargetId:     nc.TargetId,
		TypeConstant: nc.TypeConstant,
	}

	notificationMessage, err := json.Marshal(ne)
	if err != nil {
		return err
	}

	routingKey := oldAccount.Profile.Nickname

	return channel.Publish(
		"notification",
		routingKey,
		false,
		false,
		amqp.Publishing{Body: notificationMessage},
	)
}

// fetchNotifierOldAccount fetches mongo account of a given new account id.
// this function must be used under another file for further use
func fetchNotifierOldAccount(accountId int64) (*mongomodels.Account, error) {
	newAccount := models.NewAccount()
	if err := newAccount.ById(accountId); err != nil {
		return nil, err
	}

	account, err := modelhelper.GetAccountById(newAccount.OldId)
	if err != nil {
		if err == mgo.ErrNotFound {
			return nil, errors.New("old account not found")
		}

		return nil, err
	}

	return account, nil
}

// copy/pasted from realtime package
func mapMessageToMessageReply(data []byte) (*models.MessageReply, error) {
	mr := models.NewMessageReply()
	if err := json.Unmarshal(data, mr); err != nil {
		return nil, err
	}

	return mr, nil
}

// copy/pasted from realtime package
func mapMessageToInteraction(data []byte) (*models.Interaction, error) {
	i := models.NewInteraction()
	if err := json.Unmarshal(data, i); err != nil {
		return nil, err
	}

	return i, nil
}

func mapMessageToNotification(data []byte) (*models.Notification, error) {
	n := models.NewNotification()
	if err := json.Unmarshal(data, n); err != nil {
		return nil, err
	}

	return n, nil
}
<|endoftext|>


<commit_before>package server

import "net/http"

type getMappingsResponse struct {
	Mappings []mapping `json:"mappings"`
}

type mapping struct {
	Measurement string `json:"measurement"` // The measurement where data for this mapping is found
	Name        string `json:"name"`        // The application name which will be assigned to the corresponding measurement
}

// GetMappings returns the known mappings of measurements to applications
func (h *Service) GetMappings(w http.ResponseWriter, r *http.Request) {
	ctx := r.Context()
	layouts, err := h.LayoutStore.All(ctx)
	if err != nil {
		Error(w, http.StatusInternalServerError, "Error loading layouts")
	}

	mp := getMappingsResponse{
		Mappings: []mapping{},
	}

	seen := make(map[string]bool)

	for _, layout := range layouts {
		if seen[layout.Measurement+layout.ID] {
			continue
		}
		mp.Mappings = append(mp.Mappings, mapping{layout.Measurement, layout.ID})
		seen[layout.Measurement+layout.ID] = true
	}

	encodeJSON(w, http.StatusOK, mp, h.Logger)
}
<commit_msg>Make /mappings use App Name instead of ID<commit_after>package server

import "net/http"

type getMappingsResponse struct {
	Mappings []mapping `json:"mappings"`
}

type mapping struct {
	Measurement string `json:"measurement"` // The measurement where data for this mapping is found
	Name        string `json:"name"`        // The application name which will be assigned to the corresponding measurement
}

// GetMappings returns the known mappings of measurements to applications
func (h *Service) GetMappings(w http.ResponseWriter, r *http.Request) {
	ctx := r.Context()
	layouts, err := h.LayoutStore.All(ctx)
	if err != nil {
		Error(w, http.StatusInternalServerError, "Error loading layouts")
	}

	mp := getMappingsResponse{
		Mappings: []mapping{},
	}

	seen := make(map[string]bool)

	for _, layout := range layouts {
		if seen[layout.Measurement+layout.ID] {
			continue
		}
		mp.Mappings = append(mp.Mappings, mapping{layout.Measurement, layout.Application})
		seen[layout.Measurement+layout.ID] = true
	}

	encodeJSON(w, http.StatusOK, mp, h.Logger)
}
<|endoftext|>


<commit_before>package server

import (
	"errors"
	"github.com/julienschmidt/httprouter"
	. "gopkg.in/check.v1"
	"io/ioutil"
	"log"
	"net/http"
	"net/http/httptest"
	"os"
)

type WebServer struct{}

var _ = Suite(&WebServer{})

func (*WebServer) TestFrontpageRedirect(c *C) {
	config = serverConfigs{}
	config.Frontpage = "./test/frontpage.html"
	server := httptest.NewServer(http.HandlerFunc(redirectToDefault))
	defer server.Close()
	res, err := http.Get(server.URL)
	c.Assert(err, IsNil)
	frontpage, err := ioutil.ReadAll(res.Body)
	c.Assert(err, IsNil)
	c.Assert(res.Body.Close(), IsNil)
	c.Assert(string(frontpage), Equals, "<!doctype html>\n")
}

func (*WebServer) TestDefaultBoardRedirect(c *C) {
	config = serverConfigs{}
	config.Boards.Default = "a"
	rec := runHandler(c, redirectToDefault)
	c.Assert(rec.Code, Equals, 302)
	c.Assert(rec.Header().Get("Location"), Equals, "/a/")
}

func runHandler(c *C, h http.HandlerFunc) *httptest.ResponseRecorder {
	req := newRequest(c)
	rec := httptest.NewRecorder()
	h(rec, req)
	return rec
}

func newRequest(c *C) *http.Request {
	req, err := http.NewRequest("GET", "/", nil)
	c.Assert(err, IsNil)
	return req
}

func (*WebServer) TestConfigServing(c *C) {
	configHash = "foo"
	clientConfig = clientConfigs{}
	etag := "W/" + configHash
	rec := runHandler(c, serveConfigs)
	c.Assert(rec.Code, Equals, 200)
	c.Assert(rec.Body.String(), Equals, string(marshalJSON(clientConfig)))
	c.Assert(rec.Header().Get("ETag"), Equals, etag)

	// And with etag
	rec = httptest.NewRecorder()
	req := newRequest(c)
	req.Header.Set("If-None-Match", etag)
	serveConfigs(rec, req)
	c.Assert(rec.Code, Equals, 304)
}

func (*WebServer) TestEtagComparison(c *C) {
	req := newRequest(c)
	const etag = "foo"
	req.Header.Set("If-None-Match", etag)
	rec := httptest.NewRecorder()
	c.Assert(checkClientEtag(rec, req, etag), Equals, true)
}

func (*WebServer) TestNotFoundHandler(c *C) {
	webRoot = "./test"
	rec := runHandler(c, notFoundHandler)
	c.Assert(
		rec.Body.String(),
		Equals,
		"<!doctype html>404\n",
	)
	c.Assert(rec.Code, Equals, 404)
	headers := map[string]string{
		"Content-Type":           "text/html; charset=UTF-8",
		"X-Content-Type-Options": "nosniff",
	}
	assertHeaders(c, rec, headers)
}

func assertHeaders(c *C, rec *httptest.ResponseRecorder, h map[string]string) {
	for key, val := range h {
		c.Assert(rec.Header().Get(key), Equals, val)
	}
}

func (*WebServer) TestText404(c *C) {
	rec := runHandler(c, func(res http.ResponseWriter, _ *http.Request) {
		text404(res)
	})
	c.Assert(rec.Code, Equals, 404)
	c.Assert(rec.Body.String(), Equals, "404 Not found\n")
}

func (*WebServer) TestPanicHandler(c *C) {
	webRoot = "./test"
	err := errors.New("foo")

	// Prevent printing stack trace to terminal
	log.SetOutput(ioutil.Discard)
	rec := runHandler(c, func(res http.ResponseWriter, req *http.Request) {
		panicHandler(res, req, err)
	})
	log.SetOutput(os.Stdout)
	c.Assert(rec.Code, Equals, 500)
	c.Assert(rec.Body.String(), Equals, "<!doctype html>50x\n")
}

func (*WebServer) TestSetHeaders(c *C) {
	// HTML
	rec := httptest.NewRecorder()
	const etag = "foo"
	headers := map[string]string{
		"X-Frame-Options": "sameorigin",
		"Cache-Control":   "max-age=0, must-revalidate",
		"Expires":         "Fri, 01 Jan 1990 00:00:00 GMT",
		"ETag":            etag,
		"Content-Type":    "text/html; charset=UTF-8",
	}
	setHeaders(rec, etag, false)
	assertHeaders(c, rec, headers)

	// JSON
	headers["Content-Type"] = "application/json; charset=UTF-8"
	rec = httptest.NewRecorder()
	setHeaders(rec, etag, true)
	assertHeaders(c, rec, headers)
}

func (*WebServer) TestDetectLastN(c *C) {
	// No lastN query string
	req := customRequest(c, "/a/1")
	c.Assert(detectLastN(req), Equals, 0)

	// ?lastN value within bounds
	req = customRequest(c, "/a/1?lastN=100")
	c.Assert(detectLastN(req), Equals, 100)

	// ?lastN value beyond max
	req = customRequest(c, "/a/1?lastN=1000")
	c.Assert(detectLastN(req), Equals, 0)
}

func customRequest(c *C, url string) *http.Request {
	req, err := http.NewRequest("GET", url, nil)
	c.Assert(err, IsNil)
	return req
}

func (*WebServer) TestImageServer(c *C) {
	const (
		truncated         = "/src/tis life.gif"
		notFoundTruncated = "src/nobody here.gif"
	)
	imageWebRoot = "./test"
	path := imageWebRoot + truncated
	notFound := imageWebRoot + notFoundTruncated

	// Succesful first serve
	req := customRequest(c, path)
	rec := httptest.NewRecorder()
	params := httprouter.Params{
		httprouter.Param{
			Key:   "filepath",
			Value: truncated,
		},
	}
	serveImages(rec, req, params)
	buf, err := ioutil.ReadFile(path)
	c.Assert(err, IsNil)
	c.Assert(rec.Body.String(), Equals, string(buf))
	headers := map[string]string{
		"Cache-Control":   "max-age=30240000",
		"X-Frame-Options": "sameorigin",
		"ETag":            "0",
	}
	assertHeaders(c, rec, headers)

	// Fake etag validation
	req = customRequest(c, path)
	rec = httptest.NewRecorder()
	req.Header.Set("If-None-Match", "0")
	serveImages(rec, req, params)
	c.Assert(rec.Code, Equals, 304)

	// Non-existing file
	req = customRequest(c, notFound)
	rec = httptest.NewRecorder()
	params[0].Value = notFoundTruncated
	serveImages(rec, req, params)
	c.Assert(rec.Code, Equals, 404)
}

func (*WebServer) TestCompareEtag(c *C) {
	// Etag comparison
	rec := httptest.NewRecorder()
	req := newRequest(c)
	const etag = "foo"
	req.Header.Set("If-None-Match", etag)
	c.Assert(compareEtag(rec, req, etag, false), Equals, false)

	rec = httptest.NewRecorder()
	req = newRequest(c)
	headers := map[string]string{
		"ETag":          etag,
		"Content-Type":  "text/html; charset=UTF-8",
		"Cache-Control": "max-age=0, must-revalidate",
	}
	compareEtag(rec, req, etag, false)
	assertHeaders(c, rec, headers)
}

func (*WebServer) TestEtagStart(c *C) {
	c.Assert(etagStart(1), Equals, "W/1")
}

// func (*WebServer) TestDetectMobile(c *C) {
// 	req := newRequest(c)
// 	req.Header.Set(
// 		"User-Agent",
// 		"Mozilla/5.0 (Windows NT 6.1) AppleWebKit/537.36 (KHTML, like Gecko)"+
// 			" Chrome/41.0.2228.0 Safari/537.36",
// 	)
// 	c.Assert(detectMobile(req), Equals, false)
//
// 	req = newRequest(c)
// 	req.Header.Set(
// 		"User-Agent",
// 		"Mozilla/5.0 (Linux; Android 4.1.1; Galaxy Nexus Build/JRO03C)"+
// 			" AppleWebKit/535.19 (KHTML, like Gecko) Chrome/18.0.1025.166"+
// 			" Mobile Safari/535.19",
// 	)
// 	c.Assert(detectMobile(req), Equals, true)
// }
<commit_msg>server/web.go HTML serving tests<commit_after>package server

import (
	"errors"
	"github.com/julienschmidt/httprouter"
	. "gopkg.in/check.v1"
	"io/ioutil"
	"log"
	"net/http"
	"net/http/httptest"
	"os"
)

type WebServer struct{}

var _ = Suite(&WebServer{})

func (*WebServer) TestFrontpageRedirect(c *C) {
	config = serverConfigs{}
	config.Frontpage = "./test/frontpage.html"
	server := httptest.NewServer(http.HandlerFunc(redirectToDefault))
	defer server.Close()
	res, err := http.Get(server.URL)
	c.Assert(err, IsNil)
	frontpage, err := ioutil.ReadAll(res.Body)
	c.Assert(err, IsNil)
	c.Assert(res.Body.Close(), IsNil)
	c.Assert(string(frontpage), Equals, "<!doctype html>\n")
}

func (*WebServer) TestDefaultBoardRedirect(c *C) {
	config = serverConfigs{}
	config.Boards.Default = "a"
	rec := runHandler(c, redirectToDefault)
	c.Assert(rec.Code, Equals, 302)
	c.Assert(rec.Header().Get("Location"), Equals, "/a/")
}

func runHandler(c *C, h http.HandlerFunc) *httptest.ResponseRecorder {
	req := newRequest(c)
	rec := httptest.NewRecorder()
	h(rec, req)
	return rec
}

func newRequest(c *C) *http.Request {
	req, err := http.NewRequest("GET", "/", nil)
	c.Assert(err, IsNil)
	return req
}

func (*WebServer) TestConfigServing(c *C) {
	configHash = "foo"
	clientConfig = clientConfigs{}
	etag := "W/" + configHash
	rec := runHandler(c, serveConfigs)
	c.Assert(rec.Code, Equals, 200)
	c.Assert(rec.Body.String(), Equals, string(marshalJSON(clientConfig)))
	c.Assert(rec.Header().Get("ETag"), Equals, etag)

	// And with etag
	rec = httptest.NewRecorder()
	req := newRequest(c)
	req.Header.Set("If-None-Match", etag)
	serveConfigs(rec, req)
	c.Assert(rec.Code, Equals, 304)
}

func (*WebServer) TestEtagComparison(c *C) {
	req := newRequest(c)
	const etag = "foo"
	req.Header.Set("If-None-Match", etag)
	rec := httptest.NewRecorder()
	c.Assert(checkClientEtag(rec, req, etag), Equals, true)
}

func (*WebServer) TestNotFoundHandler(c *C) {
	webRoot = "./test"
	rec := runHandler(c, notFoundHandler)
	c.Assert(
		rec.Body.String(),
		Equals,
		"<!doctype html>404\n",
	)
	c.Assert(rec.Code, Equals, 404)
	headers := map[string]string{
		"Content-Type":           "text/html; charset=UTF-8",
		"X-Content-Type-Options": "nosniff",
	}
	assertHeaders(c, rec, headers)
}

func assertHeaders(c *C, rec *httptest.ResponseRecorder, h map[string]string) {
	for key, val := range h {
		c.Assert(rec.Header().Get(key), Equals, val)
	}
}

func (*WebServer) TestText404(c *C) {
	rec := runHandler(c, func(res http.ResponseWriter, _ *http.Request) {
		text404(res)
	})
	c.Assert(rec.Code, Equals, 404)
	c.Assert(rec.Body.String(), Equals, "404 Not found\n")
}

func (*WebServer) TestPanicHandler(c *C) {
	webRoot = "./test"
	err := errors.New("foo")

	// Prevent printing stack trace to terminal
	log.SetOutput(ioutil.Discard)
	rec := runHandler(c, func(res http.ResponseWriter, req *http.Request) {
		panicHandler(res, req, err)
	})
	log.SetOutput(os.Stdout)
	c.Assert(rec.Code, Equals, 500)
	c.Assert(rec.Body.String(), Equals, "<!doctype html>50x\n")
}

func (*WebServer) TestSetHeaders(c *C) {
	// HTML
	rec := httptest.NewRecorder()
	const etag = "foo"
	headers := map[string]string{
		"X-Frame-Options": "sameorigin",
		"Cache-Control":   "max-age=0, must-revalidate",
		"Expires":         "Fri, 01 Jan 1990 00:00:00 GMT",
		"ETag":            etag,
		"Content-Type":    "text/html; charset=UTF-8",
	}
	setHeaders(rec, etag, false)
	assertHeaders(c, rec, headers)

	// JSON
	headers["Content-Type"] = "application/json; charset=UTF-8"
	rec = httptest.NewRecorder()
	setHeaders(rec, etag, true)
	assertHeaders(c, rec, headers)
}

func (*WebServer) TestDetectLastN(c *C) {
	// No lastN query string
	req := customRequest(c, "/a/1")
	c.Assert(detectLastN(req), Equals, 0)

	// ?lastN value within bounds
	req = customRequest(c, "/a/1?lastN=100")
	c.Assert(detectLastN(req), Equals, 100)

	// ?lastN value beyond max
	req = customRequest(c, "/a/1?lastN=1000")
	c.Assert(detectLastN(req), Equals, 0)
}

func customRequest(c *C, url string) *http.Request {
	req, err := http.NewRequest("GET", url, nil)
	c.Assert(err, IsNil)
	return req
}

func (*WebServer) TestImageServer(c *C) {
	const (
		truncated         = "/src/tis life.gif"
		notFoundTruncated = "src/nobody here.gif"
	)
	imageWebRoot = "./test"
	path := imageWebRoot + truncated
	notFound := imageWebRoot + notFoundTruncated

	// Succesful first serve
	req := customRequest(c, path)
	rec := httptest.NewRecorder()
	params := httprouter.Params{
		httprouter.Param{
			Key:   "filepath",
			Value: truncated,
		},
	}
	serveImages(rec, req, params)
	buf, err := ioutil.ReadFile(path)
	c.Assert(err, IsNil)
	c.Assert(rec.Body.String(), Equals, string(buf))
	headers := map[string]string{
		"Cache-Control":   "max-age=30240000",
		"X-Frame-Options": "sameorigin",
		"ETag":            "0",
	}
	assertHeaders(c, rec, headers)

	// Fake etag validation
	req = customRequest(c, path)
	rec = httptest.NewRecorder()
	req.Header.Set("If-None-Match", "0")
	serveImages(rec, req, params)
	c.Assert(rec.Code, Equals, 304)

	// Non-existing file
	req = customRequest(c, notFound)
	rec = httptest.NewRecorder()
	params[0].Value = notFoundTruncated
	serveImages(rec, req, params)
	c.Assert(rec.Code, Equals, 404)
}

func (*WebServer) TestCompareEtag(c *C) {
	// Etag comparison
	rec := httptest.NewRecorder()
	req := newRequest(c)
	const etag = "foo"
	req.Header.Set("If-None-Match", etag)
	c.Assert(compareEtag(rec, req, etag, false), Equals, false)

	rec = httptest.NewRecorder()
	req = newRequest(c)
	headers := map[string]string{
		"ETag":          etag,
		"Content-Type":  "text/html; charset=UTF-8",
		"Cache-Control": "max-age=0, must-revalidate",
	}
	compareEtag(rec, req, etag, false)
	assertHeaders(c, rec, headers)
}

func (*WebServer) TestEtagStart(c *C) {
	c.Assert(etagStart(1), Equals, "W/1")
}

func (*WebServer) TestServeIndexTemplate(c *C) {
	const (
		desktopUA = "Mozilla/5.0 (Windows NT 6.1) AppleWebKit/537.36 " +
			"(KHTML, like Gecko) Chrome/41.0.2228.0 Safari/537.36"
		mobileUA = "Mozilla/5.0 (Linux; Android 4.1.1; Galaxy Nexus" +
			" Build/JRO03C) AppleWebKit/535.19 (KHTML, like Gecko)" +
			" Chrome/18.0.1025.166 Mobile Safari/535.19"
	)
	desktop := templateStore{[]byte("desktop"), "dhash"}
	mobile := templateStore{[]byte("mobile"), "mhash"}
	resources = templateMap{}
	resources["index"] = desktop
	resources["mobile"] = mobile

	// Desktop
	req := newRequest(c)
	req.Header.Set("User-Agent", desktopUA)
	rec := httptest.NewRecorder()
	serveIndexTemplate(rec, req)
	c.Assert(rec.Body.Bytes(), DeepEquals, desktop.HTML)
	c.Assert(rec.Header().Get("ETag"), Equals, desktop.Hash)

	// Mobile
	req = newRequest(c)
	req.Header.Set("User-Agent", mobileUA)
	rec = httptest.NewRecorder()
	serveIndexTemplate(rec, req)
	c.Assert(rec.Body.Bytes(), DeepEquals, mobile.HTML)
	c.Assert(rec.Header().Get("ETag"), Equals, mobile.Hash+"-mobile")

	// Etag matches
	req = newRequest(c)
	req.Header.Set("If-None-Match", desktop.Hash)
	rec = httptest.NewRecorder()
	serveIndexTemplate(rec, req)
	c.Assert(rec.Code, Equals, 304)
}

func (*DB) TestThreadHTML(c *C) {
	// Non-existant thread
	rec := httptest.NewRecorder()
	req := customRequest(c, "/a/22")
	body := []byte("body")
	resources = templateMap{
		"index": templateStore{
			HTML: body,
			Hash: "hash",
		},
	}
	webRoot = "./test"
	threadHTML("a")(rec, req, httprouter.Params{{Value: "22"}})
	c.Assert(rec.Code, Equals, 404)

	// Thread exists
	setupBoardAccess()
	setupPosts()
	rec = httptest.NewRecorder()
	req = customRequest(c, "/a/1")
	threadHTML("a")(rec, req, httprouter.Params{{Value: "1"}})
	c.Assert(rec.Body.Bytes(), DeepEquals, body)
}
<|endoftext|>


<commit_before>// Copyright (c) 2017, Janoš Guljaš <janos@resenje.org>
// All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package servers

import (
	"context"
	"fmt"
	"log"
	"net"
	"sync"
)

// Logger defines methods required for logging.
type Logger interface {
	Infof(format string, a ...interface{})
	Errorf(format string, a ...interface{})
}

// stdLogger is a simple implementation of Logger interface
// that uses log package for logging messages.
type stdLogger struct{}

func (l stdLogger) Infof(format string, a ...interface{}) {
	log.Printf("INFO "+format, a...)
}

func (l stdLogger) Errorf(format string, a ...interface{}) {
	log.Printf("ERROR "+format, a...)
}

// Option is a function that sets optional parameters for Servers.
type Option func(*Servers)

// WithLogger sets the Logger instance for logging messages.
func WithLogger(logger Logger) Option { return func(o *Servers) { o.logger = logger } }

// WithRecoverFunc sets a function that will be used to recover
// from panic inside a goroutune that servers are serving requests.
func WithRecoverFunc(recover func()) Option { return func(o *Servers) { o.recover = recover } }

// Servers holds a list of servers and their options.
// It provides a simple way to construct server group with Add method,
// to start them with Serve method, and stop them with Close or Shutdown methods.
type Servers struct {
	servers []*server
	mu      sync.Mutex
	logger  Logger
	recover func()
}

// New creates a new instance of Servers with applied options.
func New(opts ...Option) (s *Servers) {
	s = &Servers{
		logger:  stdLogger{},
		recover: func() {},
	}
	for _, opt := range opts {
		opt(s)
	}
	return
}

// Server defines required methods for a type that can be added to
// the Servers.
// In addition to this methods, a Server should implement TCPServer
// or UDPServer to be able to serve requests.
type Server interface {
	// Close should stop server from serving all existing requests
	// and stop accepting new ones.
	// The listener provided in Serve method must stop listening.
	Close() error
	// Shutdown should gracefully stop server. All existing requests
	// should be processed within a deadline provided by the context.
	// No new requests should be accepted.
	// The listener provided in Serve method must stop listening.
	Shutdown(ctx context.Context) error
}

// TCPServer defines methods for a server that accepts requests
// over TCP listener.
type TCPServer interface {
	// Serve should start server responding to requests.
	// The listener is initialized and already listening.
	ServeTCP(ln net.Listener) error
}

// UDPServer defines methods for a server that accepts requests
// over UDP listener.
type UDPServer interface {
	ServeUDP(conn *net.UDPConn) error
}

type server struct {
	Server
	name    string
	address string
	tcpAddr *net.TCPAddr
	udpAddr *net.UDPAddr
}

func (s *server) label() string {
	if s.name == "" {
		return "server"
	}
	return s.name + " server"
}

func (s *server) isTCP() (srv TCPServer, yes bool) {
	srv, yes = s.Server.(TCPServer)
	return
}

func (s *server) isUDP() (srv UDPServer, yes bool) {
	srv, yes = s.Server.(UDPServer)
	return
}

// Add adds a new server instance by a custom name and with
// address to listen to.
func (s *Servers) Add(name, address string, srv Server) {
	s.mu.Lock()
	s.servers = append(s.servers, &server{
		Server:  srv,
		name:    name,
		address: address,
	})
	s.mu.Unlock()
}

// Serve starts all added servers.
// New new servers must be added after this methid is called.
func (s *Servers) Serve() (err error) {
	lns := make([]net.Listener, len(s.servers))
	conns := make([]*net.UDPConn, len(s.servers))
	for i, srv := range s.servers {
		if _, yes := srv.isTCP(); yes {
			ln, err := net.Listen("tcp", srv.address)
			if err != nil {
				for _, l := range lns {
					if l == nil {
						continue
					}
					if err := l.Close(); err != nil {
						s.logger.Errorf("%s tcp listener %q close: %v", srv.label(), srv.address, err)
					}
				}
				return fmt.Errorf("%s tcp listener %q: %v", srv.label(), srv.address, err)
			}
			lns[i] = ln
		}
		if _, yes := srv.isUDP(); yes {
			addr, err := net.ResolveUDPAddr("udp", srv.address)
			if err != nil {
				return fmt.Errorf("%s resolve udp address %q: %v", srv.label(), srv.address, err)
			}
			conn, err := net.ListenUDP("udp", addr)
			if err != nil {
				return fmt.Errorf("%s udp listener %q: %v", srv.label(), srv.address, err)
			}
			conns[i] = conn
		}
	}
	for i, srv := range s.servers {
		if tcpSrv, yes := srv.isTCP(); yes {
			go func(srv *server, ln net.Listener) {
				defer s.recover()

				s.mu.Lock()
				srv.tcpAddr = ln.Addr().(*net.TCPAddr)
				s.mu.Unlock()

				s.logger.Infof("%s listening on %q", srv.label(), srv.tcpAddr.String())
				if err = tcpSrv.ServeTCP(ln); err != nil {
					s.logger.Errorf("%s serve %q: %v", srv.label(), srv.tcpAddr.String(), err)
				}
			}(srv, lns[i])
		}
		if udpSrv, yes := srv.isUDP(); yes {
			go func(srv *server, conn *net.UDPConn) {
				defer s.recover()

				s.mu.Lock()
				srv.udpAddr = conn.LocalAddr().(*net.UDPAddr)
				s.mu.Unlock()

				s.logger.Infof("%s listening on %q", srv.label(), srv.tcpAddr.String())
				if err = udpSrv.ServeUDP(conn); err != nil {
					s.logger.Errorf("%s serve %q: %v", srv.label(), srv.tcpAddr.String(), err)
				}
			}(srv, conns[i])
		}
	}
	return nil
}

// TCPAddr returns a TCP address of the listener that a server
// with a specific name is using. If there are more servers
// with the same name, the address of the first started server
// is returned.
func (s *Servers) TCPAddr(name string) (a *net.TCPAddr) {
	s.mu.Lock()
	defer s.mu.Unlock()

	for _, srv := range s.servers {
		if srv.name == name {
			return srv.tcpAddr
		}
	}
	return nil
}

// UDPAddr returns a UDP address of the listener that a server
// with a specific name is using. If there are more servers
// with the same name, the address of the first started server
// is returned.
func (s *Servers) UDPAddr(name string) (a *net.UDPAddr) {
	s.mu.Lock()
	defer s.mu.Unlock()

	for _, srv := range s.servers {
		if srv.name == name {
			return srv.udpAddr
		}
	}
	return nil
}

// Close stops all servers, by calling Close method on each of them.
func (s *Servers) Close() {
	wg := &sync.WaitGroup{}
	for _, srv := range s.servers {
		wg.Add(1)
		go func(srv *server) {
			defer s.recover()
			defer wg.Done()

			s.logger.Infof("%s closing", srv.label())
			if err := srv.Close(); err != nil {
				s.logger.Errorf("%s close: %v", srv.label(), err)
			}
		}(srv)
	}
	wg.Wait()
	return
}

// Shutdown gracefully stops all servers, by calling Shutdown method on each of them.
func (s *Servers) Shutdown(ctx context.Context) {
	wg := &sync.WaitGroup{}
	for _, srv := range s.servers {
		wg.Add(1)
		go func(srv *server) {
			defer s.recover()
			defer wg.Done()

			s.logger.Infof("%s shutting down", srv.label())
			if err := srv.Shutdown(ctx); err != nil {
				s.logger.Errorf("%s shutdown: %v", srv.label(), err)
			}
		}(srv)
	}
	wg.Wait()
	return
}
<commit_msg>servers: keep error variable local to a goroutine<commit_after>// Copyright (c) 2017, Janoš Guljaš <janos@resenje.org>
// All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package servers

import (
	"context"
	"fmt"
	"log"
	"net"
	"sync"
)

// Logger defines methods required for logging.
type Logger interface {
	Infof(format string, a ...interface{})
	Errorf(format string, a ...interface{})
}

// stdLogger is a simple implementation of Logger interface
// that uses log package for logging messages.
type stdLogger struct{}

func (l stdLogger) Infof(format string, a ...interface{}) {
	log.Printf("INFO "+format, a...)
}

func (l stdLogger) Errorf(format string, a ...interface{}) {
	log.Printf("ERROR "+format, a...)
}

// Option is a function that sets optional parameters for Servers.
type Option func(*Servers)

// WithLogger sets the Logger instance for logging messages.
func WithLogger(logger Logger) Option { return func(o *Servers) { o.logger = logger } }

// WithRecoverFunc sets a function that will be used to recover
// from panic inside a goroutune that servers are serving requests.
func WithRecoverFunc(recover func()) Option { return func(o *Servers) { o.recover = recover } }

// Servers holds a list of servers and their options.
// It provides a simple way to construct server group with Add method,
// to start them with Serve method, and stop them with Close or Shutdown methods.
type Servers struct {
	servers []*server
	mu      sync.Mutex
	logger  Logger
	recover func()
}

// New creates a new instance of Servers with applied options.
func New(opts ...Option) (s *Servers) {
	s = &Servers{
		logger:  stdLogger{},
		recover: func() {},
	}
	for _, opt := range opts {
		opt(s)
	}
	return
}

// Server defines required methods for a type that can be added to
// the Servers.
// In addition to this methods, a Server should implement TCPServer
// or UDPServer to be able to serve requests.
type Server interface {
	// Close should stop server from serving all existing requests
	// and stop accepting new ones.
	// The listener provided in Serve method must stop listening.
	Close() error
	// Shutdown should gracefully stop server. All existing requests
	// should be processed within a deadline provided by the context.
	// No new requests should be accepted.
	// The listener provided in Serve method must stop listening.
	Shutdown(ctx context.Context) error
}

// TCPServer defines methods for a server that accepts requests
// over TCP listener.
type TCPServer interface {
	// Serve should start server responding to requests.
	// The listener is initialized and already listening.
	ServeTCP(ln net.Listener) error
}

// UDPServer defines methods for a server that accepts requests
// over UDP listener.
type UDPServer interface {
	ServeUDP(conn *net.UDPConn) error
}

type server struct {
	Server
	name    string
	address string
	tcpAddr *net.TCPAddr
	udpAddr *net.UDPAddr
}

func (s *server) label() string {
	if s.name == "" {
		return "server"
	}
	return s.name + " server"
}

func (s *server) isTCP() (srv TCPServer, yes bool) {
	srv, yes = s.Server.(TCPServer)
	return
}

func (s *server) isUDP() (srv UDPServer, yes bool) {
	srv, yes = s.Server.(UDPServer)
	return
}

// Add adds a new server instance by a custom name and with
// address to listen to.
func (s *Servers) Add(name, address string, srv Server) {
	s.mu.Lock()
	s.servers = append(s.servers, &server{
		Server:  srv,
		name:    name,
		address: address,
	})
	s.mu.Unlock()
}

// Serve starts all added servers.
// New new servers must be added after this methid is called.
func (s *Servers) Serve() (err error) {
	lns := make([]net.Listener, len(s.servers))
	conns := make([]*net.UDPConn, len(s.servers))
	for i, srv := range s.servers {
		if _, yes := srv.isTCP(); yes {
			ln, err := net.Listen("tcp", srv.address)
			if err != nil {
				for _, l := range lns {
					if l == nil {
						continue
					}
					if err := l.Close(); err != nil {
						s.logger.Errorf("%s tcp listener %q close: %v", srv.label(), srv.address, err)
					}
				}
				return fmt.Errorf("%s tcp listener %q: %v", srv.label(), srv.address, err)
			}
			lns[i] = ln
		}
		if _, yes := srv.isUDP(); yes {
			addr, err := net.ResolveUDPAddr("udp", srv.address)
			if err != nil {
				return fmt.Errorf("%s resolve udp address %q: %v", srv.label(), srv.address, err)
			}
			conn, err := net.ListenUDP("udp", addr)
			if err != nil {
				return fmt.Errorf("%s udp listener %q: %v", srv.label(), srv.address, err)
			}
			conns[i] = conn
		}
	}
	for i, srv := range s.servers {
		if tcpSrv, yes := srv.isTCP(); yes {
			go func(srv *server, ln net.Listener) {
				defer s.recover()

				s.mu.Lock()
				srv.tcpAddr = ln.Addr().(*net.TCPAddr)
				s.mu.Unlock()

				s.logger.Infof("%s listening on %q", srv.label(), srv.tcpAddr.String())
				if err := tcpSrv.ServeTCP(ln); err != nil {
					s.logger.Errorf("%s serve %q: %v", srv.label(), srv.tcpAddr.String(), err)
				}
			}(srv, lns[i])
		}
		if udpSrv, yes := srv.isUDP(); yes {
			go func(srv *server, conn *net.UDPConn) {
				defer s.recover()

				s.mu.Lock()
				srv.udpAddr = conn.LocalAddr().(*net.UDPAddr)
				s.mu.Unlock()

				s.logger.Infof("%s listening on %q", srv.label(), srv.tcpAddr.String())
				if err := udpSrv.ServeUDP(conn); err != nil {
					s.logger.Errorf("%s serve %q: %v", srv.label(), srv.tcpAddr.String(), err)
				}
			}(srv, conns[i])
		}
	}
	return nil
}

// TCPAddr returns a TCP address of the listener that a server
// with a specific name is using. If there are more servers
// with the same name, the address of the first started server
// is returned.
func (s *Servers) TCPAddr(name string) (a *net.TCPAddr) {
	s.mu.Lock()
	defer s.mu.Unlock()

	for _, srv := range s.servers {
		if srv.name == name {
			return srv.tcpAddr
		}
	}
	return nil
}

// UDPAddr returns a UDP address of the listener that a server
// with a specific name is using. If there are more servers
// with the same name, the address of the first started server
// is returned.
func (s *Servers) UDPAddr(name string) (a *net.UDPAddr) {
	s.mu.Lock()
	defer s.mu.Unlock()

	for _, srv := range s.servers {
		if srv.name == name {
			return srv.udpAddr
		}
	}
	return nil
}

// Close stops all servers, by calling Close method on each of them.
func (s *Servers) Close() {
	wg := &sync.WaitGroup{}
	for _, srv := range s.servers {
		wg.Add(1)
		go func(srv *server) {
			defer s.recover()
			defer wg.Done()

			s.logger.Infof("%s closing", srv.label())
			if err := srv.Close(); err != nil {
				s.logger.Errorf("%s close: %v", srv.label(), err)
			}
		}(srv)
	}
	wg.Wait()
	return
}

// Shutdown gracefully stops all servers, by calling Shutdown method on each of them.
func (s *Servers) Shutdown(ctx context.Context) {
	wg := &sync.WaitGroup{}
	for _, srv := range s.servers {
		wg.Add(1)
		go func(srv *server) {
			defer s.recover()
			defer wg.Done()

			s.logger.Infof("%s shutting down", srv.label())
			if err := srv.Shutdown(ctx); err != nil {
				s.logger.Errorf("%s shutdown: %v", srv.label(), err)
			}
		}(srv)
	}
	wg.Wait()
	return
}
<|endoftext|>


<commit_before>// Copyright 2019, OpenTelemetry Authors
//
// 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 service handles the command-line, configuration, and runs the
// OpenTelemetry Collector.
package service

import (
	"context"
	"flag"
	"fmt"
	"os"
	"os/signal"
	"runtime"
	"syscall"

	"github.com/pkg/errors"
	"github.com/spf13/cobra"
	"github.com/spf13/viper"
	"go.uber.org/zap"

	"github.com/open-telemetry/opentelemetry-collector/component"
	"github.com/open-telemetry/opentelemetry-collector/component/componenterror"
	"github.com/open-telemetry/opentelemetry-collector/config"
	"github.com/open-telemetry/opentelemetry-collector/config/configcheck"
	"github.com/open-telemetry/opentelemetry-collector/config/configmodels"
	"github.com/open-telemetry/opentelemetry-collector/internal/collector/telemetry"
	"github.com/open-telemetry/opentelemetry-collector/service/builder"
)

// Application represents a collector application
type Application struct {
	info            ApplicationStartInfo
	rootCmd         *cobra.Command
	v               *viper.Viper
	logger          *zap.Logger
	builtExporters  builder.Exporters
	builtReceivers  builder.Receivers
	builtPipelines  builder.BuiltPipelines
	builtExtensions builder.Extensions

	factories config.Factories
	config    *configmodels.Config

	// stopTestChan is used to terminate the application in end to end tests.
	stopTestChan chan struct{}
	// readyChan is used in tests to indicate that the application is ready.
	readyChan chan struct{}

	// asyncErrorChannel is used to signal a fatal error from any component.
	asyncErrorChannel chan error
}

// Command returns Application's root command.
func (app *Application) Command() *cobra.Command {
	return app.rootCmd
}

// ApplicationStartInfo is the information that is logged at the application start.
// This information can be overridden in custom builds.
type ApplicationStartInfo struct {
	// Executable file name, e.g. "otelcol".
	ExeName string

	// Long name, used e.g. in the logs.
	LongName string

	// Version string.
	Version string

	// Git hash of the source code.
	GitHash string
}

// Parameters holds configuration for creating a new Application.
type Parameters struct {
	// Factories component factories.
	Factories config.Factories
	// ApplicationStartInfo provides application start information.
	ApplicationStartInfo ApplicationStartInfo
	// ConfigFactory that creates the configuration.
	// If it is not provided the default factory will be used.
	// The default factory loads the configuration specified as a command line flag.
	ConfigFactory ConfigFactory
}

// ConfigFactory creates config.
type ConfigFactory func(v *viper.Viper, factories config.Factories) (*configmodels.Config, error)

func fileLoaderConfigFactory(v *viper.Viper, factories config.Factories) (*configmodels.Config, error) {
	file := builder.GetConfigFile()
	if file == "" {
		return nil, errors.New("config file not specified")
	}
	v.SetConfigFile(file)
	err := v.ReadInConfig()
	if err != nil {
		return nil, fmt.Errorf("error loading config file %q: %v", file, err)
	}
	return config.Load(v, factories)
}

// New creates and returns a new instance of Application.
func New(params Parameters) (*Application, error) {
	app := &Application{
		info:      params.ApplicationStartInfo,
		v:         config.NewViper(),
		readyChan: make(chan struct{}),
		factories: params.Factories,
	}

	factory := params.ConfigFactory
	if factory == nil {
		// use default factory that loads the configuration file
		factory = fileLoaderConfigFactory
	}

	rootCmd := &cobra.Command{
		Use:  params.ApplicationStartInfo.ExeName,
		Long: params.ApplicationStartInfo.LongName,
		RunE: func(cmd *cobra.Command, args []string) error {
			err := app.init()
			if err != nil {
				return err
			}

			err = app.execute(context.Background(), factory)
			if err != nil {
				return err
			}

			return nil
		},
	}

	// TODO: coalesce this code and expose this information to other components.
	flagSet := new(flag.FlagSet)
	addFlagsFns := []func(*flag.FlagSet){
		telemetry.Flags,
		builder.Flags,
		loggerFlags,
	}
	for _, addFlags := range addFlagsFns {
		addFlags(flagSet)
	}
	rootCmd.Flags().AddGoFlagSet(flagSet)

	app.rootCmd = rootCmd

	return app, nil
}

// ReportFatalError is used to report to the host that the receiver encountered
// a fatal error (i.e.: an error that the instance can't recover from) after
// its start function has already returned.
func (app *Application) ReportFatalError(err error) {
	app.asyncErrorChannel <- err
}

func (app *Application) GetFactory(kind component.Kind, componentType configmodels.Type) component.Factory {
	switch kind {
	case component.KindReceiver:
		return app.factories.Receivers[componentType]
	case component.KindProcessor:
		return app.factories.Processors[componentType]
	case component.KindExporter:
		return app.factories.Exporters[componentType]
	case component.KindExtension:
		return app.factories.Extensions[componentType]
	}
	return nil
}

func (app *Application) GetExtensions() map[configmodels.Extension]component.ServiceExtension {
	return app.builtExtensions.ToMap()
}

func (app *Application) GetExporters() map[configmodels.DataType]map[configmodels.Exporter]component.Exporter {
	return app.builtExporters.ToMapByDataType()
}

func (app *Application) init() error {
	l, err := newLogger()
	if err != nil {
		return errors.Wrap(err, "failed to get logger")
	}
	app.logger = l
	return nil
}

func (app *Application) setupTelemetry(ballastSizeBytes uint64) error {
	app.logger.Info("Setting up own telemetry...")

	err := AppTelemetry.init(app.asyncErrorChannel, ballastSizeBytes, app.logger)
	if err != nil {
		return errors.Wrap(err, "failed to initialize telemetry")
	}

	return nil
}

// runAndWaitForShutdownEvent waits for one of the shutdown events that can happen.
func (app *Application) runAndWaitForShutdownEvent() {
	app.logger.Info("Everything is ready. Begin running and processing data.")

	// Plug SIGTERM signal into a channel.
	signalsChannel := make(chan os.Signal, 1)
	signal.Notify(signalsChannel, os.Interrupt, syscall.SIGTERM)

	// set the channel to stop testing.
	app.stopTestChan = make(chan struct{})
	// notify tests that it is ready.
	close(app.readyChan)

	select {
	case err := <-app.asyncErrorChannel:
		app.logger.Error("Asynchronous error received, terminating process", zap.Error(err))
	case s := <-signalsChannel:
		app.logger.Info("Received signal from OS", zap.String("signal", s.String()))
	case <-app.stopTestChan:
		app.logger.Info("Received stop test request")
	}
}

func (app *Application) setupConfigurationComponents(ctx context.Context, factory ConfigFactory) error {
	if err := configcheck.ValidateConfigFromFactories(app.factories); err != nil {
		return err
	}

	app.logger.Info("Loading configuration...")
	cfg, err := factory(app.v, app.factories)
	if err != nil {
		return errors.Wrap(err, "cannot load configuration")
	}
	err = config.ValidateConfig(cfg, app.logger)
	if err != nil {
		return errors.Wrap(err, "cannot load configuration")
	}

	app.config = cfg
	app.logger.Info("Applying configuration...")

	err = app.setupExtensions(ctx)
	if err != nil {
		return errors.Wrap(err, "cannot setup extensions")
	}

	err = app.setupPipelines(ctx)
	if err != nil {
		return errors.Wrap(err, "cannot setup pipelines")
	}

	return nil
}

func (app *Application) setupExtensions(ctx context.Context) error {
	var err error
	app.builtExtensions, err = builder.NewExtensionsBuilder(app.logger, app.config, app.factories.Extensions).Build()
	if err != nil {
		return errors.Wrap(err, "cannot build builtExtensions")
	}
	app.logger.Info("Starting extensions...")
	return app.builtExtensions.StartAll(ctx, app)
}

func (app *Application) setupPipelines(ctx context.Context) error {
	// Pipeline is built backwards, starting from exporters, so that we create objects
	// which are referenced before objects which reference them.

	// First create exporters.
	var err error
	app.builtExporters, err = builder.NewExportersBuilder(app.logger, app.config, app.factories.Exporters).Build()
	if err != nil {
		return errors.Wrap(err, "cannot build builtExporters")
	}

	app.logger.Info("Starting exporters...")
	err = app.builtExporters.StartAll(ctx, app)
	if err != nil {
		return errors.Wrap(err, "cannot start builtExporters")
	}

	// Create pipelines and their processors and plug exporters to the
	// end of the pipelines.
	app.builtPipelines, err = builder.NewPipelinesBuilder(app.logger, app.config, app.builtExporters, app.factories.Processors).Build()
	if err != nil {
		return errors.Wrap(err, "cannot build pipelines")
	}

	app.logger.Info("Starting processors...")
	err = app.builtPipelines.StartProcessors(ctx, app)
	if err != nil {
		return errors.Wrap(err, "cannot start processors")
	}

	// Create receivers and plug them into the start of the pipelines.
	app.builtReceivers, err = builder.NewReceiversBuilder(app.logger, app.config, app.builtPipelines, app.factories.Receivers).Build()
	if err != nil {
		return errors.Wrap(err, "cannot build receivers")
	}

	app.logger.Info("Starting receivers...")
	err = app.builtReceivers.StartAll(ctx, app)
	if err != nil {
		return errors.Wrap(err, "cannot start receivers")
	}

	return nil
}

func (app *Application) shutdownPipelines(ctx context.Context) error {
	// Shutdown order is the reverse of building: first receivers, then flushing pipelines
	// giving senders a chance to send all their data. This may take time, the allowed
	// time should be part of configuration.

	var errs []error

	app.logger.Info("Stopping receivers...")
	err := app.builtReceivers.ShutdownAll(ctx)
	if err != nil {
		errs = append(errs, errors.Wrap(err, "failed to stop receivers"))
	}

	app.logger.Info("Stopping processors...")
	err = app.builtPipelines.ShutdownProcessors(ctx)
	if err != nil {
		errs = append(errs, errors.Wrap(err, "failed to shutdown processors"))
	}

	app.logger.Info("Stopping exporters...")
	err = app.builtExporters.ShutdownAll(ctx)
	if err != nil {
		errs = append(errs, errors.Wrap(err, "failed to shutdown exporters"))
	}

	if len(errs) != 0 {
		return componenterror.CombineErrors(errs)
	}

	return nil
}

func (app *Application) shutdownExtensions(ctx context.Context) error {
	app.logger.Info("Stopping extensions...")
	err := app.builtExtensions.ShutdownAll(ctx)
	if err != nil {
		return errors.Wrap(err, "failed to shutdown extensions")
	}
	return nil
}

func (app *Application) execute(ctx context.Context, factory ConfigFactory) error {
	app.logger.Info("Starting "+app.info.LongName+"...",
		zap.String("Version", app.info.Version),
		zap.String("GitHash", app.info.GitHash),
		zap.Int("NumCPU", runtime.NumCPU()),
	)

	// Set memory ballast
	ballast, ballastSizeBytes := app.createMemoryBallast()

	app.asyncErrorChannel = make(chan error)

	// Setup everything.
	err := app.setupTelemetry(ballastSizeBytes)
	if err != nil {
		return err
	}

	err = app.setupConfigurationComponents(ctx, factory)
	if err != nil {
		return err
	}

	err = app.builtExtensions.NotifyPipelineReady()
	if err != nil {
		return err
	}

	// Everything is ready, now run until an event requiring shutdown happens.
	app.runAndWaitForShutdownEvent()

	// Accumulate errors and proceed with shutting down remaining components.
	var errs []error

	// Begin shutdown sequence.
	runtime.KeepAlive(ballast)
	app.logger.Info("Starting shutdown...")

	err = app.builtExtensions.NotifyPipelineNotReady()
	if err != nil {
		errs = append(errs, errors.Wrap(err, "failed to notify that pipeline is not ready"))
	}

	err = app.shutdownPipelines(ctx)
	if err != nil {
		errs = append(errs, errors.Wrap(err, "failed to shutdown pipelines"))
	}

	err = app.shutdownExtensions(ctx)
	if err != nil {
		errs = append(errs, errors.Wrap(err, "failed to shutdown extensions"))
	}

	AppTelemetry.shutdown()

	app.logger.Info("Shutdown complete.")

	if len(errs) != 0 {
		return componenterror.CombineErrors(errs)
	}
	return nil
}

// Start starts the collector according to the command and configuration
// given by the user.
func (app *Application) Start() error {
	// From this point on do not show usage in case of error.
	app.rootCmd.SilenceUsage = true

	return app.rootCmd.Execute()
}

func (app *Application) createMemoryBallast() ([]byte, uint64) {
	ballastSizeMiB := builder.MemBallastSize()
	if ballastSizeMiB > 0 {
		ballastSizeBytes := uint64(ballastSizeMiB) * 1024 * 1024
		ballast := make([]byte, ballastSizeBytes)
		app.logger.Info("Using memory ballast", zap.Int("MiBs", ballastSizeMiB))
		return ballast, ballastSizeBytes
	}
	return nil, 0
}
<commit_msg>Make file config factory public (#887)<commit_after>// Copyright 2019, OpenTelemetry Authors
//
// 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 service handles the command-line, configuration, and runs the
// OpenTelemetry Collector.
package service

import (
	"context"
	"flag"
	"fmt"
	"os"
	"os/signal"
	"runtime"
	"syscall"

	"github.com/pkg/errors"
	"github.com/spf13/cobra"
	"github.com/spf13/viper"
	"go.uber.org/zap"

	"github.com/open-telemetry/opentelemetry-collector/component"
	"github.com/open-telemetry/opentelemetry-collector/component/componenterror"
	"github.com/open-telemetry/opentelemetry-collector/config"
	"github.com/open-telemetry/opentelemetry-collector/config/configcheck"
	"github.com/open-telemetry/opentelemetry-collector/config/configmodels"
	"github.com/open-telemetry/opentelemetry-collector/internal/collector/telemetry"
	"github.com/open-telemetry/opentelemetry-collector/service/builder"
)

// Application represents a collector application
type Application struct {
	info            ApplicationStartInfo
	rootCmd         *cobra.Command
	v               *viper.Viper
	logger          *zap.Logger
	builtExporters  builder.Exporters
	builtReceivers  builder.Receivers
	builtPipelines  builder.BuiltPipelines
	builtExtensions builder.Extensions

	factories config.Factories
	config    *configmodels.Config

	// stopTestChan is used to terminate the application in end to end tests.
	stopTestChan chan struct{}
	// readyChan is used in tests to indicate that the application is ready.
	readyChan chan struct{}

	// asyncErrorChannel is used to signal a fatal error from any component.
	asyncErrorChannel chan error
}

// Command returns Application's root command.
func (app *Application) Command() *cobra.Command {
	return app.rootCmd
}

// ApplicationStartInfo is the information that is logged at the application start.
// This information can be overridden in custom builds.
type ApplicationStartInfo struct {
	// Executable file name, e.g. "otelcol".
	ExeName string

	// Long name, used e.g. in the logs.
	LongName string

	// Version string.
	Version string

	// Git hash of the source code.
	GitHash string
}

// Parameters holds configuration for creating a new Application.
type Parameters struct {
	// Factories component factories.
	Factories config.Factories
	// ApplicationStartInfo provides application start information.
	ApplicationStartInfo ApplicationStartInfo
	// ConfigFactory that creates the configuration.
	// If it is not provided the default factory (FileLoaderConfigFactory) is used.
	// The default factory loads the configuration specified as a command line flag.
	ConfigFactory ConfigFactory
}

// ConfigFactory creates config.
type ConfigFactory func(v *viper.Viper, factories config.Factories) (*configmodels.Config, error)

// FileLoaderConfigFactory implements ConfigFactory and it creates configuration from file.
func FileLoaderConfigFactory(v *viper.Viper, factories config.Factories) (*configmodels.Config, error) {
	file := builder.GetConfigFile()
	if file == "" {
		return nil, errors.New("config file not specified")
	}
	v.SetConfigFile(file)
	err := v.ReadInConfig()
	if err != nil {
		return nil, fmt.Errorf("error loading config file %q: %v", file, err)
	}
	return config.Load(v, factories)
}

// New creates and returns a new instance of Application.
func New(params Parameters) (*Application, error) {
	app := &Application{
		info:      params.ApplicationStartInfo,
		v:         config.NewViper(),
		readyChan: make(chan struct{}),
		factories: params.Factories,
	}

	factory := params.ConfigFactory
	if factory == nil {
		// use default factory that loads the configuration file
		factory = FileLoaderConfigFactory
	}

	rootCmd := &cobra.Command{
		Use:  params.ApplicationStartInfo.ExeName,
		Long: params.ApplicationStartInfo.LongName,
		RunE: func(cmd *cobra.Command, args []string) error {
			err := app.init()
			if err != nil {
				return err
			}

			err = app.execute(context.Background(), factory)
			if err != nil {
				return err
			}

			return nil
		},
	}

	// TODO: coalesce this code and expose this information to other components.
	flagSet := new(flag.FlagSet)
	addFlagsFns := []func(*flag.FlagSet){
		telemetry.Flags,
		builder.Flags,
		loggerFlags,
	}
	for _, addFlags := range addFlagsFns {
		addFlags(flagSet)
	}
	rootCmd.Flags().AddGoFlagSet(flagSet)

	app.rootCmd = rootCmd

	return app, nil
}

// ReportFatalError is used to report to the host that the receiver encountered
// a fatal error (i.e.: an error that the instance can't recover from) after
// its start function has already returned.
func (app *Application) ReportFatalError(err error) {
	app.asyncErrorChannel <- err
}

func (app *Application) GetFactory(kind component.Kind, componentType configmodels.Type) component.Factory {
	switch kind {
	case component.KindReceiver:
		return app.factories.Receivers[componentType]
	case component.KindProcessor:
		return app.factories.Processors[componentType]
	case component.KindExporter:
		return app.factories.Exporters[componentType]
	case component.KindExtension:
		return app.factories.Extensions[componentType]
	}
	return nil
}

func (app *Application) GetExtensions() map[configmodels.Extension]component.ServiceExtension {
	return app.builtExtensions.ToMap()
}

func (app *Application) GetExporters() map[configmodels.DataType]map[configmodels.Exporter]component.Exporter {
	return app.builtExporters.ToMapByDataType()
}

func (app *Application) init() error {
	l, err := newLogger()
	if err != nil {
		return errors.Wrap(err, "failed to get logger")
	}
	app.logger = l
	return nil
}

func (app *Application) setupTelemetry(ballastSizeBytes uint64) error {
	app.logger.Info("Setting up own telemetry...")

	err := AppTelemetry.init(app.asyncErrorChannel, ballastSizeBytes, app.logger)
	if err != nil {
		return errors.Wrap(err, "failed to initialize telemetry")
	}

	return nil
}

// runAndWaitForShutdownEvent waits for one of the shutdown events that can happen.
func (app *Application) runAndWaitForShutdownEvent() {
	app.logger.Info("Everything is ready. Begin running and processing data.")

	// Plug SIGTERM signal into a channel.
	signalsChannel := make(chan os.Signal, 1)
	signal.Notify(signalsChannel, os.Interrupt, syscall.SIGTERM)

	// set the channel to stop testing.
	app.stopTestChan = make(chan struct{})
	// notify tests that it is ready.
	close(app.readyChan)

	select {
	case err := <-app.asyncErrorChannel:
		app.logger.Error("Asynchronous error received, terminating process", zap.Error(err))
	case s := <-signalsChannel:
		app.logger.Info("Received signal from OS", zap.String("signal", s.String()))
	case <-app.stopTestChan:
		app.logger.Info("Received stop test request")
	}
}

func (app *Application) setupConfigurationComponents(ctx context.Context, factory ConfigFactory) error {
	if err := configcheck.ValidateConfigFromFactories(app.factories); err != nil {
		return err
	}

	app.logger.Info("Loading configuration...")
	cfg, err := factory(app.v, app.factories)
	if err != nil {
		return errors.Wrap(err, "cannot load configuration")
	}
	err = config.ValidateConfig(cfg, app.logger)
	if err != nil {
		return errors.Wrap(err, "cannot load configuration")
	}

	app.config = cfg
	app.logger.Info("Applying configuration...")

	err = app.setupExtensions(ctx)
	if err != nil {
		return errors.Wrap(err, "cannot setup extensions")
	}

	err = app.setupPipelines(ctx)
	if err != nil {
		return errors.Wrap(err, "cannot setup pipelines")
	}

	return nil
}

func (app *Application) setupExtensions(ctx context.Context) error {
	var err error
	app.builtExtensions, err = builder.NewExtensionsBuilder(app.logger, app.config, app.factories.Extensions).Build()
	if err != nil {
		return errors.Wrap(err, "cannot build builtExtensions")
	}
	app.logger.Info("Starting extensions...")
	return app.builtExtensions.StartAll(ctx, app)
}

func (app *Application) setupPipelines(ctx context.Context) error {
	// Pipeline is built backwards, starting from exporters, so that we create objects
	// which are referenced before objects which reference them.

	// First create exporters.
	var err error
	app.builtExporters, err = builder.NewExportersBuilder(app.logger, app.config, app.factories.Exporters).Build()
	if err != nil {
		return errors.Wrap(err, "cannot build builtExporters")
	}

	app.logger.Info("Starting exporters...")
	err = app.builtExporters.StartAll(ctx, app)
	if err != nil {
		return errors.Wrap(err, "cannot start builtExporters")
	}

	// Create pipelines and their processors and plug exporters to the
	// end of the pipelines.
	app.builtPipelines, err = builder.NewPipelinesBuilder(app.logger, app.config, app.builtExporters, app.factories.Processors).Build()
	if err != nil {
		return errors.Wrap(err, "cannot build pipelines")
	}

	app.logger.Info("Starting processors...")
	err = app.builtPipelines.StartProcessors(ctx, app)
	if err != nil {
		return errors.Wrap(err, "cannot start processors")
	}

	// Create receivers and plug them into the start of the pipelines.
	app.builtReceivers, err = builder.NewReceiversBuilder(app.logger, app.config, app.builtPipelines, app.factories.Receivers).Build()
	if err != nil {
		return errors.Wrap(err, "cannot build receivers")
	}

	app.logger.Info("Starting receivers...")
	err = app.builtReceivers.StartAll(ctx, app)
	if err != nil {
		return errors.Wrap(err, "cannot start receivers")
	}

	return nil
}

func (app *Application) shutdownPipelines(ctx context.Context) error {
	// Shutdown order is the reverse of building: first receivers, then flushing pipelines
	// giving senders a chance to send all their data. This may take time, the allowed
	// time should be part of configuration.

	var errs []error

	app.logger.Info("Stopping receivers...")
	err := app.builtReceivers.ShutdownAll(ctx)
	if err != nil {
		errs = append(errs, errors.Wrap(err, "failed to stop receivers"))
	}

	app.logger.Info("Stopping processors...")
	err = app.builtPipelines.ShutdownProcessors(ctx)
	if err != nil {
		errs = append(errs, errors.Wrap(err, "failed to shutdown processors"))
	}

	app.logger.Info("Stopping exporters...")
	err = app.builtExporters.ShutdownAll(ctx)
	if err != nil {
		errs = append(errs, errors.Wrap(err, "failed to shutdown exporters"))
	}

	if len(errs) != 0 {
		return componenterror.CombineErrors(errs)
	}

	return nil
}

func (app *Application) shutdownExtensions(ctx context.Context) error {
	app.logger.Info("Stopping extensions...")
	err := app.builtExtensions.ShutdownAll(ctx)
	if err != nil {
		return errors.Wrap(err, "failed to shutdown extensions")
	}
	return nil
}

func (app *Application) execute(ctx context.Context, factory ConfigFactory) error {
	app.logger.Info("Starting "+app.info.LongName+"...",
		zap.String("Version", app.info.Version),
		zap.String("GitHash", app.info.GitHash),
		zap.Int("NumCPU", runtime.NumCPU()),
	)

	// Set memory ballast
	ballast, ballastSizeBytes := app.createMemoryBallast()

	app.asyncErrorChannel = make(chan error)

	// Setup everything.
	err := app.setupTelemetry(ballastSizeBytes)
	if err != nil {
		return err
	}

	err = app.setupConfigurationComponents(ctx, factory)
	if err != nil {
		return err
	}

	err = app.builtExtensions.NotifyPipelineReady()
	if err != nil {
		return err
	}

	// Everything is ready, now run until an event requiring shutdown happens.
	app.runAndWaitForShutdownEvent()

	// Accumulate errors and proceed with shutting down remaining components.
	var errs []error

	// Begin shutdown sequence.
	runtime.KeepAlive(ballast)
	app.logger.Info("Starting shutdown...")

	err = app.builtExtensions.NotifyPipelineNotReady()
	if err != nil {
		errs = append(errs, errors.Wrap(err, "failed to notify that pipeline is not ready"))
	}

	err = app.shutdownPipelines(ctx)
	if err != nil {
		errs = append(errs, errors.Wrap(err, "failed to shutdown pipelines"))
	}

	err = app.shutdownExtensions(ctx)
	if err != nil {
		errs = append(errs, errors.Wrap(err, "failed to shutdown extensions"))
	}

	AppTelemetry.shutdown()

	app.logger.Info("Shutdown complete.")

	if len(errs) != 0 {
		return componenterror.CombineErrors(errs)
	}
	return nil
}

// Start starts the collector according to the command and configuration
// given by the user.
func (app *Application) Start() error {
	// From this point on do not show usage in case of error.
	app.rootCmd.SilenceUsage = true

	return app.rootCmd.Execute()
}

func (app *Application) createMemoryBallast() ([]byte, uint64) {
	ballastSizeMiB := builder.MemBallastSize()
	if ballastSizeMiB > 0 {
		ballastSizeBytes := uint64(ballastSizeMiB) * 1024 * 1024
		ballast := make([]byte, ballastSizeBytes)
		app.logger.Info("Using memory ballast", zap.Int("MiBs", ballastSizeMiB))
		return ballast, ballastSizeBytes
	}
	return nil, 0
}
<|endoftext|>


<commit_before>package service

import (
	"errors"
	"net/http"

	"github.com/NYTimes/gizmo/config"
	"github.com/clawio/data/datacontroller"
	"github.com/clawio/keys"
	"github.com/clawio/sdk"
	"github.com/gorilla/context"
	"github.com/prometheus/client_golang/prometheus"
)

const (
	// userKey can be used to store/retrieve a user ID in a request context.
	userKey contextKey = iota
)

type (
	// contextKey is a type to use as a key for storing data in the request context.
	contextKey int

	// Service implements server.Service and
	// handle all requests to the server.
	Service struct {
		Config         *Config
		SDK            *sdk.SDK
		DataController datacontroller.DataController
	}

	// Config is a struct that holds the
	// configuration for Service
	Config struct {
		Server         *config.Server
		General        *GeneralConfig
		DataController *DataControllerConfig
	}

	// GeneralConfig contains configuration parameters
	// for general parts of the service.
	GeneralConfig struct {
		AuthenticationServiceBaseURL string
		RequestBodyMaxSize           int64
	}

	// DataControllerConfig is a struct that holds
	// configuration parameters for a data controller.
	DataControllerConfig struct {
		Type                       string
		SimpleDataDir              string
		SimpleTempDir              string
		SimpleChecksum             string
		SimpleVerifyClientChecksum bool
	}
)

// New will instantiate and return
// a new Service that implements server.Service.
func New(cfg *Config) (*Service, error) {
	if cfg == nil {
		return nil, errors.New("config is nil")
	}
	if cfg.General == nil {
		return nil, errors.New("config.General is nil")
	}
	if cfg.DataController == nil {
		return nil, errors.New("config.DataController is  nil")
	}

	urls := &sdk.ServiceEndpoints{}
	urls.AuthServiceBaseURL = cfg.General.AuthenticationServiceBaseURL
	s := sdk.New(urls, nil)

	dataController := getDataController(cfg.DataController)
	return &Service{Config: cfg, SDK: s, DataController: dataController}, nil
}

func getDataController(cfg *DataControllerConfig) datacontroller.DataController {
	opts := &datacontroller.SimpleDataControllerOptions{
		DataDir:              cfg.SimpleDataDir,
		TempDir:              cfg.SimpleTempDir,
		Checksum:             cfg.SimpleChecksum,
		VerifyClientChecksum: cfg.SimpleVerifyClientChecksum,
	}
	return datacontroller.NewSimpleDataController(opts)
}

// Prefix returns the string prefix used for all endpoints within
// this service.
func (s *Service) Prefix() string {
	return "/clawio/v1/data"
}

// Middleware provides an http.Handler hook wrapped around all requests.
// In this implementation, we authenticate the request.
func (s *Service) Middleware(h http.Handler) http.Handler {
	return h
}

// Endpoints is a listing of all endpoints available in the Service.
func (s *Service) Endpoints() map[string]map[string]http.HandlerFunc {
	return map[string]map[string]http.HandlerFunc{
		"/metrics": {
			"GET": func(w http.ResponseWriter, r *http.Request) {
				prometheus.Handler().ServeHTTP(w, r)
			},
		},
		"/upload/{path:.*}": {
			"PUT": prometheus.InstrumentHandlerFunc("/upload", s.authenticateHandlerFunc(s.Upload)),
		},
		"/download/{path:.*}": {
			"GET": prometheus.InstrumentHandlerFunc("/download", s.authenticateHandlerFunc(s.Download)),
		},
	}
}

func (s *Service) getTokenFromRequest(r *http.Request) string {
	if t := r.Header.Get("token"); t != "" {
		return t
	}
	return r.URL.Query().Get("token")
}

func (s *Service) authenticateHandlerFunc(handler http.HandlerFunc) http.HandlerFunc {
	return func(w http.ResponseWriter, r *http.Request) {
		token := s.getTokenFromRequest(r)
		user, _, err := s.SDK.Auth.Verify(token)
		if err != nil {
			http.Error(w, http.StatusText(http.StatusUnauthorized), http.StatusUnauthorized)
			return
		}
		context.Set(r, keys.UserKey, user)
		handler(w, r)
	}
}
<commit_msg>Create data and temp dirs on initialisation<commit_after>package service

import (
	"errors"
	"net/http"
	"os"

	"github.com/NYTimes/gizmo/config"
	"github.com/clawio/data/datacontroller"
	"github.com/clawio/keys"
	"github.com/clawio/sdk"
	"github.com/gorilla/context"
	"github.com/prometheus/client_golang/prometheus"
)

type (
	// Service implements server.Service and
	// handle all requests to the server.
	Service struct {
		Config         *Config
		SDK            *sdk.SDK
		DataController datacontroller.DataController
	}

	// Config is a struct that holds the
	// configuration for Service
	Config struct {
		Server         *config.Server
		General        *GeneralConfig
		DataController *DataControllerConfig
	}

	// GeneralConfig contains configuration parameters
	// for general parts of the service.
	GeneralConfig struct {
		AuthenticationServiceBaseURL string
		RequestBodyMaxSize           int64
	}

	// DataControllerConfig is a struct that holds
	// configuration parameters for a data controller.
	DataControllerConfig struct {
		Type                       string
		SimpleDataDir              string
		SimpleTempDir              string
		SimpleChecksum             string
		SimpleVerifyClientChecksum bool
	}
)

// New will instantiate and return
// a new Service that implements server.Service.
func New(cfg *Config) (*Service, error) {
	if cfg == nil {
		return nil, errors.New("config is nil")
	}
	if cfg.General == nil {
		return nil, errors.New("config.General is nil")
	}
	if cfg.DataController == nil {
		return nil, errors.New("config.DataController is  nil")
	}

	urls := &sdk.ServiceEndpoints{}
	urls.AuthServiceBaseURL = cfg.General.AuthenticationServiceBaseURL
	s := sdk.New(urls, nil)

	dataController, err := getDataController(cfg.DataController)
	if err != nil {
		return nil, err
	}
	return &Service{Config: cfg, SDK: s, DataController: dataController}, nil
}

func getDataController(cfg *DataControllerConfig) (datacontroller.DataController, error) {
	opts := &datacontroller.SimpleDataControllerOptions{
		DataDir:              cfg.SimpleDataDir,
		TempDir:              cfg.SimpleTempDir,
		Checksum:             cfg.SimpleChecksum,
		VerifyClientChecksum: cfg.SimpleVerifyClientChecksum,
	}
	// create DataDir and TempDir
	if err := os.MkdirAll(opts.DataDir, 0755); err != nil {
		return nil, err
	}
	if err := os.MkdirAll(opts.TempDir, 0755); err != nil {
		return nil, err
	}
	return datacontroller.NewSimpleDataController(opts), nil
}

// Prefix returns the string prefix used for all endpoints within
// this service.
func (s *Service) Prefix() string {
	return "/clawio/v1/data"
}

// Middleware provides an http.Handler hook wrapped around all requests.
// In this implementation, we authenticate the request.
func (s *Service) Middleware(h http.Handler) http.Handler {
	return h
}

// Endpoints is a listing of all endpoints available in the Service.
func (s *Service) Endpoints() map[string]map[string]http.HandlerFunc {
	return map[string]map[string]http.HandlerFunc{
		"/metrics": {
			"GET": func(w http.ResponseWriter, r *http.Request) {
				prometheus.Handler().ServeHTTP(w, r)
			},
		},
		"/upload/{path:.*}": {
			"PUT": prometheus.InstrumentHandlerFunc("/upload", s.authenticateHandlerFunc(s.Upload)),
		},
		"/download/{path:.*}": {
			"GET": prometheus.InstrumentHandlerFunc("/download", s.authenticateHandlerFunc(s.Download)),
		},
	}
}

func (s *Service) getTokenFromRequest(r *http.Request) string {
	if t := r.Header.Get("token"); t != "" {
		return t
	}
	return r.URL.Query().Get("token")
}

func (s *Service) authenticateHandlerFunc(handler http.HandlerFunc) http.HandlerFunc {
	return func(w http.ResponseWriter, r *http.Request) {
		token := s.getTokenFromRequest(r)
		user, _, err := s.SDK.Auth.Verify(token)
		if err != nil {
			http.Error(w, http.StatusText(http.StatusUnauthorized), http.StatusUnauthorized)
			return
		}
		context.Set(r, keys.UserKey, user)
		handler(w, r)
	}
}
<|endoftext|>


<commit_before>package service

import (
	"io/ioutil"
	"fmt"
	"strings"
	"regexp"
	"github.com/stormcat24/ecs-formation/aws"
	"os"
	"github.com/str1ngs/ansi/color"
	"github.com/aws/aws-sdk-go/service/ecs"
	"github.com/stormcat24/ecs-formation/logger"
	"time"
	"errors"
	"github.com/stormcat24/ecs-formation/util"
)

type TaskWatchStatus int

const (
	WatchContinue TaskWatchStatus = iota
	WatchFinish
	WatchTerminate
)

type ServiceController struct {
	Ecs            *aws.AwsManager
	TargetResource string
	clusters       []Cluster
}

func NewServiceController(ecs *aws.AwsManager, projectDir string, targetResource string) (*ServiceController, error) {

	con := &ServiceController{
		Ecs: ecs,
	}

	clusters, err := con.searchServices(projectDir)
	if err != nil {
		return nil, err
	}

	con.clusters = clusters

	if targetResource != "" {
		con.TargetResource = targetResource
	}

	return con, nil
}

func (self *ServiceController) searchServices(projectDir string) ([]Cluster, error) {

	clusterDir := projectDir + "/service"
	files, err := ioutil.ReadDir(clusterDir)

	clusters := []Cluster{}

	if err != nil {
		return clusters, err
	}

	filePattern := regexp.MustCompile("^(.+)\\.yml$")

	for _, file := range files {
		if !file.IsDir() && strings.HasSuffix(file.Name(), ".yml") {
			content, _ := ioutil.ReadFile(clusterDir + "/" + file.Name())

			tokens := filePattern.FindStringSubmatch(file.Name())
			clusterName := tokens[1]

			serviceMap, _ := CreateServiceMap(content)
			cluster := Cluster{
				Name: clusterName,
				Services: serviceMap,
			}

			clusters = append(clusters, cluster)
		}
	}

	return clusters, nil
}

func (self *ServiceController) GetClusters() []Cluster {
	return self.clusters
}

func (self *ServiceController) CreateServiceUpdatePlans() ([]*ServiceUpdatePlan, error) {

	plans := []*ServiceUpdatePlan{}
	for _, cluster := range self.GetClusters() {
		if len(self.TargetResource) == 0 || self.TargetResource == cluster.Name {
			cp, err := self.CreateServiceUpdatePlan(cluster)
			if err != nil {
				return plans, err
			}

			plans = append(plans, cp)
		}
	}
	return plans, nil
}

func (self *ServiceController) CreateServiceUpdatePlan(cluster Cluster) (*ServiceUpdatePlan, error) {

	clusterApi := self.Ecs.ClusterApi()
	output, errdc := clusterApi.DescribeClusters([]*string{&cluster.Name})

	if errdc != nil {
		return &ServiceUpdatePlan{}, errdc
	}

	if len(output.Failures) > 0 {
		return &ServiceUpdatePlan{}, errors.New(fmt.Sprintf("Cluster '%s' not found", cluster.Name))
	}

	rlci, errlci := clusterApi.ListContainerInstances(cluster.Name)
	if errlci != nil {
		return &ServiceUpdatePlan{}, errlci
	}

	if len(rlci.ContainerInstanceArns) == 0 {
		return &ServiceUpdatePlan{}, errors.New(fmt.Sprintf("ECS instances not found in cluster '%s' not found", cluster.Name))
	}

	target := output.Clusters[0]

	if *target.Status != "ACTIVE" {
		return &ServiceUpdatePlan{}, errors.New(fmt.Sprintf("Cluster '%s' is not ACTIVE.", cluster.Name))
	}

	serviceApi := self.Ecs.ServiceApi()

	resListServices, errls := serviceApi.ListServices(cluster.Name)
	if errls != nil {
		return &ServiceUpdatePlan{}, errls
	}

	currentServices := map[string]*ecs.Service{}
	if len(resListServices.ServiceArns) > 0 {
		resDescribeService, errds := serviceApi.DescribeService(cluster.Name, resListServices.ServiceArns)
		if errds != nil {
			return &ServiceUpdatePlan{}, errds
		}

		for _, service := range resDescribeService.Services {
			currentServices[*service.ServiceName] = service
		}
	}

	newServices := map[string]*Service{}
	for name, newService := range cluster.Services {
		s := newService
		newServices[name] = &s
	}

	return &ServiceUpdatePlan{
		Name: cluster.Name,
		InstanceARNs: rlci.ContainerInstanceArns,
		CurrentServices: currentServices,
		NewServices: newServices,
	}, nil
}

func (self *ServiceController) ApplyServicePlans(plans []*ServiceUpdatePlan) {

	logger.Main.Info("Start apply serivces...")

	for _, plan := range plans {
		if err := self.ApplyServicePlan(plan); err != nil {
			logger.Main.Error(color.Red(err.Error()))
			os.Exit(1)
		}
	}
}

func (self *ServiceController) ApplyServicePlan(plan *ServiceUpdatePlan) error {

	api := self.Ecs.ServiceApi()

	for _, current := range plan.CurrentServices {

		// set desired_count = 0
		if _, err := api.UpdateService(plan.Name, *current.ServiceName, 0, *current.TaskDefinition); err != nil {
			return err
		}

		// wait to stop service
		logger.Main.Infof("Waiting to stop '%s' service on '%s' ...", *current.ServiceName, plan.Name)
		if err := self.waitStoppingService(plan.Name, *current.ServiceName); err != nil {
			return err
		}
		logger.Main.Infof("Stoped '%s' service on '%s'.", *current.ServiceName, plan.Name)

		// delete service
		result, err := api.DeleteService(plan.Name, *current.ServiceArn)
		if err != nil {
			return err
		}

		if err := self.waitStoppingService(plan.Name, *current.ServiceName); err != nil {
			return err
		}

		logger.Main.Infof("Deleted service '%s'", *result.Service.ServiceArn)
	}

	for _, add := range plan.NewServices {

		result, err := api.CreateService(plan.Name, add.Name, add.DesiredCount, toLoadBalancers(&add.LoadBalancers), add.TaskDefinition, add.Role)
		if err != nil {
			return err
		}

		logger.Main.Infof("Created service '%s'", *result.Service.ServiceArn)
		errwas := self.WaitActiveService(plan.Name, add.Name)
		if errwas != nil {
			return errwas
		}
	}

	return nil
}

func toLoadBalancers(values *[]LoadBalancer) []*ecs.LoadBalancer {

	loadBalancers := []*ecs.LoadBalancer{}
	for _, lb := range *values {
		loadBalancers = append(loadBalancers, &ecs.LoadBalancer{
			LoadBalancerName: &lb.Name,
			ContainerName: &lb.ContainerName,
			ContainerPort: &lb.ContainerPort,
		})
	}

	return loadBalancers
}

func (self *ServiceController) waitStoppingService(cluster string, service string) error {

	api := self.Ecs.ServiceApi()

	for {
		time.Sleep(5 * time.Second)

		result, err := api.DescribeService(cluster, []*string{&service})

		if err != nil {
			return err
		}

		if len(result.Services) == 0 {
			return nil
		}

		target := result.Services[0]

		logger.Main.Infof("service '%s@%s' current status = %s", service, cluster, *target.Status)
		if *target.RunningCount == 0 && *target.Status != "DRAINING" {
			return nil
		}

	}
}

func (self *ServiceController) WaitActiveService(cluster string, service string) error {

	api := self.Ecs.ServiceApi()
	taskApi := self.Ecs.TaskApi()

	var flag = false
	var taskARNs []*string

	for {
		time.Sleep(5 * time.Second)

		result, err := api.DescribeService(cluster, []*string{&service})

		if err != nil {
			return err
		}

		if len(result.Services) == 0 {
			continue
		}

		target := result.Services[0]

		// The status of the service. The valid values are ACTIVE, DRAINING, or INACTIVE.
		logger.Main.Infof("service '%s@%s' status = %s ...", service, cluster, *target.Status)

		if *target.Status == "ACTIVE" {

			if len(target.Events) > 0 && strings.Contains(*target.Events[0].Message, "was unable to place a task"){
				return errors.New(*target.Events[0].Message)
			}

			if !flag {
				reslt, errlt := taskApi.ListTasks(cluster, service)
				if errlt != nil {
					return errlt
				}

				if len(reslt.TaskArns) == 0 {
					continue
				} else {
					taskARNs = reslt.TaskArns
					flag = true
				}
			}

			resdt, errdt := taskApi.DescribeTasks(cluster, taskARNs)
			if errdt != nil {
				return errdt
			}

			watchStatus := self.checkRunningTask(resdt)
			if watchStatus == WatchFinish {
				logger.Main.Info("At least one of task has started successfully.")
				return nil
			} else if watchStatus == WatchTerminate {
				logger.Main.Error("Stopped watching task, because task has stopped.")
				return errors.New("Task has been stopped for some reason.")
			}

		}
	}
}

func (self *ServiceController) checkRunningTask(dto *ecs.DescribeTasksOutput) TaskWatchStatus {

	logger.Main.Info("Current task conditions as follows:")

	status := []string{}
	for _, task := range dto.Tasks {
		util.Println(fmt.Sprintf("    %s:", *task.TaskArn))
		util.Println(fmt.Sprintf("        LastStatus:%s", self.RoundColorStatus(*task.LastStatus)))
		util.Println("        Containers:")

		for _, con := range task.Containers {
			util.Println(fmt.Sprintf("            ----------[%s]----------", *con.Name))
			util.Println(fmt.Sprintf("            ContainerARN:%s", *con.ContainerArn))
			util.Println(fmt.Sprintf("            Status:%s", self.RoundColorStatus(*con.LastStatus)))
			util.Println()
		}

		status = append(status, *task.LastStatus)
	}

	// if RUNNING at least one, ecs-formation deals with ok.
	for _, s := range status {
		if s == "RUNNING" {
			return WatchFinish
		} else if s == "STOPPED" {
			return WatchTerminate
		}
	}

	return WatchContinue
}

func (self *ServiceController) RoundColorStatus(status string) *color.Escape {

	if status == "RUNNING" {
		return color.Green(status)
	} else if status == "PENDING" {
		return color.Yellow(status)
	} else if status == "STOPPED" {
		return color.Red(status)
	} else {
		return color.Magenta(status)
	}
}
<commit_msg>no instances in cluster is warn<commit_after>package service

import (
	"io/ioutil"
	"fmt"
	"strings"
	"regexp"
	"github.com/stormcat24/ecs-formation/aws"
	"os"
	"github.com/str1ngs/ansi/color"
	"github.com/aws/aws-sdk-go/service/ecs"
	"github.com/stormcat24/ecs-formation/logger"
	"time"
	"errors"
	"github.com/stormcat24/ecs-formation/util"
)

type TaskWatchStatus int

const (
	WatchContinue TaskWatchStatus = iota
	WatchFinish
	WatchTerminate
)

type ServiceController struct {
	Ecs            *aws.AwsManager
	TargetResource string
	clusters       []Cluster
}

func NewServiceController(ecs *aws.AwsManager, projectDir string, targetResource string) (*ServiceController, error) {

	con := &ServiceController{
		Ecs: ecs,
	}

	clusters, err := con.searchServices(projectDir)
	if err != nil {
		return nil, err
	}

	con.clusters = clusters

	if targetResource != "" {
		con.TargetResource = targetResource
	}

	return con, nil
}

func (self *ServiceController) searchServices(projectDir string) ([]Cluster, error) {

	clusterDir := projectDir + "/service"
	files, err := ioutil.ReadDir(clusterDir)

	clusters := []Cluster{}

	if err != nil {
		return clusters, err
	}

	filePattern := regexp.MustCompile("^(.+)\\.yml$")

	for _, file := range files {
		if !file.IsDir() && strings.HasSuffix(file.Name(), ".yml") {
			content, _ := ioutil.ReadFile(clusterDir + "/" + file.Name())

			tokens := filePattern.FindStringSubmatch(file.Name())
			clusterName := tokens[1]

			serviceMap, _ := CreateServiceMap(content)
			cluster := Cluster{
				Name: clusterName,
				Services: serviceMap,
			}

			clusters = append(clusters, cluster)
		}
	}

	return clusters, nil
}

func (self *ServiceController) GetClusters() []Cluster {
	return self.clusters
}

func (self *ServiceController) CreateServiceUpdatePlans() ([]*ServiceUpdatePlan, error) {

	plans := []*ServiceUpdatePlan{}
	for _, cluster := range self.GetClusters() {
		if len(self.TargetResource) == 0 || self.TargetResource == cluster.Name {
			cp, err := self.CreateServiceUpdatePlan(cluster)
			if err != nil {
				return plans, err
			}

			if cp != nil {
				plans = append(plans, cp)
			}
		}
	}
	return plans, nil
}

func (self *ServiceController) CreateServiceUpdatePlan(cluster Cluster) (*ServiceUpdatePlan, error) {

	clusterApi := self.Ecs.ClusterApi()
	output, errdc := clusterApi.DescribeClusters([]*string{&cluster.Name})

	if errdc != nil {
		return nil, errdc
	}

	if len(output.Failures) > 0 {
		return nil, errors.New(fmt.Sprintf("Cluster '%s' not found", cluster.Name))
	}

	rlci, errlci := clusterApi.ListContainerInstances(cluster.Name)
	if errlci != nil {
		return nil, errlci
	}

	if len(rlci.ContainerInstanceArns) == 0 {
		logger.Main.Warnf("ECS instances not found in cluster '%s' not found", cluster.Name)
		return nil, nil

	} else {
		target := output.Clusters[0]

		if *target.Status != "ACTIVE" {
			return &ServiceUpdatePlan{}, errors.New(fmt.Sprintf("Cluster '%s' is not ACTIVE.", cluster.Name))
		}

		serviceApi := self.Ecs.ServiceApi()

		resListServices, errls := serviceApi.ListServices(cluster.Name)
		if errls != nil {
			return &ServiceUpdatePlan{}, errls
		}

		currentServices := map[string]*ecs.Service{}
		if len(resListServices.ServiceArns) > 0 {
			resDescribeService, errds := serviceApi.DescribeService(cluster.Name, resListServices.ServiceArns)
			if errds != nil {
				return &ServiceUpdatePlan{}, errds
			}

			for _, service := range resDescribeService.Services {
				currentServices[*service.ServiceName] = service
			}
		}

		newServices := map[string]*Service{}
		for name, newService := range cluster.Services {
			s := newService
			newServices[name] = &s
		}

		return &ServiceUpdatePlan{
			Name: cluster.Name,
			InstanceARNs: rlci.ContainerInstanceArns,
			CurrentServices: currentServices,
			NewServices: newServices,
		}, nil
	}
}

func (self *ServiceController) ApplyServicePlans(plans []*ServiceUpdatePlan) {

	logger.Main.Info("Start apply serivces...")

	for _, plan := range plans {
		if err := self.ApplyServicePlan(plan); err != nil {
			logger.Main.Error(color.Red(err.Error()))
			os.Exit(1)
		}
	}
}

func (self *ServiceController) ApplyServicePlan(plan *ServiceUpdatePlan) error {

	api := self.Ecs.ServiceApi()

	for _, current := range plan.CurrentServices {

		// set desired_count = 0
		if _, err := api.UpdateService(plan.Name, *current.ServiceName, 0, *current.TaskDefinition); err != nil {
			return err
		}

		// wait to stop service
		logger.Main.Infof("Waiting to stop '%s' service on '%s' ...", *current.ServiceName, plan.Name)
		if err := self.waitStoppingService(plan.Name, *current.ServiceName); err != nil {
			return err
		}
		logger.Main.Infof("Stoped '%s' service on '%s'.", *current.ServiceName, plan.Name)

		// delete service
		result, err := api.DeleteService(plan.Name, *current.ServiceArn)
		if err != nil {
			return err
		}

		if err := self.waitStoppingService(plan.Name, *current.ServiceName); err != nil {
			return err
		}

		logger.Main.Infof("Deleted service '%s'", *result.Service.ServiceArn)
	}

	for _, add := range plan.NewServices {

		result, err := api.CreateService(plan.Name, add.Name, add.DesiredCount, toLoadBalancers(&add.LoadBalancers), add.TaskDefinition, add.Role)
		if err != nil {
			return err
		}

		logger.Main.Infof("Created service '%s'", *result.Service.ServiceArn)
		errwas := self.WaitActiveService(plan.Name, add.Name)
		if errwas != nil {
			return errwas
		}
	}

	return nil
}

func toLoadBalancers(values *[]LoadBalancer) []*ecs.LoadBalancer {

	loadBalancers := []*ecs.LoadBalancer{}
	for _, lb := range *values {
		loadBalancers = append(loadBalancers, &ecs.LoadBalancer{
			LoadBalancerName: &lb.Name,
			ContainerName: &lb.ContainerName,
			ContainerPort: &lb.ContainerPort,
		})
	}

	return loadBalancers
}

func (self *ServiceController) waitStoppingService(cluster string, service string) error {

	api := self.Ecs.ServiceApi()

	for {
		time.Sleep(5 * time.Second)

		result, err := api.DescribeService(cluster, []*string{&service})

		if err != nil {
			return err
		}

		if len(result.Services) == 0 {
			return nil
		}

		target := result.Services[0]

		logger.Main.Infof("service '%s@%s' current status = %s", service, cluster, *target.Status)
		if *target.RunningCount == 0 && *target.Status != "DRAINING" {
			return nil
		}

	}
}

func (self *ServiceController) WaitActiveService(cluster string, service string) error {

	api := self.Ecs.ServiceApi()
	taskApi := self.Ecs.TaskApi()

	var flag = false
	var taskARNs []*string

	for {
		time.Sleep(5 * time.Second)

		result, err := api.DescribeService(cluster, []*string{&service})

		if err != nil {
			return err
		}

		if len(result.Services) == 0 {
			continue
		}

		target := result.Services[0]

		// The status of the service. The valid values are ACTIVE, DRAINING, or INACTIVE.
		logger.Main.Infof("service '%s@%s' status = %s ...", service, cluster, *target.Status)

		if *target.Status == "ACTIVE" {

			if len(target.Events) > 0 && strings.Contains(*target.Events[0].Message, "was unable to place a task"){
				return errors.New(*target.Events[0].Message)
			}

			if !flag {
				reslt, errlt := taskApi.ListTasks(cluster, service)
				if errlt != nil {
					return errlt
				}

				if len(reslt.TaskArns) == 0 {
					continue
				} else {
					taskARNs = reslt.TaskArns
					flag = true
				}
			}

			resdt, errdt := taskApi.DescribeTasks(cluster, taskARNs)
			if errdt != nil {
				return errdt
			}

			watchStatus := self.checkRunningTask(resdt)
			if watchStatus == WatchFinish {
				logger.Main.Info("At least one of task has started successfully.")
				return nil
			} else if watchStatus == WatchTerminate {
				logger.Main.Error("Stopped watching task, because task has stopped.")
				return errors.New("Task has been stopped for some reason.")
			}

		}
	}
}

func (self *ServiceController) checkRunningTask(dto *ecs.DescribeTasksOutput) TaskWatchStatus {

	logger.Main.Info("Current task conditions as follows:")

	status := []string{}
	for _, task := range dto.Tasks {
		util.Println(fmt.Sprintf("    %s:", *task.TaskArn))
		util.Println(fmt.Sprintf("        LastStatus:%s", self.RoundColorStatus(*task.LastStatus)))
		util.Println("        Containers:")

		for _, con := range task.Containers {
			util.Println(fmt.Sprintf("            ----------[%s]----------", *con.Name))
			util.Println(fmt.Sprintf("            ContainerARN:%s", *con.ContainerArn))
			util.Println(fmt.Sprintf("            Status:%s", self.RoundColorStatus(*con.LastStatus)))
			util.Println()
		}

		status = append(status, *task.LastStatus)
	}

	// if RUNNING at least one, ecs-formation deals with ok.
	for _, s := range status {
		if s == "RUNNING" {
			return WatchFinish
		} else if s == "STOPPED" {
			return WatchTerminate
		}
	}

	return WatchContinue
}

func (self *ServiceController) RoundColorStatus(status string) *color.Escape {

	if status == "RUNNING" {
		return color.Green(status)
	} else if status == "PENDING" {
		return color.Yellow(status)
	} else if status == "STOPPED" {
		return color.Red(status)
	} else {
		return color.Magenta(status)
	}
}
<|endoftext|>