<commit_before>package collaboration

import (
	"fmt"
	"koding/db/mongodb/modelhelper"
	"math/rand"
	apimodels "socialapi/models"
	"socialapi/rest"
	"socialapi/workers/collaboration/models"
	"socialapi/workers/common/runner"
	"testing"
	"time"

	"socialapi/workers/collaboration"
	"socialapi/workers/helper"

	. "github.com/smartystreets/goconvey/convey"
	"labix.org/v2/mgo/bson"
)

var (
	AccountOldId = bson.NewObjectId()
)

func TestCollaborationPing(t *testing.T) {
	r := runner.New("collaboration-tests")
	err := r.Init()
	if err != nil {
		panic(err)
	}

	defer r.Close()

	modelhelper.Initialize(r.Conf.Mongo)
	defer modelhelper.Close()

	redisConn := helper.MustInitRedisConn(r.Conf)
	defer redisConn.Close()

	redis := helper.MustGetRedisConn()

	handler := New(r.Log, redisConn, r.Conf, r.Kite)

	Convey("while pinging collaboration", t, func() {
		// owner
		owner := apimodels.NewAccount()
		owner.OldId = AccountOldId.Hex()
		owner, err := rest.CreateAccount(owner)
		So(err, ShouldBeNil)
		So(owner, ShouldNotBeNil)

		ownerSession, err := apimodels.FetchOrCreateSession(owner.Nick)
		So(err, ShouldBeNil)
		So(ownerSession, ShouldNotBeNil)

		rand.Seed(time.Now().UnixNano())

		req := &models.Ping{
			AccountId: 1,
			FileId:    fmt.Sprintf("%d", rand.Int63()),
		}

		Convey("while testing Ping", func() {
			Convey("reponse should be success with valid ping", func() {
				err := handler.Ping(req)
				So(err, ShouldBeNil)
			})

			Convey("reponse should be success with invalid FileId", func() {
				req.FileId = ""
				err := handler.Ping(req)
				So(err, ShouldBeNil)
			})

			Convey("reponse should be success with invalid AccountId", func() {
				req.AccountId = 0
				err := handler.Ping(req)
				So(err, ShouldBeNil)
			})

			Convey("reponse should be success with invalid session", func() {
				req := req
				// prepare an invalid session here
				req.CreatedAt = time.Now().UTC()
				err := redis.Setex(
					PrepareFileKey(req.FileId),
					collaboration.ExpireSessionKeyDuration, // expire the key after this period
					req.CreatedAt.Add(-terminateSessionDuration),
				)

				err = handler.Ping(req)
				So(err, ShouldBeNil)
			})

			Convey("after sleep time", func() {
				req := req

				Convey("expired session should get invalidSessoin", func() {
					st := sleepTime
					sleepTime = time.Millisecond * 110

					tsd := terminateSessionDuration
					terminateSessionDuration = 100

					// set durations back to the original value
					defer func() {
						sleepTime = st
						terminateSessionDuration = tsd
					}()

					req.CreatedAt = time.Now().UTC()
					// prepare a valid key
					err := redis.Setex(
						PrepareFileKey(req.FileId),
						terminateSessionDuration, // expire the key after this period
						req.CreatedAt.Unix(),     // value - unix time
					)

					// while sleeping here, redis key should be removed
					// and we can understand that the Collab session is expired
					time.Sleep(sleepTime)

					req := req
					err = handler.wait(req)
					So(err, ShouldEqual, errSessionInvalid)
				})

				Convey("deadlined session should get errDeadlineReached", func() {
					st := sleepTime
					sleepTime = time.Millisecond * 110

					dd := deadLineDuration
					deadLineDuration = 100

					// set durations back to the original value
					defer func() {
						sleepTime = st
						deadLineDuration = dd
					}()

					req := req
					err := handler.wait(req)
					So(err, ShouldEqual, errDeadlineReached)
				})
			})
		})

		Convey("while testing checkIfKeyIsValid", func() {

			req := req
			req.CreatedAt = time.Now().UTC()

			// prepare a valid key
			err := redis.Setex(
				PrepareFileKey(req.FileId),
				collaboration.ExpireSessionKeyDuration, // expire the key after this period
				req.CreatedAt.Unix(),                   // value - unix time
			)

			So(err, ShouldBeNil)

			Convey("valid key should return nil", func() {
				err := handler.checkIfKeyIsValid(req)
				So(err, ShouldBeNil)
			})

			Convey("invalid key should return errSessionInvalid", func() {
				req := req
				// override fileId
				req.FileId = fmt.Sprintf("%d", rand.Int63())
				err := handler.checkIfKeyIsValid(req)
				So(err, ShouldEqual, errSessionInvalid)
			})

			Convey("invalid (non-timestamp) value should return errSessionInvalid", func() {
				req := req
				req.CreatedAt = time.Now().UTC()
				err := redis.Setex(
					PrepareFileKey(req.FileId),
					collaboration.ExpireSessionKeyDuration, // expire the key after this period
					"req.CreatedAt.Unix()",                 // replace timestamp with unix time
				)

				err = handler.checkIfKeyIsValid(req)
				So(err, ShouldEqual, errSessionInvalid)
			})

			Convey("old ping time should return errSessionInvalid", func() {
				req := req
				req.CreatedAt = time.Now().UTC()
				err := redis.Setex(
					PrepareFileKey(req.FileId),
					collaboration.ExpireSessionKeyDuration, // expire the key after this period
					req.CreatedAt.Add(-terminateSessionDuration),
				)

				err = handler.checkIfKeyIsValid(req)
				So(err, ShouldEqual, errSessionInvalid)
			})
		})

		Convey("while testing drive operations", func() {
			req := req
			req.CreatedAt = time.Now().UTC()
			Convey("should be able to create the file", func() {
				f, err := handler.createFile()
				So(err, ShouldBeNil)
				req.FileId = f.Id
				Convey("should be able to get the created file", func() {
					f2, err := handler.getFile(f.Id)
					So(err, ShouldBeNil)
					So(f2, ShouldNotBeNil)
					Convey("should be able to delete the created file", func() {
						err = handler.deleteFile(req.FileId)
						So(err, ShouldBeNil)
						Convey("should not be able to get the deleted file", func() {
							f2, err = handler.getFile(f.Id)
							So(err, ShouldNotBeNil)
							So(f2, ShouldBeNil)
						})
						Convey("deleting the deleted file should not give error", func() {
							err = handler.deleteFile(req.FileId)
							So(err, ShouldBeNil)
						})
					})
				})
			})
		})
	})
}
<commit_msg>Socialapi: added better test cases for ping system<commit_after>package collaboration

import (
	"fmt"
	"koding/db/mongodb/modelhelper"
	"math/rand"
	apimodels "socialapi/models"
	"socialapi/rest"
	"socialapi/workers/collaboration/models"
	"socialapi/workers/common/runner"
	"strconv"
	"testing"
	"time"

	"socialapi/workers/collaboration"
	"socialapi/workers/helper"

	"github.com/koding/redis"

	. "github.com/smartystreets/goconvey/convey"
	"labix.org/v2/mgo/bson"
)

var (
	AccountOldId = bson.NewObjectId()
)

func TestCollaboration(t *testing.T) {
	r := runner.New("collaboration-tests")
	err := r.Init()
	if err != nil {
		panic(err)
	}

	defer r.Close()

	modelhelper.Initialize(r.Conf.Mongo)
	defer modelhelper.Close()

	redisConn := helper.MustInitRedisConn(r.Conf)
	defer redisConn.Close()

	redis := helper.MustGetRedisConn()

	handler := New(r.Log, redisConn, r.Conf, r.Kite)

	Convey("while pinging collaboration", t, func() {
		// owner
		owner := apimodels.NewAccount()
		owner.OldId = AccountOldId.Hex()
		owner, err := rest.CreateAccount(owner)
		So(err, ShouldBeNil)
		So(owner, ShouldNotBeNil)

		ownerSession, err := apimodels.FetchOrCreateSession(owner.Nick)
		So(err, ShouldBeNil)
		So(ownerSession, ShouldNotBeNil)

		rand.Seed(time.Now().UnixNano())

		req := &models.Ping{
			AccountId: 1,
			FileId:    fmt.Sprintf("%d", rand.Int63()),
		}

		Convey("while testing Ping", func() {
			Convey("reponse should be success with valid ping", func() {
				err := handler.Ping(req)
				So(err, ShouldBeNil)
			})

			Convey("reponse should be success with invalid FileId", func() {
				req.FileId = ""
				err := handler.Ping(req)
				So(err, ShouldBeNil)
			})

			Convey("reponse should be success with invalid AccountId", func() {
				req.AccountId = 0
				err := handler.Ping(req)
				So(err, ShouldBeNil)
			})

			Convey("reponse should be success with invalid session", func() {
				req := req
				// prepare an invalid session here
				req.CreatedAt = time.Now().UTC()
				err := redis.Setex(
					PrepareFileKey(req.FileId),
					collaboration.ExpireSessionKeyDuration, // expire the key after this period
					req.CreatedAt.Add(-terminateSessionDuration),
				)

				err = handler.Ping(req)
				So(err, ShouldBeNil)
			})

			Convey("after sleep time", func() {
				req := req

				Convey("expired session should get invalidSessoin", func() {
					st := sleepTime
					sleepTime = time.Millisecond * 110

					tsd := terminateSessionDuration
					terminateSessionDuration = 100

					// set durations back to the original value
					defer func() {
						sleepTime = st
						terminateSessionDuration = tsd
					}()

					req.CreatedAt = time.Now().UTC()
					// prepare a valid key
					err := redis.Setex(
						PrepareFileKey(req.FileId),
						terminateSessionDuration, // expire the key after this period
						req.CreatedAt.Unix(),     // value - unix time
					)

					// while sleeping here, redis key should be removed
					// and we can understand that the Collab session is expired
					time.Sleep(sleepTime)

					req := req
					err = handler.wait(req)
					So(err, ShouldEqual, errSessionInvalid)
				})

				Convey("deadlined session should get errDeadlineReached", func() {
					st := sleepTime
					sleepTime = time.Millisecond * 110

					dd := deadLineDuration
					deadLineDuration = 100

					// set durations back to the original value
					defer func() {
						sleepTime = st
						deadLineDuration = dd
					}()

					req := req
					err := handler.wait(req)
					So(err, ShouldEqual, errDeadlineReached)
				})
			})
		})

		Convey("while testing checkIfKeyIsValid", func() {

			req := req
			req.CreatedAt = time.Now().UTC()

			// prepare a valid key
			err := redis.Setex(
				PrepareFileKey(req.FileId),
				collaboration.ExpireSessionKeyDuration, // expire the key after this period
				req.CreatedAt.Unix(),                   // value - unix time
			)

			So(err, ShouldBeNil)

			Convey("valid key should return nil", func() {
				err := handler.checkIfKeyIsValid(req)
				So(err, ShouldBeNil)
			})

			Convey("invalid key should return errSessionInvalid", func() {
				req := req
				// override fileId
				req.FileId = fmt.Sprintf("%d", rand.Int63())
				err := handler.checkIfKeyIsValid(req)
				So(err, ShouldEqual, errSessionInvalid)
			})

			Convey("invalid (non-timestamp) value should return errSessionInvalid", func() {
				req := req
				req.CreatedAt = time.Now().UTC()
				err := redis.Setex(
					PrepareFileKey(req.FileId),
					collaboration.ExpireSessionKeyDuration, // expire the key after this period
					"req.CreatedAt.Unix()",                 // replace timestamp with unix time
				)

				err = handler.checkIfKeyIsValid(req)
				So(err, ShouldEqual, errSessionInvalid)
			})

			Convey("old ping time should return errSessionInvalid", func() {
				req := req
				req.CreatedAt = time.Now().UTC()
				err := redis.Setex(
					PrepareFileKey(req.FileId),
					collaboration.ExpireSessionKeyDuration, // expire the key after this period
					req.CreatedAt.Add(-terminateSessionDuration).Unix(),
				)

				err = handler.checkIfKeyIsValid(req)
				So(err, ShouldEqual, errSessionInvalid)
			})

			Convey("previous ping time is in safe area", func() {
				req := req
				testPingTimes(req, -1, redis, handler, nil)
			})

			Convey("0 ping time is in safe area", func() {
				req := req
				testPingTimes(req, 0, redis, handler, nil)
			})

			Convey("2 ping time is in safe area", func() {
				req := req
				testPingTimes(req, 2, redis, handler, nil)
			})

			Convey("3 ping time is in safe area", func() {
				req := req
				testPingTimes(req, 3, redis, handler, nil)
			})

			Convey("4 ping time is not in safe area - because we already reverted the time ", func() {
				req := req
				testPingTimes(req, 4, redis, handler, errSessionInvalid)
			})

			Convey("5 ping time is not in safe area ", func() {
				req := req
				testPingTimes(req, 5, redis, handler, errSessionInvalid)
			})
		})
	})
}

func testPingTimes(
	req *models.Ping,
	pingCount int,
	redis *redis.RedisSession,
	handler *Controller,
	expectedErr error,
) {
	req.FileId = req.FileId + strconv.Itoa(pingCount)
	req.CreatedAt = time.
		Now().
		UTC().
		Add(-pingDuration * time.Duration(pingCount))

	err := redis.Setex(
		PrepareFileKey(req.FileId),
		collaboration.ExpireSessionKeyDuration, // expire the key after this period
		req.CreatedAt.Unix(),
	)

	err = handler.checkIfKeyIsValid(req)
	So(err, ShouldEqual, expectedErr)
}
<|endoftext|>


<commit_before>package algoliaconnector

import (
	"socialapi/models"
	"strconv"
)

// ChannelCreated handles the channel create events, for now only handles the
// channels that are topic channels,
func (f *Controller) ChannelCreated(data *models.Channel) error {
	if data.TypeConstant != models.Channel_TYPE_TOPIC {
		return nil
	}

	return f.insert(IndexTopics, map[string]interface{}{
		"objectID": strconv.FormatInt(data.Id, 10),
		"name":     data.Name,
		"purpose":  data.Purpose,
	})
}

// ChannelUpdated handles the channel update events, for now only handles the
// channels that are topic channels, we can link channels together in any point
// of time, after linking, leaf channel is removed from search engine. But it is
// still searchable via its root channel, because we are adding it as synonym to
// the root of it
func (f *Controller) ChannelUpdated(data *models.Channel) error {
	if data.TypeConstant != models.Channel_TYPE_LINKED_TOPIC {
		f.log.Debug("unsuported channel for topic update type: %s id: %d", data.TypeConstant, data.Id)
		return nil
	}

	return f.delete(IndexTopics, strconv.FormatInt(data.Id, 10))
}
<commit_msg>Socialapi: added main group initial channel id into all channels for security<commit_after>package algoliaconnector

import (
	"socialapi/models"
	"strconv"
)

// ChannelCreated handles the channel create events, for now only handles the
// channels that are topic channels,
func (f *Controller) ChannelCreated(data *models.Channel) error {
	if data.TypeConstant != models.Channel_TYPE_TOPIC {
		return nil
	}

	// add public group channel id into tags for security
	publicChannel := models.NewChannel()
	err := publicChannel.FetchPublicChannel(data.GroupName)
	if err != nil {
		return err
	}

	return f.insert(IndexTopics, map[string]interface{}{
		"objectID": strconv.FormatInt(data.Id, 10),
		"name":     data.Name,
		"purpose":  data.Purpose,
		"_tags":    []string{strconv.FormatInt(publicChannel.Id, 10)},
	})
}

// ChannelUpdated handles the channel update events, for now only handles the
// channels that are topic channels, we can link channels together in any point
// of time, after linking, leaf channel is removed from search engine. But it is
// still searchable via its root channel, because we are adding it as synonym to
// the root of it
func (f *Controller) ChannelUpdated(data *models.Channel) error {
	if data.TypeConstant != models.Channel_TYPE_LINKED_TOPIC {
		f.log.Debug("unsuported channel for topic update type: %s id: %d", data.TypeConstant, data.Id)
		return nil
	}

	return f.delete(IndexTopics, strconv.FormatInt(data.Id, 10))
}
<|endoftext|>


<commit_before>package webhook

import (
	"net/http"
	"socialapi/config"
	"socialapi/workers/common/mux"
	"socialapi/workers/integration/webhook"
	"testing"

	"github.com/koding/logging"
	"github.com/koding/runner"
	"github.com/rcrowley/go-tigertonic/mocking"
	. "github.com/smartystreets/goconvey/convey"
)

func newRequest(body, channelName string) *webhook.WebhookRequest {
	return &webhook.WebhookRequest{
		Body:        body,
		ChannelName: channelName,
	}
}

func TestWebhookListen(t *testing.T) {
	r := runner.New("test")
	if err := r.Init(); err != nil {
		t.Fatalf("something went wrong: %s", err)
	}
	r.Log.SetLevel(logging.CRITICAL)

	defer r.Close()
	config.MustRead(r.Conf.Path)
	mc := mux.NewConfig("testing", "", "")
	m := mux.New(mc, r.Log)

	h := webhook.NewHandler(r.Log)
	h.AddHandlers(m)

	Convey("while testing incoming webhook", t, func() {

		Convey("users should not be able to send any message when they don't have valid token", func() {
			token := "123123"
			s, _, _, err := h.Push(
				mocking.URL(m, "POST", "/webhook/push/"+token),
				mocking.Header(nil),
				newRequest("hey", "testingchannel"),
			)
			So(err, ShouldNotBeNil)
			So(s, ShouldEqual, http.StatusNotFound)

			token = ""
			s, _, _, err = h.Push(
				mocking.URL(m, "POST", "/webhook/push/"+token),
				mocking.Header(nil),
				newRequest("hey", "testingchannel"),
			)
			So(err.Error(), ShouldEqual, ErrTokenNotSet.Error())
			So(s, ShouldEqual, http.StatusBadRequest)
		})

		Convey("users should not be able to send any message when their request does not include body or channel name", func() {

			token := "123123"
			s, _, _, err := h.Push(
				mocking.URL(m, "POST", "/webhook/push/"+token),
				mocking.Header(nil),
				newRequest("", "testingchannel"),
			)
			So(err.Error(), ShouldEqual, ErrBodyNotSet.Error())
			So(s, ShouldEqual, http.StatusBadRequest)
		})

		Convey("users should be able to send message when token is valid", nil)

	})
}
<commit_msg>webhook: add channel validation test<commit_after>package webhook

import (
	"net/http"
	"socialapi/config"
	"socialapi/workers/common/mux"
	"socialapi/workers/integration/webhook"
	"testing"

	"github.com/koding/logging"
	"github.com/koding/runner"
	"github.com/rcrowley/go-tigertonic/mocking"
	. "github.com/smartystreets/goconvey/convey"
)

func newRequest(body, channelName string) *webhook.WebhookRequest {
	return &webhook.WebhookRequest{
		Body:        body,
		ChannelName: channelName,
	}
}

func TestWebhookListen(t *testing.T) {
	r := runner.New("test")
	if err := r.Init(); err != nil {
		t.Fatalf("something went wrong: %s", err)
	}
	r.Log.SetLevel(logging.CRITICAL)

	defer r.Close()
	config.MustRead(r.Conf.Path)
	mc := mux.NewConfig("testing", "", "")
	m := mux.New(mc, r.Log)

	h := webhook.NewHandler(r.Log)
	h.AddHandlers(m)

	Convey("while testing incoming webhook", t, func() {

		Convey("users should not be able to send any message when they don't have valid token", func() {
			token := "123123"
			s, _, _, err := h.Push(
				mocking.URL(m, "POST", "/webhook/push/"+token),
				mocking.Header(nil),
				newRequest("hey", "testingchannel"),
			)
			So(err, ShouldNotBeNil)
			So(s, ShouldEqual, http.StatusNotFound)

			token = ""
			s, _, _, err = h.Push(
				mocking.URL(m, "POST", "/webhook/push/"+token),
				mocking.Header(nil),
				newRequest("hey", "testingchannel"),
			)
			So(err.Error(), ShouldEqual, ErrTokenNotSet.Error())
			So(s, ShouldEqual, http.StatusBadRequest)
		})

		Convey("users should not be able to send any message when their request does not include body or channel name", func() {

			token := "123123"
			s, _, _, err := h.Push(
				mocking.URL(m, "POST", "/webhook/push/"+token),
				mocking.Header(nil),
				newRequest("", "testingchannel"),
			)
			So(err.Error(), ShouldEqual, ErrBodyNotSet.Error())
			So(s, ShouldEqual, http.StatusBadRequest)

			s, _, _, err = h.Push(
				mocking.URL(m, "POST", "/webhook/push/"+token),
				mocking.Header(nil),
				newRequest("hey", ""),
			)
			So(err.Error(), ShouldEqual, ErrChannelNotSet.Error())
			So(s, ShouldEqual, http.StatusBadRequest)
		})

		Convey("users should be able to send message when token is valid", nil)

	})
}
<|endoftext|>


<commit_before>// Copyright 2014 The Go Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package runtime

import "unsafe"

// Solaris runtime-integrated network poller.
//
// Solaris uses event ports for scalable network I/O. Event
// ports are level-triggered, unlike epoll and kqueue which
// can be configured in both level-triggered and edge-triggered
// mode. Level triggering means we have to keep track of a few things
// ourselves. After we receive an event for a file descriptor,
// it's our responsibility to ask again to be notified for future
// events for that descriptor. When doing this we must keep track of
// what kind of events the goroutines are currently interested in,
// for example a fd may be open both for reading and writing.
//
// A description of the high level operation of this code
// follows. Networking code will get a file descriptor by some means
// and will register it with the netpolling mechanism by a code path
// that eventually calls runtime·netpollopen. runtime·netpollopen
// calls port_associate with an empty event set. That means that we
// will not receive any events at this point. The association needs
// to be done at this early point because we need to process the I/O
// readiness notification at some point in the future. If I/O becomes
// ready when nobody is listening, when we finally care about it,
// nobody will tell us anymore.
//
// Beside calling runtime·netpollopen, the networking code paths
// will call runtime·netpollarm each time goroutines are interested
// in doing network I/O. Because now we know what kind of I/O we
// are interested in (reading/writing), we can call port_associate
// passing the correct type of event set (POLLIN/POLLOUT). As we made
// sure to have already associated the file descriptor with the port,
// when we now call port_associate, we will unblock the main poller
// loop (in runtime·netpoll) right away if the socket is actually
// ready for I/O.
//
// The main poller loop runs in its own thread waiting for events
// using port_getn. When an event happens, it will tell the scheduler
// about it using runtime·netpollready. Besides doing this, it must
// also re-associate the events that were not part of this current
// notification with the file descriptor. Failing to do this would
// mean each notification will prevent concurrent code using the
// same file descriptor in parallel.
//
// The logic dealing with re-associations is encapsulated in
// runtime·netpollupdate. This function takes care to associate the
// descriptor only with the subset of events that were previously
// part of the association, except the one that just happened. We
// can't re-associate with that right away, because event ports
// are level triggered so it would cause a busy loop. Instead, that
// association is effected only by the runtime·netpollarm code path,
// when Go code actually asks for I/O.
//
// The open and arming mechanisms are serialized using the lock
// inside PollDesc. This is required because the netpoll loop runs
// asynchronously in respect to other Go code and by the time we get
// to call port_associate to update the association in the loop, the
// file descriptor might have been closed and reopened already. The
// lock allows runtime·netpollupdate to be called synchronously from
// the loop thread while preventing other threads operating to the
// same PollDesc, so once we unblock in the main loop, until we loop
// again we know for sure we are always talking about the same file
// descriptor and can safely access the data we want (the event set).

//go:cgo_import_dynamic libc_port_create port_create "libc.so"
//go:cgo_import_dynamic libc_port_associate port_associate "libc.so"
//go:cgo_import_dynamic libc_port_dissociate port_dissociate "libc.so"
//go:cgo_import_dynamic libc_port_getn port_getn "libc.so"

//go:linkname libc_port_create libc_port_create
//go:linkname libc_port_associate libc_port_associate
//go:linkname libc_port_dissociate libc_port_dissociate
//go:linkname libc_port_getn libc_port_getn

var (
	libc_port_create,
	libc_port_associate,
	libc_port_dissociate,
	libc_port_getn libcFunc
)

func errno() int32 {
	return *getg().m.perrno
}

func fcntl(fd, cmd int32, arg uintptr) int32 {
	return int32(sysvicall3(&libc_fcntl, uintptr(fd), uintptr(cmd), arg))
}

func port_create() int32 {
	return int32(sysvicall0(&libc_port_create))
}

func port_associate(port, source int32, object uintptr, events uint32, user uintptr) int32 {
	return int32(sysvicall5(&libc_port_associate, uintptr(port), uintptr(source), object, uintptr(events), user))
}

func port_dissociate(port, source int32, object uintptr) int32 {
	return int32(sysvicall3(&libc_port_dissociate, uintptr(port), uintptr(source), object))
}

func port_getn(port int32, evs *portevent, max uint32, nget *uint32, timeout *timespec) int32 {
	return int32(sysvicall5(&libc_port_getn, uintptr(port), uintptr(unsafe.Pointer(evs)), uintptr(max), uintptr(unsafe.Pointer(nget)), uintptr(unsafe.Pointer(timeout))))
}

var portfd int32 = -1

func netpollinit() {
	portfd = port_create()
	if portfd >= 0 {
		fcntl(portfd, _F_SETFD, _FD_CLOEXEC)
		return
	}

	print("runtime: port_create failed (errno=", errno(), ")\n")
	throw("runtime: netpollinit failed")
}

func netpolldescriptor() uintptr {
	return uintptr(portfd)
}

func netpollopen(fd uintptr, pd *pollDesc) int32 {
	lock(&pd.lock)
	// We don't register for any specific type of events yet, that's
	// netpollarm's job. We merely ensure we call port_associate before
	// asynchronous connect/accept completes, so when we actually want
	// to do any I/O, the call to port_associate (from netpollarm,
	// with the interested event set) will unblock port_getn right away
	// because of the I/O readiness notification.
	pd.user = 0
	r := port_associate(portfd, _PORT_SOURCE_FD, fd, 0, uintptr(unsafe.Pointer(pd)))
	unlock(&pd.lock)
	return r
}

func netpollclose(fd uintptr) int32 {
	return port_dissociate(portfd, _PORT_SOURCE_FD, fd)
}

// Updates the association with a new set of interested events. After
// this call, port_getn will return one and only one event for that
// particular descriptor, so this function needs to be called again.
func netpollupdate(pd *pollDesc, set, clear uint32) {
	if pd.closing {
		return
	}

	old := pd.user
	events := (old & ^clear) | set
	if old == events {
		return
	}

	if events != 0 && port_associate(portfd, _PORT_SOURCE_FD, pd.fd, events, uintptr(unsafe.Pointer(pd))) != 0 {
		print("runtime: port_associate failed (errno=", errno(), ")\n")
		throw("runtime: netpollupdate failed")
	}
	pd.user = events
}

// subscribe the fd to the port such that port_getn will return one event.
func netpollarm(pd *pollDesc, mode int) {
	lock(&pd.lock)
	switch mode {
	case 'r':
		netpollupdate(pd, _POLLIN, 0)
	case 'w':
		netpollupdate(pd, _POLLOUT, 0)
	default:
		throw("runtime: bad mode")
	}
	unlock(&pd.lock)
}

// polls for ready network connections
// returns list of goroutines that become runnable
func netpoll(block bool) gList {
	if portfd == -1 {
		return gList{}
	}

	var wait *timespec
	var zero timespec
	if !block {
		wait = &zero
	}

	var events [128]portevent
retry:
	var n uint32 = 1
	if port_getn(portfd, &events[0], uint32(len(events)), &n, wait) < 0 {
		if e := errno(); e != _EINTR {
			print("runtime: port_getn on fd ", portfd, " failed (errno=", e, ")\n")
			throw("runtime: netpoll failed")
		}
		goto retry
	}

	var toRun gList
	for i := 0; i < int(n); i++ {
		ev := &events[i]

		if ev.portev_events == 0 {
			continue
		}
		pd := (*pollDesc)(unsafe.Pointer(ev.portev_user))

		var mode, clear int32
		if (ev.portev_events & (_POLLIN | _POLLHUP | _POLLERR)) != 0 {
			mode += 'r'
			clear |= _POLLIN
		}
		if (ev.portev_events & (_POLLOUT | _POLLHUP | _POLLERR)) != 0 {
			mode += 'w'
			clear |= _POLLOUT
		}
		// To effect edge-triggered events, we need to be sure to
		// update our association with whatever events were not
		// set with the event. For example if we are registered
		// for POLLIN|POLLOUT, and we get POLLIN, besides waking
		// the goroutine interested in POLLIN we have to not forget
		// about the one interested in POLLOUT.
		if clear != 0 {
			lock(&pd.lock)
			netpollupdate(pd, 0, uint32(clear))
			unlock(&pd.lock)
		}

		if mode != 0 {
			pd.everr = false
			if ev.portev_events == _POLLERR {
				pd.everr = true
			}
			netpollready(&toRun, pd, mode)
		}
	}

	if block && toRun.empty() {
		goto retry
	}
	return toRun
}
<commit_msg>runtime: disable event scanning error reporting on solaris<commit_after>// Copyright 2014 The Go Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package runtime

import "unsafe"

// Solaris runtime-integrated network poller.
//
// Solaris uses event ports for scalable network I/O. Event
// ports are level-triggered, unlike epoll and kqueue which
// can be configured in both level-triggered and edge-triggered
// mode. Level triggering means we have to keep track of a few things
// ourselves. After we receive an event for a file descriptor,
// it's our responsibility to ask again to be notified for future
// events for that descriptor. When doing this we must keep track of
// what kind of events the goroutines are currently interested in,
// for example a fd may be open both for reading and writing.
//
// A description of the high level operation of this code
// follows. Networking code will get a file descriptor by some means
// and will register it with the netpolling mechanism by a code path
// that eventually calls runtime·netpollopen. runtime·netpollopen
// calls port_associate with an empty event set. That means that we
// will not receive any events at this point. The association needs
// to be done at this early point because we need to process the I/O
// readiness notification at some point in the future. If I/O becomes
// ready when nobody is listening, when we finally care about it,
// nobody will tell us anymore.
//
// Beside calling runtime·netpollopen, the networking code paths
// will call runtime·netpollarm each time goroutines are interested
// in doing network I/O. Because now we know what kind of I/O we
// are interested in (reading/writing), we can call port_associate
// passing the correct type of event set (POLLIN/POLLOUT). As we made
// sure to have already associated the file descriptor with the port,
// when we now call port_associate, we will unblock the main poller
// loop (in runtime·netpoll) right away if the socket is actually
// ready for I/O.
//
// The main poller loop runs in its own thread waiting for events
// using port_getn. When an event happens, it will tell the scheduler
// about it using runtime·netpollready. Besides doing this, it must
// also re-associate the events that were not part of this current
// notification with the file descriptor. Failing to do this would
// mean each notification will prevent concurrent code using the
// same file descriptor in parallel.
//
// The logic dealing with re-associations is encapsulated in
// runtime·netpollupdate. This function takes care to associate the
// descriptor only with the subset of events that were previously
// part of the association, except the one that just happened. We
// can't re-associate with that right away, because event ports
// are level triggered so it would cause a busy loop. Instead, that
// association is effected only by the runtime·netpollarm code path,
// when Go code actually asks for I/O.
//
// The open and arming mechanisms are serialized using the lock
// inside PollDesc. This is required because the netpoll loop runs
// asynchronously in respect to other Go code and by the time we get
// to call port_associate to update the association in the loop, the
// file descriptor might have been closed and reopened already. The
// lock allows runtime·netpollupdate to be called synchronously from
// the loop thread while preventing other threads operating to the
// same PollDesc, so once we unblock in the main loop, until we loop
// again we know for sure we are always talking about the same file
// descriptor and can safely access the data we want (the event set).

//go:cgo_import_dynamic libc_port_create port_create "libc.so"
//go:cgo_import_dynamic libc_port_associate port_associate "libc.so"
//go:cgo_import_dynamic libc_port_dissociate port_dissociate "libc.so"
//go:cgo_import_dynamic libc_port_getn port_getn "libc.so"

//go:linkname libc_port_create libc_port_create
//go:linkname libc_port_associate libc_port_associate
//go:linkname libc_port_dissociate libc_port_dissociate
//go:linkname libc_port_getn libc_port_getn

var (
	libc_port_create,
	libc_port_associate,
	libc_port_dissociate,
	libc_port_getn libcFunc
)

func errno() int32 {
	return *getg().m.perrno
}

func fcntl(fd, cmd int32, arg uintptr) int32 {
	return int32(sysvicall3(&libc_fcntl, uintptr(fd), uintptr(cmd), arg))
}

func port_create() int32 {
	return int32(sysvicall0(&libc_port_create))
}

func port_associate(port, source int32, object uintptr, events uint32, user uintptr) int32 {
	return int32(sysvicall5(&libc_port_associate, uintptr(port), uintptr(source), object, uintptr(events), user))
}

func port_dissociate(port, source int32, object uintptr) int32 {
	return int32(sysvicall3(&libc_port_dissociate, uintptr(port), uintptr(source), object))
}

func port_getn(port int32, evs *portevent, max uint32, nget *uint32, timeout *timespec) int32 {
	return int32(sysvicall5(&libc_port_getn, uintptr(port), uintptr(unsafe.Pointer(evs)), uintptr(max), uintptr(unsafe.Pointer(nget)), uintptr(unsafe.Pointer(timeout))))
}

var portfd int32 = -1

func netpollinit() {
	portfd = port_create()
	if portfd >= 0 {
		fcntl(portfd, _F_SETFD, _FD_CLOEXEC)
		return
	}

	print("runtime: port_create failed (errno=", errno(), ")\n")
	throw("runtime: netpollinit failed")
}

func netpolldescriptor() uintptr {
	return uintptr(portfd)
}

func netpollopen(fd uintptr, pd *pollDesc) int32 {
	lock(&pd.lock)
	// We don't register for any specific type of events yet, that's
	// netpollarm's job. We merely ensure we call port_associate before
	// asynchronous connect/accept completes, so when we actually want
	// to do any I/O, the call to port_associate (from netpollarm,
	// with the interested event set) will unblock port_getn right away
	// because of the I/O readiness notification.
	pd.user = 0
	r := port_associate(portfd, _PORT_SOURCE_FD, fd, 0, uintptr(unsafe.Pointer(pd)))
	unlock(&pd.lock)
	return r
}

func netpollclose(fd uintptr) int32 {
	return port_dissociate(portfd, _PORT_SOURCE_FD, fd)
}

// Updates the association with a new set of interested events. After
// this call, port_getn will return one and only one event for that
// particular descriptor, so this function needs to be called again.
func netpollupdate(pd *pollDesc, set, clear uint32) {
	if pd.closing {
		return
	}

	old := pd.user
	events := (old & ^clear) | set
	if old == events {
		return
	}

	if events != 0 && port_associate(portfd, _PORT_SOURCE_FD, pd.fd, events, uintptr(unsafe.Pointer(pd))) != 0 {
		print("runtime: port_associate failed (errno=", errno(), ")\n")
		throw("runtime: netpollupdate failed")
	}
	pd.user = events
}

// subscribe the fd to the port such that port_getn will return one event.
func netpollarm(pd *pollDesc, mode int) {
	lock(&pd.lock)
	switch mode {
	case 'r':
		netpollupdate(pd, _POLLIN, 0)
	case 'w':
		netpollupdate(pd, _POLLOUT, 0)
	default:
		throw("runtime: bad mode")
	}
	unlock(&pd.lock)
}

// polls for ready network connections
// returns list of goroutines that become runnable
func netpoll(block bool) gList {
	if portfd == -1 {
		return gList{}
	}

	var wait *timespec
	var zero timespec
	if !block {
		wait = &zero
	}

	var events [128]portevent
retry:
	var n uint32 = 1
	if port_getn(portfd, &events[0], uint32(len(events)), &n, wait) < 0 {
		if e := errno(); e != _EINTR {
			print("runtime: port_getn on fd ", portfd, " failed (errno=", e, ")\n")
			throw("runtime: netpoll failed")
		}
		goto retry
	}

	var toRun gList
	for i := 0; i < int(n); i++ {
		ev := &events[i]

		if ev.portev_events == 0 {
			continue
		}
		pd := (*pollDesc)(unsafe.Pointer(ev.portev_user))

		var mode, clear int32
		if (ev.portev_events & (_POLLIN | _POLLHUP | _POLLERR)) != 0 {
			mode += 'r'
			clear |= _POLLIN
		}
		if (ev.portev_events & (_POLLOUT | _POLLHUP | _POLLERR)) != 0 {
			mode += 'w'
			clear |= _POLLOUT
		}
		// To effect edge-triggered events, we need to be sure to
		// update our association with whatever events were not
		// set with the event. For example if we are registered
		// for POLLIN|POLLOUT, and we get POLLIN, besides waking
		// the goroutine interested in POLLIN we have to not forget
		// about the one interested in POLLOUT.
		if clear != 0 {
			lock(&pd.lock)
			netpollupdate(pd, 0, uint32(clear))
			unlock(&pd.lock)
		}

		if mode != 0 {
			// TODO(mikio): Consider implementing event
			// scanning error reporting once we are sure
			// about the event port on SmartOS.
			//
			// See golang.org/x/issue/30840.
			netpollready(&toRun, pd, mode)
		}
	}

	if block && toRun.empty() {
		goto retry
	}
	return toRun
}
<|endoftext|>


<commit_before>// Copyright 2017 The Go Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

// +build linux
// +build 386 amd64 arm arm64 ppc64 ppc64le

package runtime_test

import (
	"testing"
	"time"
	_ "unsafe"
)

// These tests are a little risky because they overwrite the vdsoClockgettimeSym value.
// It's normally initialized at startup and remains unchanged after that.

//go:linkname vdsoClockgettimeSym runtime.vdsoClockgettimeSym
var vdsoClockgettimeSym uintptr

func TestClockVDSOAndFallbackPaths(t *testing.T) {
	// Check that we can call walltime() and nanotime() with and without their (1st) fast-paths.
	// This just checks that fast and fallback paths can be called, rather than testing their
	// results.
	//
	// Call them indirectly via time.Now(), so we don't need auxiliary .s files to allow us to
	// use go:linkname to refer to the functions directly.

	save := vdsoClockgettimeSym
	if save == 0 {
		t.Log("vdsoClockgettime symbol not found; fallback path will be used by default")
	}

	// Call with fast-path enabled (if vDSO symbol found at startup)
	time.Now()

	// Call with fast-path disabled
	vdsoClockgettimeSym = 0
	time.Now()
	vdsoClockgettimeSym = save
}

func BenchmarkClockVDSOAndFallbackPaths(b *testing.B) {
	run := func(b *testing.B) {
		for i := 0; i < b.N; i++ {
			// Call via time.Now() - see comment in test above.
			time.Now()
		}
	}

	save := vdsoClockgettimeSym
	b.Run("vDSO", run)
	vdsoClockgettimeSym = 0
	b.Run("Fallback", run)
	vdsoClockgettimeSym = save
}

func BenchmarkTimeNow(b *testing.B) {
	for i := 0; i < b.N; i++ {
		time.Now()
	}
}
<commit_msg>runtime: remove VDSO fallback test and benchmarks<commit_after><|endoftext|>


<commit_before>// Copyright 2015 Google Inc. All Rights Reserved.
//
// 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 samples

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

	"github.com/googlecloudplatform/gcsfuse/timeutil"
	"github.com/jacobsa/fuse"
	"github.com/jacobsa/ogletest"
	"golang.org/x/net/context"
)

// A struct that implements common behavior needed by tests in the samples/
// directory. Use it as an embedded field in your test fixture, calling its
// SetUp method from your SetUp method after setting the FileSystem field.
type SampleTest struct {
	// The file system under test and the configuration with which it should be
	// mounted. These must be set by the user of this type before calling SetUp;
	// all the other fields below are set by SetUp itself.
	FileSystem  fuse.FileSystem
	MountConfig fuse.MountConfig

	// A context object that can be used for long-running operations.
	Ctx context.Context

	// A clock with a fixed initial time. The test's set up method may use this
	// to wire the file system with a clock, if desired.
	Clock timeutil.SimulatedClock

	// The directory at which the file system is mounted.
	Dir string

	// Anothing non-nil in this slice will be closed by TearDown. The test will
	// fail if closing fails.
	ToClose []io.Closer

	mfs *fuse.MountedFileSystem
}

// Mount t.FileSystem and initialize the other exported fields of the struct.
// Panics on error.
//
// REQUIRES: t.FileSystem has been set.
func (t *SampleTest) SetUp(ti *ogletest.TestInfo) {
	err := t.initialize(t.FileSystem, &t.MountConfig)
	if err != nil {
		panic(err)
	}
}

// Like SetUp, but doens't panic.
func (t *SampleTest) initialize(
	fs fuse.FileSystem,
	config *fuse.MountConfig) (err error) {
	// Initialize the context.
	t.Ctx = context.Background()

	// Initialize the clock.
	t.Clock.SetTime(time.Date(2012, 8, 15, 22, 56, 0, 0, time.Local))

	// Set up a temporary directory.
	t.Dir, err = ioutil.TempDir("", "sample_test")
	if err != nil {
		err = fmt.Errorf("TempDir: %v", err)
		return
	}

	// Mount the file system.
	t.mfs, err = fuse.Mount(t.Dir, fs, config)
	if err != nil {
		err = fmt.Errorf("Mount: %v", err)
		return
	}

	// Wait for it to be read.
	err = t.mfs.WaitForReady(t.Ctx)
	if err != nil {
		err = fmt.Errorf("WaitForReady: %v", err)
		return
	}

	return
}

// Unmount the file system and clean up. Panics on error.
func (t *SampleTest) TearDown() {
	err := t.destroy()
	if err != nil {
		panic(err)
	}
}

// Like TearDown, but doesn't panic.
func (t *SampleTest) destroy() (err error) {
	// Close what is necessary.
	for _, c := range t.ToClose {
		if c == nil {
			continue
		}

		ogletest.ExpectEq(nil, c.Close())
	}

	// Was the file system mounted?
	if t.mfs == nil {
		return
	}

	// Unmount the file system. Try again on "resource busy" errors.
	delay := 10 * time.Millisecond
	for {
		err = t.mfs.Unmount()
		if err == nil {
			break
		}

		if strings.Contains(err.Error(), "resource busy") {
			log.Println("Resource busy error while unmounting; trying again")
			time.Sleep(delay)
			delay = time.Duration(1.3 * float64(delay))
			continue
		}

		err = fmt.Errorf("MountedFileSystem.Unmount: %v", err)
		return
	}

	if err = t.mfs.Join(t.Ctx); err != nil {
		err = fmt.Errorf("MountedFileSystem.Join: %v", err)
		return
	}

	return
}
<commit_msg>Fixed a comment.<commit_after>// Copyright 2015 Google Inc. All Rights Reserved.
//
// 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 samples

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

	"github.com/googlecloudplatform/gcsfuse/timeutil"
	"github.com/jacobsa/fuse"
	"github.com/jacobsa/ogletest"
	"golang.org/x/net/context"
)

// A struct that implements common behavior needed by tests in the samples/
// directory. Use it as an embedded field in your test fixture, calling its
// SetUp method from your SetUp method after setting the FileSystem field.
type SampleTest struct {
	// The file system under test and the configuration with which it should be
	// mounted. These must be set by the user of this type before calling SetUp;
	// all the other fields below are set by SetUp itself.
	FileSystem  fuse.FileSystem
	MountConfig fuse.MountConfig

	// A context object that can be used for long-running operations.
	Ctx context.Context

	// A clock with a fixed initial time. The test's set up method may use this
	// to wire the file system with a clock, if desired.
	Clock timeutil.SimulatedClock

	// The directory at which the file system is mounted.
	Dir string

	// Anothing non-nil in this slice will be closed by TearDown. The test will
	// fail if closing fails.
	ToClose []io.Closer

	mfs *fuse.MountedFileSystem
}

// Mount t.FileSystem and initialize the other exported fields of the struct.
// Panics on error.
//
// REQUIRES: t.FileSystem has been set.
func (t *SampleTest) SetUp(ti *ogletest.TestInfo) {
	err := t.initialize(t.FileSystem, &t.MountConfig)
	if err != nil {
		panic(err)
	}
}

// Like SetUp, but doens't panic.
func (t *SampleTest) initialize(
	fs fuse.FileSystem,
	config *fuse.MountConfig) (err error) {
	// Initialize the context.
	t.Ctx = context.Background()

	// Initialize the clock.
	t.Clock.SetTime(time.Date(2012, 8, 15, 22, 56, 0, 0, time.Local))

	// Set up a temporary directory.
	t.Dir, err = ioutil.TempDir("", "sample_test")
	if err != nil {
		err = fmt.Errorf("TempDir: %v", err)
		return
	}

	// Mount the file system.
	t.mfs, err = fuse.Mount(t.Dir, fs, config)
	if err != nil {
		err = fmt.Errorf("Mount: %v", err)
		return
	}

	// Wait for it to be ready.
	err = t.mfs.WaitForReady(t.Ctx)
	if err != nil {
		err = fmt.Errorf("WaitForReady: %v", err)
		return
	}

	return
}

// Unmount the file system and clean up. Panics on error.
func (t *SampleTest) TearDown() {
	err := t.destroy()
	if err != nil {
		panic(err)
	}
}

// Like TearDown, but doesn't panic.
func (t *SampleTest) destroy() (err error) {
	// Close what is necessary.
	for _, c := range t.ToClose {
		if c == nil {
			continue
		}

		ogletest.ExpectEq(nil, c.Close())
	}

	// Was the file system mounted?
	if t.mfs == nil {
		return
	}

	// Unmount the file system. Try again on "resource busy" errors.
	delay := 10 * time.Millisecond
	for {
		err = t.mfs.Unmount()
		if err == nil {
			break
		}

		if strings.Contains(err.Error(), "resource busy") {
			log.Println("Resource busy error while unmounting; trying again")
			time.Sleep(delay)
			delay = time.Duration(1.3 * float64(delay))
			continue
		}

		err = fmt.Errorf("MountedFileSystem.Unmount: %v", err)
		return
	}

	if err = t.mfs.Join(t.Ctx); err != nil {
		err = fmt.Errorf("MountedFileSystem.Join: %v", err)
		return
	}

	return
}
<|endoftext|>


<commit_before>/*
Copyright 2017 by the contributors.

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

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

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

package main

import (
	"errors"
	"fmt"
	"os"

	"github.com/kubernetes-sigs/aws-iam-authenticator/pkg/config"

	"github.com/sirupsen/logrus"
	"github.com/spf13/cobra"
	"github.com/spf13/viper"
)

var cfgFile string

var rootCmd = &cobra.Command{
	Use:   "aws-iam-authenticator",
	Short: "A tool to authenticate to Kubernetes using AWS IAM credentials",
}

func main() {
	Execute()
}

// Execute the CLI entrypoint
func Execute() {
	if err := rootCmd.Execute(); err != nil {
		fmt.Println(err)
		os.Exit(1)
	}
}

func init() {
	logrus.SetFormatter(&logrus.TextFormatter{FullTimestamp: true})
	cobra.OnInitialize(initConfig)
	rootCmd.PersistentFlags().StringVarP(&cfgFile, "config", "c", "", "Load configuration from `filename`")

	rootCmd.PersistentFlags().StringP(
		"cluster-id",
		"i",
		"",
		"Specify the cluster `ID`, a unique-per-cluster identifier for your aws-iam-authenticator installation.",
	)
	viper.BindPFlag("clusterID", rootCmd.PersistentFlags().Lookup("cluster-id"))
	viper.BindEnv("clusterID", "KUBERNETES_AWS_AUTHENTICATOR_CLUSTER_ID")
}

func initConfig() {
	if cfgFile == "" {
		return
	}
	viper.SetConfigFile(cfgFile)
	if err := viper.ReadInConfig(); err != nil {
		fmt.Printf("Can't read configuration file %q: %v\n", cfgFile, err)
		os.Exit(1)
	}
}

func getConfig() (config.Config, error) {
	config := config.Config{
		ClusterID:                         viper.GetString("clusterID"),
		ServerEC2DescribeInstancesRoleARN: viper.GetString("server.ec2DescribeInstancesRoleARN"),
		LocalhostPort:                     viper.GetInt("server.port"),
		GenerateKubeconfigPath:            viper.GetString("server.generateKubeconfig"),
		KubeconfigPregenerated:            viper.GetBool("server.kubeconfigPregenerated"),
		StateDir:                          viper.GetString("server.stateDir"),
	}
	if err := viper.UnmarshalKey("server.mapRoles", &config.RoleMappings); err != nil {
		return config, fmt.Errorf("invalid server role mappings: %v", err)
	}
	if err := viper.UnmarshalKey("server.mapUsers", &config.UserMappings); err != nil {
		logrus.WithError(err).Fatal("invalid server user mappings")
	}
	if err := viper.UnmarshalKey("server.mapAccounts", &config.AutoMappedAWSAccounts); err != nil {
		logrus.WithError(err).Fatal("invalid server account mappings")
	}

	if config.ClusterID == "" {
		return config, errors.New("cluster ID cannot be empty")
	}

	return config, nil
}
<commit_msg>Add --log-format server flag to configure output format<commit_after>/*
Copyright 2017 by the contributors.

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

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

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

package main

import (
	"errors"
	"fmt"
	"os"

	"github.com/kubernetes-sigs/aws-iam-authenticator/pkg/config"

	"github.com/sirupsen/logrus"
	"github.com/spf13/cobra"
	"github.com/spf13/viper"
)

var cfgFile string

var rootCmd = &cobra.Command{
	Use:   "aws-iam-authenticator",
	Short: "A tool to authenticate to Kubernetes using AWS IAM credentials",
}

func main() {
	Execute()
}

// Execute the CLI entrypoint
func Execute() {
	if err := rootCmd.Execute(); err != nil {
		fmt.Println(err)
		os.Exit(1)
	}
}

func init() {
	cobra.OnInitialize(initConfig)
	rootCmd.PersistentFlags().StringVarP(&cfgFile, "config", "c", "", "Load configuration from `filename`")

	rootCmd.PersistentFlags().StringP("log-format", "l", "text", "Specify log format to use when logging to stderr [text or json]")

	rootCmd.PersistentFlags().StringP(
		"cluster-id",
		"i",
		"",
		"Specify the cluster `ID`, a unique-per-cluster identifier for your aws-iam-authenticator installation.",
	)
	viper.BindPFlag("clusterID", rootCmd.PersistentFlags().Lookup("cluster-id"))
	viper.BindEnv("clusterID", "KUBERNETES_AWS_AUTHENTICATOR_CLUSTER_ID")
}

func initConfig() {
	logrus.SetFormatter(getLogFormatter())
	if cfgFile == "" {
		return
	}
	viper.SetConfigFile(cfgFile)
	if err := viper.ReadInConfig(); err != nil {
		fmt.Printf("Can't read configuration file %q: %v\n", cfgFile, err)
		os.Exit(1)
	}
}

func getConfig() (config.Config, error) {
	config := config.Config{
		ClusterID:                         viper.GetString("clusterID"),
		ServerEC2DescribeInstancesRoleARN: viper.GetString("server.ec2DescribeInstancesRoleARN"),
		LocalhostPort:                     viper.GetInt("server.port"),
		GenerateKubeconfigPath:            viper.GetString("server.generateKubeconfig"),
		KubeconfigPregenerated:            viper.GetBool("server.kubeconfigPregenerated"),
		StateDir:                          viper.GetString("server.stateDir"),
	}
	if err := viper.UnmarshalKey("server.mapRoles", &config.RoleMappings); err != nil {
		return config, fmt.Errorf("invalid server role mappings: %v", err)
	}
	if err := viper.UnmarshalKey("server.mapUsers", &config.UserMappings); err != nil {
		logrus.WithError(err).Fatal("invalid server user mappings")
	}
	if err := viper.UnmarshalKey("server.mapAccounts", &config.AutoMappedAWSAccounts); err != nil {
		logrus.WithError(err).Fatal("invalid server account mappings")
	}

	if config.ClusterID == "" {
		return config, errors.New("cluster ID cannot be empty")
	}

	return config, nil
}

func getLogFormatter() logrus.Formatter {
	format, _ := rootCmd.PersistentFlags().GetString("log-format")

	if format == "json" {
		return &logrus.JSONFormatter{}
	} else if format != "text" {
		logrus.Warnf("Unknown log format specified (%s), will use default text formatter instead.", format)
	}

	return &logrus.TextFormatter{FullTimestamp: true}
}
<|endoftext|>


<commit_before>package gateway

import (
	"fmt"
	"net"
	"net/http"
	"runtime"
	"strings"
	"sync"
	"sync/atomic"
	"time"

	"github.com/funkygao/httprouter"
	log "github.com/funkygao/log4go"
)

type webServer struct {
	gw *Gateway

	name       string
	maxClients int

	router *httprouter.Router

	httpListener net.Listener
	httpServer   *http.Server

	httpsListener net.Listener
	httpsServer   *http.Server

	waitExitFunc    waitExitFunc
	connStateFunc   connStateFunc
	onConnNewFunc   onConnNewFunc
	onConnCloseFunc onConnCloseFunc

	onStop        func()
	mu            sync.Mutex
	waiterStarted bool

	// FIXME if http/https listener both enabled, must able to tell them apart
	activeConnN int32

	// TODO channel performance is frustrating, no better than mutex/map use ring buffer
	stateIdleCh, stateRemoveCh, stateActiveCh chan net.Conn

	closed chan struct{}
}

func newWebServer(name string, httpAddr, httpsAddr string, maxClients int,
	gw *Gateway) *webServer {
	const initialConnBuckets = 200
	this := &webServer{
		name:          name,
		gw:            gw,
		maxClients:    maxClients,
		stateActiveCh: make(chan net.Conn, initialConnBuckets),
		stateIdleCh:   make(chan net.Conn, initialConnBuckets),
		stateRemoveCh: make(chan net.Conn, initialConnBuckets),
		router:        httprouter.New(),
		closed:        make(chan struct{}),
	}

	if Options.EnableHttpPanicRecover {
		this.router.PanicHandler = func(w http.ResponseWriter, r *http.Request, err interface{}) {
			log.Error("PANIC %s %s(%s) %s %s: %+v", this.name, r.RemoteAddr, getHttpRemoteIp(r), r.Method, r.RequestURI, err)

			writeServerError(w, http.StatusText(http.StatusInternalServerError))
		}
	}

	if httpAddr != "" {
		this.httpServer = &http.Server{
			Addr:           httpAddr,
			Handler:        this.router,
			ReadTimeout:    Options.HttpReadTimeout,
			WriteTimeout:   Options.HttpWriteTimeout,
			MaxHeaderBytes: Options.HttpHeaderMaxBytes,
		}
	}

	if httpsAddr != "" {
		this.httpsServer = &http.Server{
			Addr:           httpsAddr,
			Handler:        this.router,
			ReadTimeout:    Options.HttpReadTimeout,
			WriteTimeout:   Options.HttpWriteTimeout,
			MaxHeaderBytes: Options.HttpHeaderMaxBytes,
		}
	}

	return this
}

func (this *webServer) Router() *httprouter.Router {
	return this.router
}

func (this *webServer) Start() {
	if this.waitExitFunc == nil {
		this.waitExitFunc = this.defaultWaitExit
	}
	if this.connStateFunc == nil {
		this.connStateFunc = this.defaultConnStateMachine

		go this.manageIdleConns()
	}

	if this.httpsServer != nil {
		this.httpsServer.ConnState = this.connStateFunc
		this.startServer(true)
	}

	if this.httpServer != nil {
		this.httpServer.ConnState = this.connStateFunc
		this.startServer(false)
	}

}

func (this *webServer) startServer(https bool) {
	var err error
	waitListenerUp := make(chan struct{})
	go func() {
		if Options.CpuAffinity {
			runtime.LockOSThread()
		}

		var (
			retryDelay         time.Duration
			theListener        net.Listener
			waitListenerUpOnce sync.Once
		)
		for {
			if https {
				this.httpsListener, err = net.Listen("tcp", this.httpsServer.Addr)
				if err != nil {
					if strings.HasSuffix(err.Error(), "address already in use") {
						// non-retriable error encountered
						panic(fmt.Errorf("%s listener: %v", this.name, err))
					}

					if retryDelay == 0 {
						retryDelay = 50 * time.Millisecond
					} else {
						retryDelay = 2 * retryDelay
					}
					if maxDelay := time.Second; retryDelay > maxDelay {
						retryDelay = maxDelay
					}
					log.Error("%s listener %v, retry in %v", this.name, err, retryDelay)
					time.Sleep(retryDelay)
					continue
				}

				theListener, err = setupHttpsListener(this.httpsListener, this.gw.certFile, this.gw.keyFile)
				if err != nil {
					panic(err)
				}
			} else {
				theListener, err = net.Listen("tcp", this.httpServer.Addr)
			}

			if err != nil {
				if retryDelay == 0 {
					retryDelay = 50 * time.Millisecond
				} else {
					retryDelay = 2 * retryDelay
				}
				if maxDelay := time.Second; retryDelay > maxDelay {
					retryDelay = maxDelay
				}
				log.Error("%s listener %v, retry in %v", this.name, err, retryDelay)
				time.Sleep(retryDelay)
				continue
			}

			theListener = LimitListener(this.name, this.gw, theListener, this.maxClients)
			waitListenerUpOnce.Do(func() {
				close(waitListenerUp)
			})

			// on non-temporary err, net/http will close the listener
			if https {
				this.mu.Lock()
				this.httpsListener = theListener
				this.mu.Unlock()

				err = this.httpsServer.Serve(theListener)
			} else {
				this.mu.Lock()
				this.httpListener = theListener
				this.mu.Unlock()

				err = this.httpServer.Serve(theListener)
			}

			select {
			case <-this.gw.shutdownCh:
				return

			default:
				log.Error("%s server: %v", this.name, err)
			}
		}
	}()

	// wait for the listener up
	<-waitListenerUp

	this.mu.Lock()
	if !this.waiterStarted {
		this.waiterStarted = true

		this.gw.wg.Add(1)
		go this.waitExitFunc(this.gw.shutdownCh)
	}
	this.mu.Unlock()

	if https {
		log.Info("%s https server ready on %s", this.name, this.httpsServer.Addr)
	} else {
		log.Info("%s http server ready on %s", this.name, this.httpServer.Addr)
	}
}

func (this *webServer) defaultConnStateMachine(c net.Conn, cs http.ConnState) {
	switch cs {
	case http.StateNew:
		atomic.AddInt32(&this.activeConnN, 1)

		if this.onConnNewFunc != nil {
			this.onConnNewFunc(c)
		}

	case http.StateIdle:
		this.stateIdleCh <- c

	case http.StateActive:
		this.stateActiveCh <- c

	case http.StateClosed, http.StateHijacked:
		atomic.AddInt32(&this.activeConnN, -1)
		this.stateRemoveCh <- c

		if this.onConnCloseFunc != nil {
			this.onConnCloseFunc(c)
		}
	}
}

func (this *webServer) manageIdleConns() {
	var (
		idleConns     = make(map[net.Conn]struct{}, 200)
		c             net.Conn
		waitNextRound = make(chan struct{}, 10)
	)
	defer close(waitNextRound)

	log.Debug("%s is managing idle connections", this.name)

	for {
		select {
		case <-this.gw.shutdownCh:
			if len(idleConns) == 0 {
				// happy ending
				log.Debug("%s closed all idle conns", this.name)
				return
			}

			t := time.Now().Add(time.Millisecond * 100)
			for conn := range idleConns {
				if conn == nil {
					continue
				}

				log.Debug("%s closing %s", this.name, conn.RemoteAddr())
				conn.SetDeadline(t)
			}

			// wait for next loop
			waitNextRound <- struct{}{}

		case <-waitNextRound:
			if len(idleConns) == 0 {
				log.Debug("%s closed all idle conns", this.name)
				return
			}

			t := time.Now().Add(time.Millisecond * 100)
			for conn := range idleConns {
				if conn == nil {
					continue
				}

				log.Debug("%s closing %s", this.name, conn.RemoteAddr())
				conn.SetDeadline(t)
			}

			// wait for next loop
			waitNextRound <- struct{}{}

		case c = <-this.stateActiveCh:
			delete(idleConns, c)

		case c = <-this.stateIdleCh:
			idleConns[c] = struct{}{}

		case c = <-this.stateRemoveCh:
			delete(idleConns, c)
		}
	}
}

func (this *webServer) defaultWaitExit(exit <-chan struct{}) {
	log.Debug("%s enter default wait exit", this.name)

	<-exit

	var err error
	if this.httpServer != nil {
		// HTTP response will have "Connection: close"
		this.httpServer.SetKeepAlivesEnabled(false)

		// avoid new connections
		this.mu.Lock()
		if this.httpListener != nil {
			err = this.httpListener.Close()
		}
		this.mu.Unlock()

		if err != nil {
			log.Error("%s on %s: %+v", this.name, this.httpServer.Addr, err)
		}

		log.Trace("%s on %s listener closed", this.name, this.httpServer.Addr)
	}

	if this.httpsServer != nil {
		// HTTP response will have "Connection: close"
		this.httpsServer.SetKeepAlivesEnabled(false)

		// avoid new connections
		this.mu.Lock()
		if this.httpsListener != nil {
			err = this.httpsListener.Close()
		}
		this.mu.Unlock()

		if err != nil {
			log.Error("%s on %s: %+v", this.name, this.httpsServer.Addr, err)
		}

		log.Trace("%s on %s listener closed", this.name, this.httpsServer.Addr)
	}

	// wait for all established http/https conns close
	waitStart := time.Now()
	var prompt sync.Once
	for {
		activeConnN := atomic.LoadInt32(&this.activeConnN)
		if activeConnN == 0 {
			// good luck, all connections finished
			break
		}

		prompt.Do(func() {
			log.Trace("%s waiting for %d clients shutdown...", this.name, activeConnN)
		})

		// timeout mechanism
		if time.Since(waitStart) > Options.SubTimeout+time.Second {
			log.Warn("%s still left %d conns after %s, forced to shutdown",
				this.name, activeConnN, Options.SubTimeout)
			break
		}

		time.Sleep(time.Millisecond * 50)
	}
	log.Trace("%s all connections finished", this.name)

	if this.httpsServer != nil {
		this.httpsServer.ConnState = nil
	}
	if this.httpServer != nil {
		this.httpServer.ConnState = nil
	}

	// TODO close will lead to race condition
	//close(this.stateActiveCh)
	//close(this.stateIdleCh)
	//close(this.stateRemoveCh)

	if this.onStop != nil {
		this.onStop()
	}

	this.gw.wg.Done()
	close(this.closed)
}

func (this *webServer) Closed() <-chan struct{} {
	return this.closed
}

func (this *webServer) notFoundHandler(w http.ResponseWriter, r *http.Request) {
	log.Error("%s: not found %s", this.name, r.RequestURI)

	writeNotFound(w)
}
<commit_msg>log the request method in 404 handler<commit_after>package gateway

import (
	"fmt"
	"net"
	"net/http"
	"runtime"
	"strings"
	"sync"
	"sync/atomic"
	"time"

	"github.com/funkygao/httprouter"
	log "github.com/funkygao/log4go"
)

type webServer struct {
	gw *Gateway

	name       string
	maxClients int

	router *httprouter.Router

	httpListener net.Listener
	httpServer   *http.Server

	httpsListener net.Listener
	httpsServer   *http.Server

	waitExitFunc    waitExitFunc
	connStateFunc   connStateFunc
	onConnNewFunc   onConnNewFunc
	onConnCloseFunc onConnCloseFunc

	onStop        func()
	mu            sync.Mutex
	waiterStarted bool

	// FIXME if http/https listener both enabled, must able to tell them apart
	activeConnN int32

	// TODO channel performance is frustrating, no better than mutex/map use ring buffer
	stateIdleCh, stateRemoveCh, stateActiveCh chan net.Conn

	closed chan struct{}
}

func newWebServer(name string, httpAddr, httpsAddr string, maxClients int,
	gw *Gateway) *webServer {
	const initialConnBuckets = 200
	this := &webServer{
		name:          name,
		gw:            gw,
		maxClients:    maxClients,
		stateActiveCh: make(chan net.Conn, initialConnBuckets),
		stateIdleCh:   make(chan net.Conn, initialConnBuckets),
		stateRemoveCh: make(chan net.Conn, initialConnBuckets),
		router:        httprouter.New(),
		closed:        make(chan struct{}),
	}

	if Options.EnableHttpPanicRecover {
		this.router.PanicHandler = func(w http.ResponseWriter, r *http.Request, err interface{}) {
			log.Error("PANIC %s %s(%s) %s %s: %+v", this.name, r.RemoteAddr, getHttpRemoteIp(r), r.Method, r.RequestURI, err)

			writeServerError(w, http.StatusText(http.StatusInternalServerError))
		}
	}

	if httpAddr != "" {
		this.httpServer = &http.Server{
			Addr:           httpAddr,
			Handler:        this.router,
			ReadTimeout:    Options.HttpReadTimeout,
			WriteTimeout:   Options.HttpWriteTimeout,
			MaxHeaderBytes: Options.HttpHeaderMaxBytes,
		}
	}

	if httpsAddr != "" {
		this.httpsServer = &http.Server{
			Addr:           httpsAddr,
			Handler:        this.router,
			ReadTimeout:    Options.HttpReadTimeout,
			WriteTimeout:   Options.HttpWriteTimeout,
			MaxHeaderBytes: Options.HttpHeaderMaxBytes,
		}
	}

	return this
}

func (this *webServer) Router() *httprouter.Router {
	return this.router
}

func (this *webServer) Start() {
	if this.waitExitFunc == nil {
		this.waitExitFunc = this.defaultWaitExit
	}
	if this.connStateFunc == nil {
		this.connStateFunc = this.defaultConnStateMachine

		go this.manageIdleConns()
	}

	if this.httpsServer != nil {
		this.httpsServer.ConnState = this.connStateFunc
		this.startServer(true)
	}

	if this.httpServer != nil {
		this.httpServer.ConnState = this.connStateFunc
		this.startServer(false)
	}

}

func (this *webServer) startServer(https bool) {
	var err error
	waitListenerUp := make(chan struct{})
	go func() {
		if Options.CpuAffinity {
			runtime.LockOSThread()
		}

		var (
			retryDelay         time.Duration
			theListener        net.Listener
			waitListenerUpOnce sync.Once
		)
		for {
			if https {
				this.httpsListener, err = net.Listen("tcp", this.httpsServer.Addr)
				if err != nil {
					if strings.HasSuffix(err.Error(), "address already in use") {
						// non-retriable error encountered
						panic(fmt.Errorf("%s listener: %v", this.name, err))
					}

					if retryDelay == 0 {
						retryDelay = 50 * time.Millisecond
					} else {
						retryDelay = 2 * retryDelay
					}
					if maxDelay := time.Second; retryDelay > maxDelay {
						retryDelay = maxDelay
					}
					log.Error("%s listener %v, retry in %v", this.name, err, retryDelay)
					time.Sleep(retryDelay)
					continue
				}

				theListener, err = setupHttpsListener(this.httpsListener, this.gw.certFile, this.gw.keyFile)
				if err != nil {
					panic(err)
				}
			} else {
				theListener, err = net.Listen("tcp", this.httpServer.Addr)
			}

			if err != nil {
				if retryDelay == 0 {
					retryDelay = 50 * time.Millisecond
				} else {
					retryDelay = 2 * retryDelay
				}
				if maxDelay := time.Second; retryDelay > maxDelay {
					retryDelay = maxDelay
				}
				log.Error("%s listener %v, retry in %v", this.name, err, retryDelay)
				time.Sleep(retryDelay)
				continue
			}

			theListener = LimitListener(this.name, this.gw, theListener, this.maxClients)
			waitListenerUpOnce.Do(func() {
				close(waitListenerUp)
			})

			// on non-temporary err, net/http will close the listener
			if https {
				this.mu.Lock()
				this.httpsListener = theListener
				this.mu.Unlock()

				err = this.httpsServer.Serve(theListener)
			} else {
				this.mu.Lock()
				this.httpListener = theListener
				this.mu.Unlock()

				err = this.httpServer.Serve(theListener)
			}

			select {
			case <-this.gw.shutdownCh:
				return

			default:
				log.Error("%s server: %v", this.name, err)
			}
		}
	}()

	// wait for the listener up
	<-waitListenerUp

	this.mu.Lock()
	if !this.waiterStarted {
		this.waiterStarted = true

		this.gw.wg.Add(1)
		go this.waitExitFunc(this.gw.shutdownCh)
	}
	this.mu.Unlock()

	if https {
		log.Info("%s https server ready on %s", this.name, this.httpsServer.Addr)
	} else {
		log.Info("%s http server ready on %s", this.name, this.httpServer.Addr)
	}
}

func (this *webServer) defaultConnStateMachine(c net.Conn, cs http.ConnState) {
	switch cs {
	case http.StateNew:
		atomic.AddInt32(&this.activeConnN, 1)

		if this.onConnNewFunc != nil {
			this.onConnNewFunc(c)
		}

	case http.StateIdle:
		this.stateIdleCh <- c

	case http.StateActive:
		this.stateActiveCh <- c

	case http.StateClosed, http.StateHijacked:
		atomic.AddInt32(&this.activeConnN, -1)
		this.stateRemoveCh <- c

		if this.onConnCloseFunc != nil {
			this.onConnCloseFunc(c)
		}
	}
}

func (this *webServer) manageIdleConns() {
	var (
		idleConns     = make(map[net.Conn]struct{}, 200)
		c             net.Conn
		waitNextRound = make(chan struct{}, 10)
	)
	defer close(waitNextRound)

	log.Debug("%s is managing idle connections", this.name)

	for {
		select {
		case <-this.gw.shutdownCh:
			if len(idleConns) == 0 {
				// happy ending
				log.Debug("%s closed all idle conns", this.name)
				return
			}

			t := time.Now().Add(time.Millisecond * 100)
			for conn := range idleConns {
				if conn == nil {
					continue
				}

				log.Debug("%s closing %s", this.name, conn.RemoteAddr())
				conn.SetDeadline(t)
			}

			// wait for next loop
			waitNextRound <- struct{}{}

		case <-waitNextRound:
			if len(idleConns) == 0 {
				log.Debug("%s closed all idle conns", this.name)
				return
			}

			t := time.Now().Add(time.Millisecond * 100)
			for conn := range idleConns {
				if conn == nil {
					continue
				}

				log.Debug("%s closing %s", this.name, conn.RemoteAddr())
				conn.SetDeadline(t)
			}

			// wait for next loop
			waitNextRound <- struct{}{}

		case c = <-this.stateActiveCh:
			delete(idleConns, c)

		case c = <-this.stateIdleCh:
			idleConns[c] = struct{}{}

		case c = <-this.stateRemoveCh:
			delete(idleConns, c)
		}
	}
}

func (this *webServer) defaultWaitExit(exit <-chan struct{}) {
	log.Debug("%s enter default wait exit", this.name)

	<-exit

	var err error
	if this.httpServer != nil {
		// HTTP response will have "Connection: close"
		this.httpServer.SetKeepAlivesEnabled(false)

		// avoid new connections
		this.mu.Lock()
		if this.httpListener != nil {
			err = this.httpListener.Close()
		}
		this.mu.Unlock()

		if err != nil {
			log.Error("%s on %s: %+v", this.name, this.httpServer.Addr, err)
		}

		log.Trace("%s on %s listener closed", this.name, this.httpServer.Addr)
	}

	if this.httpsServer != nil {
		// HTTP response will have "Connection: close"
		this.httpsServer.SetKeepAlivesEnabled(false)

		// avoid new connections
		this.mu.Lock()
		if this.httpsListener != nil {
			err = this.httpsListener.Close()
		}
		this.mu.Unlock()

		if err != nil {
			log.Error("%s on %s: %+v", this.name, this.httpsServer.Addr, err)
		}

		log.Trace("%s on %s listener closed", this.name, this.httpsServer.Addr)
	}

	// wait for all established http/https conns close
	waitStart := time.Now()
	var prompt sync.Once
	for {
		activeConnN := atomic.LoadInt32(&this.activeConnN)
		if activeConnN == 0 {
			// good luck, all connections finished
			break
		}

		prompt.Do(func() {
			log.Trace("%s waiting for %d clients shutdown...", this.name, activeConnN)
		})

		// timeout mechanism
		if time.Since(waitStart) > Options.SubTimeout+time.Second {
			log.Warn("%s still left %d conns after %s, forced to shutdown",
				this.name, activeConnN, Options.SubTimeout)
			break
		}

		time.Sleep(time.Millisecond * 50)
	}
	log.Trace("%s all connections finished", this.name)

	if this.httpsServer != nil {
		this.httpsServer.ConnState = nil
	}
	if this.httpServer != nil {
		this.httpServer.ConnState = nil
	}

	// TODO close will lead to race condition
	//close(this.stateActiveCh)
	//close(this.stateIdleCh)
	//close(this.stateRemoveCh)

	if this.onStop != nil {
		this.onStop()
	}

	this.gw.wg.Done()
	close(this.closed)
}

func (this *webServer) Closed() <-chan struct{} {
	return this.closed
}

func (this *webServer) notFoundHandler(w http.ResponseWriter, r *http.Request) {
	log.Error("%s: not found %s %s", this.name, r.Method, r.RequestURI)

	writeNotFound(w)
}
<|endoftext|>


<commit_before>// Copyright The 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 trace_test

import (
	"context"
	"testing"

	"go.opentelemetry.io/otel/attribute"
	sdktrace "go.opentelemetry.io/otel/sdk/trace"
	"go.opentelemetry.io/otel/trace"
)

func BenchmarkStartEndSpan(b *testing.B) {
	traceBenchmark(b, "Benchmark StartEndSpan", func(b *testing.B, t trace.Tracer) {
		ctx := context.Background()
		b.ResetTimer()
		for i := 0; i < b.N; i++ {
			_, span := t.Start(ctx, "/foo")
			span.End()
		}
	})
}

func BenchmarkSpanWithAttributes_4(b *testing.B) {
	traceBenchmark(b, "Benchmark Start With 4 Attributes", func(b *testing.B, t trace.Tracer) {
		ctx := context.Background()
		b.ResetTimer()

		for i := 0; i < b.N; i++ {
			_, span := t.Start(ctx, "/foo")
			span.SetAttributes(
				attribute.Bool("key1", false),
				attribute.String("key2", "hello"),
				attribute.Float64("key4", 123.456),
			)
			span.End()
		}
	})
}

func BenchmarkSpanWithAttributes_8(b *testing.B) {
	traceBenchmark(b, "Benchmark Start With 8 Attributes", func(b *testing.B, t trace.Tracer) {
		ctx := context.Background()
		b.ResetTimer()

		for i := 0; i < b.N; i++ {
			_, span := t.Start(ctx, "/foo")
			span.SetAttributes(
				attribute.Bool("key1", false),
				attribute.String("key2", "hello"),
				attribute.Float64("key4", 123.456),
				attribute.Bool("key21", false),
				attribute.String("key22", "hello"),
				attribute.Float64("key24", 123.456),
			)
			span.End()
		}
	})
}

func BenchmarkSpanWithAttributes_all(b *testing.B) {
	traceBenchmark(b, "Benchmark Start With all Attribute types", func(b *testing.B, t trace.Tracer) {
		ctx := context.Background()
		b.ResetTimer()

		for i := 0; i < b.N; i++ {
			_, span := t.Start(ctx, "/foo")
			span.SetAttributes(
				attribute.Bool("key1", false),
				attribute.String("key2", "hello"),
				attribute.Int64("key3", 123),
				attribute.Float64("key7", 123.456),
				attribute.Int("key9", 123),
			)
			span.End()
		}
	})
}

func BenchmarkSpanWithAttributes_all_2x(b *testing.B) {
	traceBenchmark(b, "Benchmark Start With all Attributes types twice", func(b *testing.B, t trace.Tracer) {
		ctx := context.Background()
		b.ResetTimer()

		for i := 0; i < b.N; i++ {
			_, span := t.Start(ctx, "/foo")
			span.SetAttributes(
				attribute.Bool("key1", false),
				attribute.String("key2", "hello"),
				attribute.Int64("key3", 123),
				attribute.Float64("key7", 123.456),
				attribute.Int("key10", 123),
				attribute.Bool("key21", false),
				attribute.String("key22", "hello"),
				attribute.Int64("key23", 123),
				attribute.Float64("key27", 123.456),
				attribute.Int("key210", 123),
			)
			span.End()
		}
	})
}

func BenchmarkTraceID_DotString(b *testing.B) {
	t, _ := trace.TraceIDFromHex("0000000000000001000000000000002a")
	sc := trace.NewSpanContext(trace.SpanContextConfig{TraceID: t})

	want := "0000000000000001000000000000002a"
	for i := 0; i < b.N; i++ {
		if got := sc.TraceID().String(); got != want {
			b.Fatalf("got = %q want = %q", got, want)
		}
	}
}

func BenchmarkSpanID_DotString(b *testing.B) {
	sc := trace.NewSpanContext(trace.SpanContextConfig{SpanID: trace.SpanID{1}})
	want := "0100000000000000"
	for i := 0; i < b.N; i++ {
		if got := sc.SpanID().String(); got != want {
			b.Fatalf("got = %q want = %q", got, want)
		}
	}
}

func traceBenchmark(b *testing.B, name string, fn func(*testing.B, trace.Tracer)) {
	b.Run("AlwaysSample", func(b *testing.B) {
		b.ReportAllocs()
		fn(b, tracer(b, name, sdktrace.AlwaysSample()))
	})
	b.Run("NeverSample", func(b *testing.B) {
		b.ReportAllocs()
		fn(b, tracer(b, name, sdktrace.NeverSample()))
	})
}

func tracer(b *testing.B, name string, sampler sdktrace.Sampler) trace.Tracer {
	tp := sdktrace.NewTracerProvider(sdktrace.WithSampler(sampler))
	return tp.Tracer(name)
}
<commit_msg>Added Benchmarks around events (#2405)<commit_after>// Copyright The 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 trace_test

import (
	"context"
	"testing"
	"time"

	"go.opentelemetry.io/otel/attribute"
	sdktrace "go.opentelemetry.io/otel/sdk/trace"
	"go.opentelemetry.io/otel/trace"
)

func BenchmarkStartEndSpan(b *testing.B) {
	traceBenchmark(b, "Benchmark StartEndSpan", func(b *testing.B, t trace.Tracer) {
		ctx := context.Background()
		b.ResetTimer()
		for i := 0; i < b.N; i++ {
			_, span := t.Start(ctx, "/foo")
			span.End()
		}
	})
}

func BenchmarkSpanWithAttributes_4(b *testing.B) {
	traceBenchmark(b, "Benchmark Start With 4 Attributes", func(b *testing.B, t trace.Tracer) {
		ctx := context.Background()
		b.ResetTimer()

		for i := 0; i < b.N; i++ {
			_, span := t.Start(ctx, "/foo")
			span.SetAttributes(
				attribute.Bool("key1", false),
				attribute.String("key2", "hello"),
				attribute.Int64("key3", 123),
				attribute.Float64("key4", 123.456),
			)
			span.End()
		}
	})
}

func BenchmarkSpanWithAttributes_8(b *testing.B) {
	traceBenchmark(b, "Benchmark Start With 8 Attributes", func(b *testing.B, t trace.Tracer) {
		ctx := context.Background()
		b.ResetTimer()

		for i := 0; i < b.N; i++ {
			_, span := t.Start(ctx, "/foo")
			span.SetAttributes(
				attribute.Bool("key1", false),
				attribute.String("key2", "hello"),
				attribute.Int64("key3", 123),
				attribute.Float64("key4", 123.456),
				attribute.Bool("key21", false),
				attribute.String("key22", "hello"),
				attribute.Int64("key23", 123),
				attribute.Float64("key24", 123.456),
			)
			span.End()
		}
	})
}

func BenchmarkSpanWithAttributes_all(b *testing.B) {
	traceBenchmark(b, "Benchmark Start With all Attribute types", func(b *testing.B, t trace.Tracer) {
		ctx := context.Background()
		b.ResetTimer()

		for i := 0; i < b.N; i++ {
			_, span := t.Start(ctx, "/foo")
			span.SetAttributes(
				attribute.Bool("key1", false),
				attribute.String("key2", "hello"),
				attribute.Int64("key3", 123),
				attribute.Float64("key7", 123.456),
				attribute.Int("key9", 123),
			)
			span.End()
		}
	})
}

func BenchmarkSpanWithAttributes_all_2x(b *testing.B) {
	traceBenchmark(b, "Benchmark Start With all Attributes types twice", func(b *testing.B, t trace.Tracer) {
		ctx := context.Background()
		b.ResetTimer()

		for i := 0; i < b.N; i++ {
			_, span := t.Start(ctx, "/foo")
			span.SetAttributes(
				attribute.Bool("key1", false),
				attribute.String("key2", "hello"),
				attribute.Int64("key3", 123),
				attribute.Float64("key7", 123.456),
				attribute.Int("key10", 123),
				attribute.Bool("key21", false),
				attribute.String("key22", "hello"),
				attribute.Int64("key23", 123),
				attribute.Float64("key27", 123.456),
				attribute.Int("key210", 123),
			)
			span.End()
		}
	})
}

func BenchmarkSpanWithEvents_4(b *testing.B) {
	traceBenchmark(b, "Benchmark Start With 4 Events", func(b *testing.B, t trace.Tracer) {
		ctx := context.Background()
		b.ResetTimer()

		for i := 0; i < b.N; i++ {
			_, span := t.Start(ctx, "/foo")
			span.AddEvent("event1")
			span.AddEvent("event2")
			span.AddEvent("event3")
			span.AddEvent("event4")
			span.End()
		}
	})
}

func BenchmarkSpanWithEvents_8(b *testing.B) {
	traceBenchmark(b, "Benchmark Start With 4 Events", func(b *testing.B, t trace.Tracer) {
		ctx := context.Background()
		b.ResetTimer()

		for i := 0; i < b.N; i++ {
			_, span := t.Start(ctx, "/foo")
			span.AddEvent("event1")
			span.AddEvent("event2")
			span.AddEvent("event3")
			span.AddEvent("event4")
			span.AddEvent("event5")
			span.AddEvent("event6")
			span.AddEvent("event7")
			span.AddEvent("event8")
			span.End()
		}
	})
}

func BenchmarkSpanWithEvents_WithStackTrace(b *testing.B) {
	traceBenchmark(b, "Benchmark Start With 4 Attributes", func(b *testing.B, t trace.Tracer) {
		ctx := context.Background()
		b.ResetTimer()

		for i := 0; i < b.N; i++ {
			_, span := t.Start(ctx, "/foo")
			span.AddEvent("event1", trace.WithStackTrace(true))
			span.End()
		}
	})
}
func BenchmarkSpanWithEvents_WithTimestamp(b *testing.B) {
	traceBenchmark(b, "Benchmark Start With 4 Attributes", func(b *testing.B, t trace.Tracer) {
		ctx := context.Background()
		b.ResetTimer()

		for i := 0; i < b.N; i++ {
			_, span := t.Start(ctx, "/foo")
			span.AddEvent("event1", trace.WithTimestamp(time.Unix(0, 0)))
			span.End()
		}
	})
}

func BenchmarkTraceID_DotString(b *testing.B) {
	t, _ := trace.TraceIDFromHex("0000000000000001000000000000002a")
	sc := trace.NewSpanContext(trace.SpanContextConfig{TraceID: t})

	want := "0000000000000001000000000000002a"
	for i := 0; i < b.N; i++ {
		if got := sc.TraceID().String(); got != want {
			b.Fatalf("got = %q want = %q", got, want)
		}
	}
}

func BenchmarkSpanID_DotString(b *testing.B) {
	sc := trace.NewSpanContext(trace.SpanContextConfig{SpanID: trace.SpanID{1}})
	want := "0100000000000000"
	for i := 0; i < b.N; i++ {
		if got := sc.SpanID().String(); got != want {
			b.Fatalf("got = %q want = %q", got, want)
		}
	}
}

func traceBenchmark(b *testing.B, name string, fn func(*testing.B, trace.Tracer)) {
	b.Run("AlwaysSample", func(b *testing.B) {
		b.ReportAllocs()
		fn(b, tracer(b, name, sdktrace.AlwaysSample()))
	})
	b.Run("NeverSample", func(b *testing.B) {
		b.ReportAllocs()
		fn(b, tracer(b, name, sdktrace.NeverSample()))
	})
}

func tracer(b *testing.B, name string, sampler sdktrace.Sampler) trace.Tracer {
	tp := sdktrace.NewTracerProvider(sdktrace.WithSampler(sampler))
	return tp.Tracer(name)
}
<|endoftext|>


<commit_before>package api

import (
	"encoding/json"
	"time"
)

const (
	UNIT_BYTES = "bytes"
	UNIT_BYTES_PER_SECOND = "bytes/second"
	UNIT_BITS = "bits"
	UNIT_BITS_PER_SECOND = "bits/second"
	UNIT_SECONDS = "seconds"
	UNIT_COUNT = "count"
	UNIT_COUNT_PER_SECOND = "count/second"
	UNIT_PERCENT = "percent"
)

type Metric struct {
	Name      string
	Value     float64
	Unit      string
	Timestamp time.Time
	Metadata  map[string]string
}

// Add metadata values to the metric if they do not already exist.
func (m *Metric) Underlay(metadata map[string]string) {
	if m.Metadata == nil {
		m.Metadata = make(map[string]string, len(metadata))
	}
	for key, value := range metadata {
		if _, ok := m.Metadata[key]; !ok {
			m.Metadata[key] = value
		}
	}
}

// A collection of metric values.
type Metrics struct {
	metrics []*Metric
}

func (m *Metrics) MarshalJSON() ([]byte, error) {
	val, err := json.Marshal(m.metrics)
	return val, err
}

func (m *Metrics) UnmarshalJSON(raw []byte) error {
	err := json.Unmarshal(raw, &m.metrics)
	return err
}

// Add a single metric to the collection.
func (m *Metrics) Add(name string, value float64, unit string, timestamp time.Time, metadata map[string]string) {
	m.metrics = append(m.metrics, &Metric{name, value, unit, timestamp, metadata})
}

// Append other metrics to this metrics struct.
func (m *Metrics) Append(metrics Metrics) {
	m.metrics = append(m.metrics, metrics.metrics...)
}

// Retrieve the metrics as an array.
func (m *Metrics) Items() []*Metric {
	return m.metrics
}

<commit_msg>Add more units.<commit_after>package api

import (
	"encoding/json"
	"time"
)

const (
	UNIT_BYTES = "bytes"
	UNIT_KILOBYTES = "kilobytes"
	UNIT_MEGABYTES = "megabytes"
	UNIT_GIGABYTES = "gigabytes"
	UNIT_TERABYTES = "terabytes"
	UNIT_BYTES_PER_SECOND = "bytes/second"
	UNIT_KILOBYTES_PER_SECOND = "kilobytes/second"
	UNIT_MEGABYTES_PER_SECOND = "megabytes/second"
	UNIT_GIGABYTES_PER_SECOND = "gigabytes/second"
	UNIT_TERABYTES_PER_SECOND = "terabytes/second"
	UNIT_BITS = "bits"
	UNIT_KILOBITS = "kilobits"
	UNIT_MEGABITS = "megabits"
	UNIT_GIGABITS = "gigabits"
	UNIT_TERABITS = "terabits"
	UNIT_BITS_PER_SECOND = "bits/second"
	UNIT_KILOBITS_PER_SECOND = "kilobits/second"
	UNIT_MEGABITS_PER_SECOND = "megabits/second"
	UNIT_GIGABITS_PER_SECOND = "gigabits/second"
	UNIT_TERABITS_PER_SECOND = "terabits/second"
	UNIT_SECONDS = "seconds"
	UNIT_PERCENT = "percent"
	UNIT_COUNT = "count"
	UNIT_COUNT_PER_SECOND = "count/second"
)

type Metric struct {
	Name      string
	Value     float64
	Unit      string
	Timestamp time.Time
	Metadata  map[string]string
}

// Add metadata values to the metric if they do not already exist.
func (m *Metric) Underlay(metadata map[string]string) {
	if m.Metadata == nil {
		m.Metadata = make(map[string]string, len(metadata))
	}
	for key, value := range metadata {
		if _, ok := m.Metadata[key]; !ok {
			m.Metadata[key] = value
		}
	}
}

// A collection of metric values.
type Metrics struct {
	metrics []*Metric
}

func (m *Metrics) MarshalJSON() ([]byte, error) {
	val, err := json.Marshal(m.metrics)
	return val, err
}

func (m *Metrics) UnmarshalJSON(raw []byte) error {
	err := json.Unmarshal(raw, &m.metrics)
	return err
}

// Add a single metric to the collection.
func (m *Metrics) Add(name string, value float64, unit string, timestamp time.Time, metadata map[string]string) {
	m.metrics = append(m.metrics, &Metric{name, value, unit, timestamp, metadata})
}

// Append other metrics to this metrics struct.
func (m *Metrics) Append(metrics Metrics) {
	m.metrics = append(m.metrics, metrics.metrics...)
}

// Retrieve the metrics as an array.
func (m *Metrics) Items() []*Metric {
	return m.metrics
}

<|endoftext|>


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

package producer

import (
	"bufio"
	"fmt"
	"github.com/trivago/gollum/core"
	"github.com/trivago/gollum/core/log"
	"github.com/trivago/gollum/shared"
	"io"
	"io/ioutil"
	"os"
	"path/filepath"
	"time"
)

type spoolFile struct {
	file        *os.File
	batch       core.MessageBatch
	assembly    core.WriterAssembly
	fileCreated time.Time
	streamName  string
	prod        *Spooling
	source      core.MessageSource
}

const maxSpoolFileNumber = 99999999 // maximum file number defined by %08d -> 8 digits
const spoolFileFormatString = "%s/%08d.spl"

func newSpoolFile(prod *Spooling, streamName string, source core.MessageSource) *spoolFile {
	spool := &spoolFile{
		file:        nil,
		batch:       core.NewMessageBatch(prod.batchMaxCount),
		assembly:    core.NewWriterAssembly(nil, prod.Drop, prod.GetFormatter()),
		fileCreated: time.Now(),
		streamName:  streamName,
		prod:        prod,
		source:      source,
	}

	shared.Metric.New(spoolingMetricName + streamName)
	shared.Metric.New(spooledMetricName + streamName)
	go spool.read()
	return spool
}

func (spool *spoolFile) flush() {
	spool.batch.Flush(spool.assembly.Write)
}

func (spool *spoolFile) close() {
	spool.batch.Flush(spool.assembly.Flush)
	spool.batch.WaitForFlush(spool.prod.GetShutdownTimeout())
	spool.file.Close()
}

func (spool *spoolFile) getFileNumbering(path string) (min int, max int) {
	min, max = maxSpoolFileNumber+1, 0
	files, _ := ioutil.ReadDir(path)
	for _, file := range files {
		base := filepath.Base(file.Name())
		number, _ := shared.Btoi([]byte(base)) // Because we need leading zero support
		min = shared.MinI(min, int(number))
		max = shared.MaxI(max, int(number))
	}
	return min, max
}

func (spool *spoolFile) openOrRotate() bool {
	fileSize := int64(0)
	if spool.file != nil {
		fileInfo, _ := spool.file.Stat()
		fileSize = fileInfo.Size()
	}

	if spool.file == nil || fileSize >= spool.prod.maxFileSize || time.Since(spool.fileCreated) > spool.prod.maxFileAge {
		path := spool.prod.path + "/" + spool.streamName
		_, maxSuffix := spool.getFileNumbering(path)

		// Reopen spooling file
		var err error
		oldFile := spool.file
		spoolFileName := fmt.Sprintf(spoolFileFormatString, path, maxSuffix+1)
		spool.file, err = os.OpenFile(spoolFileName, os.O_WRONLY|os.O_CREATE, 0600)

		// Close current file
		if oldFile != nil {
			spool.batch.WaitForFlush(time.Duration(0))
			oldFile.Close()
		}

		spool.assembly.SetWriter(spool.file)
		spool.fileCreated = time.Now()
		if err != nil {
			Log.Error.Print("Spooling: ", err)
			return false // ### return, could not open file ###
		}
		Log.Debug.Print("Spooler opened ", spoolFileName, " for writing")
	}

	return true
}

func (spool *spoolFile) read() {
	spool.prod.AddWorker()
	defer spool.prod.WorkerDone()

	path := spool.prod.path + "/" + spool.streamName

	for spool.prod.IsActive() {
		minSuffix, _ := spool.getFileNumbering(path)

		spoolFileName := fmt.Sprintf(spoolFileFormatString, path, minSuffix)
		if minSuffix == 0 || minSuffix > maxSpoolFileNumber || (spool.file != nil && spool.file.Name() == spoolFileName) {
			if minSuffix > maxSpoolFileNumber {
				Log.Debug.Print("Spool read sleeps (no file)")
			} else {
				Log.Debug.Printf("Spool read waits for %s", spoolFileName)
			}
			time.Sleep(spool.prod.maxFileAge / 2)
			continue // ### continue, try again ###
		}

		file, err := os.OpenFile(spoolFileName, os.O_RDONLY, 0600)
		if err != nil {
			Log.Error.Print("Spool read open error ", err)
			continue // ### continue, try again ###
		}

		Log.Debug.Print("Spooler opened ", spoolFileName, " for reading")
		reader := bufio.NewReader(file)
		for {
			// Only spool back if target is not busy
			if spool.source != nil && spool.source.IsBlocked() {
				time.Sleep(time.Second)
				continue // ### contine, busy source ###
			}

			// Read one line (might require multiple reads)
			buffer, isPartial, err := reader.ReadLine()
			for isPartial && err == nil {
				var appendix []byte
				appendix, isPartial, err = reader.ReadLine()
				buffer = append(buffer, appendix...)
			}
			// Any error cancels the loop
			if err != nil {
				if err != io.EOF {
					Log.Error.Print("Spool read error: ", err)
				}
				break // ### break, read error or EOF ###
			}
			// Deserialize from string
			msg, err := core.DeserializeMessage(string(buffer))
			if err != nil {
				Log.Error.Print("Spool file read: ", err)
			} else {
				spool.prod.routeToOrigin(msg)
			}
		}

		// Close and remove file
		Log.Debug.Print("Spooler removes ", spoolFileName)
		file.Close()
		os.Remove(spoolFileName)
	}
}
<commit_msg>[fix] Inner spooling loop exit condition<commit_after>// Copyright 2015 trivago GmbH
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package producer

import (
	"bufio"
	"fmt"
	"github.com/trivago/gollum/core"
	"github.com/trivago/gollum/core/log"
	"github.com/trivago/gollum/shared"
	"io"
	"io/ioutil"
	"os"
	"path/filepath"
	"time"
)

type spoolFile struct {
	file        *os.File
	batch       core.MessageBatch
	assembly    core.WriterAssembly
	fileCreated time.Time
	streamName  string
	prod        *Spooling
	source      core.MessageSource
}

const maxSpoolFileNumber = 99999999 // maximum file number defined by %08d -> 8 digits
const spoolFileFormatString = "%s/%08d.spl"

func newSpoolFile(prod *Spooling, streamName string, source core.MessageSource) *spoolFile {
	spool := &spoolFile{
		file:        nil,
		batch:       core.NewMessageBatch(prod.batchMaxCount),
		assembly:    core.NewWriterAssembly(nil, prod.Drop, prod.GetFormatter()),
		fileCreated: time.Now(),
		streamName:  streamName,
		prod:        prod,
		source:      source,
	}

	shared.Metric.New(spoolingMetricName + streamName)
	shared.Metric.New(spooledMetricName + streamName)
	go spool.read()
	return spool
}

func (spool *spoolFile) flush() {
	spool.batch.Flush(spool.assembly.Write)
}

func (spool *spoolFile) close() {
	spool.batch.Flush(spool.assembly.Flush)
	spool.batch.WaitForFlush(spool.prod.GetShutdownTimeout())
	spool.file.Close()
}

func (spool *spoolFile) getFileNumbering(path string) (min int, max int) {
	min, max = maxSpoolFileNumber+1, 0
	files, _ := ioutil.ReadDir(path)
	for _, file := range files {
		base := filepath.Base(file.Name())
		number, _ := shared.Btoi([]byte(base)) // Because we need leading zero support
		min = shared.MinI(min, int(number))
		max = shared.MaxI(max, int(number))
	}
	return min, max
}

func (spool *spoolFile) openOrRotate() bool {
	fileSize := int64(0)
	if spool.file != nil {
		fileInfo, _ := spool.file.Stat()
		fileSize = fileInfo.Size()
	}

	if spool.file == nil || fileSize >= spool.prod.maxFileSize || time.Since(spool.fileCreated) > spool.prod.maxFileAge {
		path := spool.prod.path + "/" + spool.streamName
		_, maxSuffix := spool.getFileNumbering(path)

		// Reopen spooling file
		var err error
		oldFile := spool.file
		spoolFileName := fmt.Sprintf(spoolFileFormatString, path, maxSuffix+1)
		spool.file, err = os.OpenFile(spoolFileName, os.O_WRONLY|os.O_CREATE, 0600)

		// Close current file
		if oldFile != nil {
			spool.batch.WaitForFlush(time.Duration(0))
			oldFile.Close()
		}

		spool.assembly.SetWriter(spool.file)
		spool.fileCreated = time.Now()
		if err != nil {
			Log.Error.Print("Spooling: ", err)
			return false // ### return, could not open file ###
		}
		Log.Debug.Print("Spooler opened ", spoolFileName, " for writing")
	}

	return true
}

func (spool *spoolFile) read() {
	spool.prod.AddWorker()
	defer spool.prod.WorkerDone()

	path := spool.prod.path + "/" + spool.streamName

	for spool.prod.IsActive() {
		minSuffix, _ := spool.getFileNumbering(path)

		spoolFileName := fmt.Sprintf(spoolFileFormatString, path, minSuffix)
		if minSuffix == 0 || minSuffix > maxSpoolFileNumber || (spool.file != nil && spool.file.Name() == spoolFileName) {
			if minSuffix > maxSpoolFileNumber {
				Log.Debug.Print("Spool read sleeps (no file)")
			} else {
				Log.Debug.Printf("Spool read waits for %s", spoolFileName)
			}
			time.Sleep(spool.prod.maxFileAge / 2)
			continue // ### continue, try again ###
		}

		file, err := os.OpenFile(spoolFileName, os.O_RDONLY, 0600)
		if err != nil {
			Log.Error.Print("Spool read open error ", err)
			continue // ### continue, try again ###
		}

		Log.Debug.Print("Spooler opened ", spoolFileName, " for reading")
		reader := bufio.NewReader(file)

		for spool.prod.IsActive() {
			// Only spool back if target is not busy
			if spool.source != nil && spool.source.IsBlocked() {
				Log.Debug.Print("Spool read sleeps on inactive source")
				time.Sleep(time.Second)
				continue // ### contine, busy source ###
			}

			// Read one line (might require multiple reads)
			buffer, isPartial, err := reader.ReadLine()
			for isPartial && err == nil {
				var appendix []byte
				appendix, isPartial, err = reader.ReadLine()
				buffer = append(buffer, appendix...)
			}
			// Any error cancels the loop
			if err != nil {
				if err != io.EOF {
					Log.Error.Print("Spool read error: ", err)
				}
				break // ### break, read error or EOF ###
			}
			// Deserialize from string
			msg, err := core.DeserializeMessage(string(buffer))
			if err != nil {
				Log.Error.Print("Spool file read: ", err)
			} else {
				spool.prod.routeToOrigin(msg)
			}
		}

		// Close and remove file
		Log.Debug.Print("Spooler removes ", spoolFileName)
		file.Close()
		os.Remove(spoolFileName)
	}
}
<|endoftext|>