<commit_before>/*
   Nging is a toolbox for webmasters
   Copyright (C) 2018-present  Wenhui Shen <swh@admpub.com>

   This program is free software: you can redistribute it and/or modify
   it under the terms of the GNU Affero General Public License as published
   by the Free Software Foundation, either version 3 of the License, or
   (at your option) any later version.

   This program is distributed in the hope that it will be useful,
   but WITHOUT ANY WARRANTY; without even the implied warranty of
   MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
   GNU Affero General Public License for more details.

   You should have received a copy of the GNU Affero General Public License
   along with this program.  If not, see <https://www.gnu.org/licenses/>.
*/

package middleware

import (
	"fmt"
	"html/template"
	"strconv"
	"time"

	"github.com/webx-top/com"
	"github.com/webx-top/echo"
	"github.com/webx-top/echo/middleware/tplfunc"
	"github.com/webx-top/echo/subdomains"

	"github.com/admpub/nging/application/dbschema"
	"github.com/admpub/nging/application/library/common"
	"github.com/admpub/nging/application/library/config"
	"github.com/admpub/nging/application/library/license"
	"github.com/admpub/nging/application/library/modal"
	"github.com/admpub/nging/application/model"
	"github.com/admpub/nging/application/registry/dashboard"
	"github.com/admpub/nging/application/registry/navigate"
	"github.com/admpub/nging/application/registry/upload"
	"github.com/admpub/nging/application/registry/upload/helper"
)

var DefaultAvatarURL = `/public/assets/backend/images/user_128.png`

func ErrorPageFunc(c echo.Context) error {
	c.SetFunc(`URLFor`, subdomains.Default.URL)
	c.SetFunc(`IsMessage`, common.IsMessage)
	c.SetFunc(`Stored`, c.Stored)
	c.SetFunc(`Languages`, func() []string {
		return config.DefaultConfig.Language.AllList
	})
	c.SetFunc(`IsError`, common.IsError)
	c.SetFunc(`IsOk`, common.IsOk)
	c.SetFunc(`Message`, common.Message)
	c.SetFunc(`Ok`, common.OkString)
	c.SetFunc(`Version`, func() *config.VersionInfo { return config.Version })
	c.SetFunc(`VersionNumber`, func() string { return config.Version.Number })
	c.SetFunc(`CommitID`, func() string { return config.Version.CommitID })
	c.SetFunc(`BuildTime`, func() string { return config.Version.BuildTime })
	c.SetFunc(`TrackerURL`, license.TrackerURL)
	c.SetFunc(`Fetch`, func(tmpl string, data interface{}) template.HTML {
		b, e := c.Fetch(tmpl, data)
		if e != nil {
			return template.HTML(e.Error())
		}
		return template.HTML(string(b))
	})
	c.SetFunc(`Prefix`, func() string {
		return c.Route().Prefix
	})
	c.SetFunc(`Path`, c.Path)
	c.SetFunc(`Queries`, c.Queries)
	configs := config.Setting()
	c.SetFunc(`Config`, func(args ...string) echo.H {
		if len(args) > 0 {
			return config.Setting(args...)
		}
		return configs
	})
	return nil
}

func FuncMap() echo.MiddlewareFunc {
	return func(h echo.Handler) echo.Handler {
		return echo.HandlerFunc(func(c echo.Context) error {
			now := com.NewTime(time.Now())
			c.SetFunc(`Now`, func() *com.Time {
				return now
			})
			c.SetFunc(`UnixTime`, now.Local().Unix)
			c.SetFunc(`HasString`, hasString)
			c.SetFunc(`Date`, date)
			c.SetFunc(`Token`, upload.Token)
			c.SetFunc(`BackendUploadURL`, upload.BackendUploadURL)
			c.SetFunc(`FrontendUploadURL`, upload.FrontendUploadURL)
			c.SetFunc(`Modal`, func(data interface{}) template.HTML {
				return modal.Render(c, data)
			})
			ErrorPageFunc(c)
			c.SetFunc(`IndexStrSlice`, indexStrSlice)

			if !config.DefaultConfig.ConnectedDB(false) {
				return h.Handle(c)
			}
			//用户相关函数
			user, _ := c.Session().Get(`user`).(*dbschema.User)
			roleM := model.NewUserRole(c)
			var roleList []*dbschema.UserRole
			if user != nil {
				c.Set(`user`, user)
				c.SetFunc(`Username`, func() string { return user.Username })
				roleList = roleM.ListByUser(user)
				c.Set(`roleList`, roleList)
			}
			c.SetFunc(`Avatar`, func(avatar string, defaults ...string) string {
				if len(avatar) > 0 {
					return tplfunc.AddSuffix(avatar, `_200_200`)
				}
				if len(defaults) > 0 && len(defaults[0]) > 0 {
					return defaults[0]
				}
				return DefaultAvatarURL
			})
			c.SetFunc(`FileTypeByName`, helper.FileTypeByName)
			c.SetFunc(`FileTypeIcon`, helper.FileTypeIcon)
			c.SetFunc(`Project`, func(ident string) *navigate.ProjectItem {
				return navigate.ProjectGet(ident)
			})
			c.SetFunc(`ProjectIdent`, func() string {
				return navigate.ProjectIdent(c.Path())
			})
			c.SetFunc(`Projects`, func() navigate.ProjectList {
				return navigate.ProjectListAll()
			})
			c.SetFunc(`TopButtons`, func() dashboard.TopButtons {
				buttons := dashboard.TopButtonAll()
				buttons.Ready(c)
				return buttons
			})
			c.SetFunc(`GlobalFooters`, func() dashboard.GlobalFooters {
				footers := dashboard.GlobalFooterAll()
				footers.Ready(c)
				return footers
			})
			c.SetFunc(`Navigate`, func(side string) navigate.List {
				switch side {
				case `top`:
					if user != nil && user.Id == 1 {
						if navigate.TopNavigate == nil {
							return navigate.EmptyList
						}
						return *navigate.TopNavigate
					}
					return roleM.FilterNavigate(roleList, navigate.TopNavigate)
				case `left`:
					fallthrough
				default:
					var leftNav *navigate.List
					ident := navigate.ProjectIdent(c.Path())
					if len(ident) > 0 {
						if proj := navigate.ProjectGet(ident); proj != nil {
							leftNav = proj.NavList
						}
					}
					if leftNav == nil {
						if proj := navigate.ProjectFirst(); proj != nil {
							leftNav = proj.NavList
						}
					}
					if user != nil && user.Id == 1 {
						if leftNav == nil {
							return navigate.EmptyList
						}
						return *leftNav
					}
					return roleM.FilterNavigate(roleList, leftNav)
				}
			})
			return h.Handle(c)
		})
	}
}

func indexStrSlice(slice []string, index int) string {
	if slice == nil {
		return ``
	}
	if index >= len(slice) {
		return ``
	}
	return slice[index]
}

func hasString(slice []string, str string) bool {
	if slice == nil {
		return false
	}
	for _, v := range slice {
		if v == str {
			return true
		}
	}
	return false
}

func date(timestamp interface{}) time.Time {
	if v, y := timestamp.(int64); y {
		return time.Unix(v, 0)
	}
	if v, y := timestamp.(uint); y {
		return time.Unix(int64(v), 0)
	}
	v, _ := strconv.ParseInt(fmt.Sprint(timestamp), 10, 64)
	return time.Unix(v, 0)
}
<commit_msg>update<commit_after>/*
   Nging is a toolbox for webmasters
   Copyright (C) 2018-present  Wenhui Shen <swh@admpub.com>

   This program is free software: you can redistribute it and/or modify
   it under the terms of the GNU Affero General Public License as published
   by the Free Software Foundation, either version 3 of the License, or
   (at your option) any later version.

   This program is distributed in the hope that it will be useful,
   but WITHOUT ANY WARRANTY; without even the implied warranty of
   MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
   GNU Affero General Public License for more details.

   You should have received a copy of the GNU Affero General Public License
   along with this program.  If not, see <https://www.gnu.org/licenses/>.
*/

package middleware

import (
	"fmt"
	"html/template"
	"strconv"
	"time"

	"github.com/webx-top/com"
	"github.com/webx-top/echo"
	"github.com/webx-top/echo/middleware/tplfunc"
	"github.com/webx-top/echo/subdomains"

	"github.com/admpub/nging/application/dbschema"
	"github.com/admpub/nging/application/library/common"
	"github.com/admpub/nging/application/library/config"
	"github.com/admpub/nging/application/library/license"
	"github.com/admpub/nging/application/library/modal"
	"github.com/admpub/nging/application/model"
	"github.com/admpub/nging/application/registry/dashboard"
	"github.com/admpub/nging/application/registry/navigate"
	"github.com/admpub/nging/application/registry/upload"
	"github.com/admpub/nging/application/registry/upload/helper"
)

var DefaultAvatarURL = `/public/assets/backend/images/user_128.png`

func ErrorPageFunc(c echo.Context) error {
	c.SetFunc(`URLFor`, subdomains.Default.URL)
	c.SetFunc(`IsMessage`, common.IsMessage)
	c.SetFunc(`Stored`, c.Stored)
	c.SetFunc(`Languages`, func() []string {
		return config.DefaultConfig.Language.AllList
	})
	c.SetFunc(`IsError`, common.IsError)
	c.SetFunc(`IsOk`, common.IsOk)
	c.SetFunc(`Message`, common.Message)
	c.SetFunc(`Ok`, common.OkString)
	c.SetFunc(`Version`, func() *config.VersionInfo { return config.Version })
	c.SetFunc(`VersionNumber`, func() string { return config.Version.Number })
	c.SetFunc(`CommitID`, func() string { return config.Version.CommitID })
	c.SetFunc(`BuildTime`, func() string { return config.Version.BuildTime })
	c.SetFunc(`TrackerURL`, license.TrackerURL)
	c.SetFunc(`Fetch`, func(tmpl string, data interface{}) template.HTML {
		b, e := c.Fetch(tmpl, data)
		if e != nil {
			return template.HTML(e.Error())
		}
		return template.HTML(string(b))
	})
	c.SetFunc(`Prefix`, func() string {
		return c.Route().Prefix
	})
	c.SetFunc(`Path`, c.Path)
	c.SetFunc(`Queries`, c.Queries)
	c.SetFunc(`Domain`, c.Domain)
	c.SetFunc(`Port`, c.Port)
	c.SetFunc(`Scheme`, c.Scheme)
	c.SetFunc(`Site`, c.Site)
	configs := config.Setting()
	c.SetFunc(`Config`, func(args ...string) echo.H {
		if len(args) > 0 {
			return config.Setting(args...)
		}
		return configs
	})
	return nil
}

func FuncMap() echo.MiddlewareFunc {
	return func(h echo.Handler) echo.Handler {
		return echo.HandlerFunc(func(c echo.Context) error {
			now := com.NewTime(time.Now())
			c.SetFunc(`Now`, func() *com.Time {
				return now
			})
			c.SetFunc(`UnixTime`, now.Local().Unix)
			c.SetFunc(`HasString`, hasString)
			c.SetFunc(`Date`, date)
			c.SetFunc(`Token`, upload.Token)
			c.SetFunc(`BackendUploadURL`, upload.BackendUploadURL)
			c.SetFunc(`FrontendUploadURL`, upload.FrontendUploadURL)
			c.SetFunc(`Modal`, func(data interface{}) template.HTML {
				return modal.Render(c, data)
			})
			ErrorPageFunc(c)
			c.SetFunc(`IndexStrSlice`, indexStrSlice)

			if !config.DefaultConfig.ConnectedDB(false) {
				return h.Handle(c)
			}
			//用户相关函数
			user, _ := c.Session().Get(`user`).(*dbschema.User)
			roleM := model.NewUserRole(c)
			var roleList []*dbschema.UserRole
			if user != nil {
				c.Set(`user`, user)
				c.SetFunc(`Username`, func() string { return user.Username })
				roleList = roleM.ListByUser(user)
				c.Set(`roleList`, roleList)
			}
			c.SetFunc(`Avatar`, func(avatar string, defaults ...string) string {
				if len(avatar) > 0 {
					return tplfunc.AddSuffix(avatar, `_200_200`)
				}
				if len(defaults) > 0 && len(defaults[0]) > 0 {
					return defaults[0]
				}
				return DefaultAvatarURL
			})
			c.SetFunc(`FileTypeByName`, helper.FileTypeByName)
			c.SetFunc(`FileTypeIcon`, helper.FileTypeIcon)
			c.SetFunc(`Project`, func(ident string) *navigate.ProjectItem {
				return navigate.ProjectGet(ident)
			})
			c.SetFunc(`ProjectIdent`, func() string {
				return navigate.ProjectIdent(c.Path())
			})
			c.SetFunc(`Projects`, func() navigate.ProjectList {
				return navigate.ProjectListAll()
			})
			c.SetFunc(`TopButtons`, func() dashboard.TopButtons {
				buttons := dashboard.TopButtonAll()
				buttons.Ready(c)
				return buttons
			})
			c.SetFunc(`GlobalFooters`, func() dashboard.GlobalFooters {
				footers := dashboard.GlobalFooterAll()
				footers.Ready(c)
				return footers
			})
			c.SetFunc(`Navigate`, func(side string) navigate.List {
				switch side {
				case `top`:
					if user != nil && user.Id == 1 {
						if navigate.TopNavigate == nil {
							return navigate.EmptyList
						}
						return *navigate.TopNavigate
					}
					return roleM.FilterNavigate(roleList, navigate.TopNavigate)
				case `left`:
					fallthrough
				default:
					var leftNav *navigate.List
					ident := navigate.ProjectIdent(c.Path())
					if len(ident) > 0 {
						if proj := navigate.ProjectGet(ident); proj != nil {
							leftNav = proj.NavList
						}
					}
					if leftNav == nil {
						if proj := navigate.ProjectFirst(); proj != nil {
							leftNav = proj.NavList
						}
					}
					if user != nil && user.Id == 1 {
						if leftNav == nil {
							return navigate.EmptyList
						}
						return *leftNav
					}
					return roleM.FilterNavigate(roleList, leftNav)
				}
			})
			return h.Handle(c)
		})
	}
}

func indexStrSlice(slice []string, index int) string {
	if slice == nil {
		return ``
	}
	if index >= len(slice) {
		return ``
	}
	return slice[index]
}

func hasString(slice []string, str string) bool {
	if slice == nil {
		return false
	}
	for _, v := range slice {
		if v == str {
			return true
		}
	}
	return false
}

func date(timestamp interface{}) time.Time {
	if v, y := timestamp.(int64); y {
		return time.Unix(v, 0)
	}
	if v, y := timestamp.(uint); y {
		return time.Unix(int64(v), 0)
	}
	v, _ := strconv.ParseInt(fmt.Sprint(timestamp), 10, 64)
	return time.Unix(v, 0)
}
<|endoftext|>


<commit_before>package main

import (
	"koding/db/mongodb"
	"labix.org/v2/mgo"
	"labix.org/v2/mgo/bson"
	"time"
)

func init() {
	registerAnalytic(numberOfAccounts)
	registerAnalytic(numberOfUsersWhoLinkedOauth)
	registerAnalytic(numberOfEmailsImportedFromGoogle)
	registerAnalytic(numberOfInvitesSent)
	registerAnalytic(numberOfUsersWhoJoinedToday)
	registerAnalytic(numberOfGuestAccountsCreatedToday)
	registerAnalytic(numberOfUsersWhoDeletedTheirAccount)
	registerAnalytic(numberOfUsersWhoPostedContentToday)
	registerAnalytic(numberOfUsersWhoAreOnline)
	registerAnalytic(numberOfUsersWhoLoggedInToday)
	registerAnalytic(numberOfGuestAccounts)
}

func numberOfAccounts() (string, int) {
	var identifier string = "number_of_accounts"
	var count int
	var err error
	var query = func(c *mgo.Collection) error {
		count, err = c.Count()

		return err
	}

	mongodb.Run("jAccounts", query)

	return identifier, count
}

func numberOfUsersWhoLinkedOauth() (string, int) {
	var identifier string = "number_of_users_who_linked_oauth"
	var count int
	var err error
	var query = func(c *mgo.Collection) error {
		count, err = c.Find(bson.M{"foreignAuth": bson.M{"$exists": true}}).Count()

		return err
	}

	mongodb.Run("jUsers", query)

	return identifier, count
}

func numberOfEmailsImportedFromGoogle() (string, int) {
	var identifier string = "number_of_emails_imported_from_google"
	var count int
	var err error
	var query = func(c *mgo.Collection) error {
		count, err = c.Find(nil).Count()

		return err
	}

	mongodb.Run("jReferrableEmails", query)

	return identifier, count
}

func numberOfInvitesSent() (string, int) {
	var identifier string = "number_of_invites_sent"
	var count int
	var err error
	var query = func(c *mgo.Collection) error {
		count, err = c.Find(bson.M{"invited": true}).Count()

		return err
	}

	mongodb.Run("jReferrableEmails", query)

	return identifier, count
}

func numberOfUsersWhoJoinedToday() (string, int) {
	var identifier string = "number_of_users_who_joined_today"
	var count int
	var err error
	var query = func(c *mgo.Collection) error {
		var today = getTodayDate()
		var filter = bson.M{"meta.createdAt": bson.M{"$gte": today}, "type": "registered"}

		count, err = c.Find(filter).Count()

		return err
	}

	mongodb.Run("jAccounts", query)

	return identifier, count
}

func numberOfGuestAccountsCreatedToday() (string, int) {
	var identifier string = "number_of_guest_accounts_created_today"
	var count int
	var err error
	var query = func(c *mgo.Collection) error {
		var today = getTodayDate()
		var filter = bson.M{"meta.createdAt": bson.M{"$gte": today}, "type": "unregistered"}

		count, err = c.Find(filter).Count()

		return err
	}

	mongodb.Run("jAccounts", query)

	return identifier, count
}

func numberOfUsersWhoDeletedTheirAccount() (string, int) {
	var identifier string = "number_of_users_who_deleted_their_account"
	var count int
	var err error
	var query = func(c *mgo.Collection) error {
		count, err = c.Find(bson.M{"status": "deleted"}).Count()

		return err
	}

	mongodb.Run("jUsers", query)

	return identifier, count
}

func numberOfUsersWhoPostedContentToday() (string, int) {
	var identifier string = "number_of_users_who_posted_content_today"
	var count int
	var err error
	var query = func(c *mgo.Collection) error {
		var today = getTodayDate()
		var filter = bson.M{"meta.modifiedAt": bson.M{"$gte": today}}

		count, err = c.Find(filter).Count()

		return err
	}

	mongodb.Run("jAccounts", query)

	return identifier, count
}

func numberOfUsersWhoAreOnline() (string, int) {
	var identifier string = "number_of_users_who_are_online"
	var count int
	var err error
	var query = func(c *mgo.Collection) error {
		count, err = c.Find(bson.M{"onlineStatus": "online"}).Count()

		return err
	}

	mongodb.Run("jAccounts", query)

	return identifier, count
}

func numberOfUsersWhoLoggedInToday() (string, int) {
	var identifier string = "number_of_users_who_logged_in_today"
	var count int
	var err error
	var query = func(c *mgo.Collection) error {
		var today = getTodayDate()
		var filter = bson.M{"lastLoginDate": bson.M{"$gte": today}}

		count, err = c.Find(filter).Count()

		return err
	}

	mongodb.Run("jUsers", query)

	return identifier, count
}

func numberOfGuestAccounts() (string, int) {
	var identifier string = "number_of_guest_accounts"
	var count int
	var err error
	var query = func(c *mgo.Collection) error {
		var filter = bson.M{"type": "unregistered"}

		count, err = c.Find(filter).Count()

		return err
	}

	mongodb.Run("jAccounts", query)

	return identifier, count
}

func getTodayDate() time.Time {
	year, month, day := time.Now().Date()
	return time.Date(year, month, day, 0, 0, 0, 0, time.Local)
}
<commit_msg>graphite: commenting online user metric till persistence worker is fixed<commit_after>package main

import (
	"koding/db/mongodb"
	"labix.org/v2/mgo"
	"labix.org/v2/mgo/bson"
	"time"
)

func init() {
	registerAnalytic(numberOfAccounts)
	registerAnalytic(numberOfUsersWhoLinkedOauth)
	registerAnalytic(numberOfEmailsImportedFromGoogle)
	registerAnalytic(numberOfInvitesSent)
	registerAnalytic(numberOfUsersWhoJoinedToday)
	registerAnalytic(numberOfGuestAccountsCreatedToday)
	registerAnalytic(numberOfUsersWhoDeletedTheirAccount)
	registerAnalytic(numberOfUsersWhoPostedContentToday)
	// commenting this out till persistence worker is fixed
	//registerAnalytic(numberOfUsersWhoAreOnline)
	registerAnalytic(numberOfUsersWhoLoggedInToday)
	registerAnalytic(numberOfGuestAccounts)
}

func numberOfAccounts() (string, int) {
	var identifier string = "number_of_accounts"
	var count int
	var err error
	var query = func(c *mgo.Collection) error {
		count, err = c.Count()

		return err
	}

	mongodb.Run("jAccounts", query)

	return identifier, count
}

func numberOfUsersWhoLinkedOauth() (string, int) {
	var identifier string = "number_of_users_who_linked_oauth"
	var count int
	var err error
	var query = func(c *mgo.Collection) error {
		count, err = c.Find(bson.M{"foreignAuth": bson.M{"$exists": true}}).Count()

		return err
	}

	mongodb.Run("jUsers", query)

	return identifier, count
}

func numberOfEmailsImportedFromGoogle() (string, int) {
	var identifier string = "number_of_emails_imported_from_google"
	var count int
	var err error
	var query = func(c *mgo.Collection) error {
		count, err = c.Find(nil).Count()

		return err
	}

	mongodb.Run("jReferrableEmails", query)

	return identifier, count
}

func numberOfInvitesSent() (string, int) {
	var identifier string = "number_of_invites_sent"
	var count int
	var err error
	var query = func(c *mgo.Collection) error {
		count, err = c.Find(bson.M{"invited": true}).Count()

		return err
	}

	mongodb.Run("jReferrableEmails", query)

	return identifier, count
}

func numberOfUsersWhoJoinedToday() (string, int) {
	var identifier string = "number_of_users_who_joined_today"
	var count int
	var err error
	var query = func(c *mgo.Collection) error {
		var today = getTodayDate()
		var filter = bson.M{"meta.createdAt": bson.M{"$gte": today}, "type": "registered"}

		count, err = c.Find(filter).Count()

		return err
	}

	mongodb.Run("jAccounts", query)

	return identifier, count
}

func numberOfGuestAccountsCreatedToday() (string, int) {
	var identifier string = "number_of_guest_accounts_created_today"
	var count int
	var err error
	var query = func(c *mgo.Collection) error {
		var today = getTodayDate()
		var filter = bson.M{"meta.createdAt": bson.M{"$gte": today}, "type": "unregistered"}

		count, err = c.Find(filter).Count()

		return err
	}

	mongodb.Run("jAccounts", query)

	return identifier, count
}

func numberOfUsersWhoDeletedTheirAccount() (string, int) {
	var identifier string = "number_of_users_who_deleted_their_account"
	var count int
	var err error
	var query = func(c *mgo.Collection) error {
		count, err = c.Find(bson.M{"status": "deleted"}).Count()

		return err
	}

	mongodb.Run("jUsers", query)

	return identifier, count
}

func numberOfUsersWhoPostedContentToday() (string, int) {
	var identifier string = "number_of_users_who_posted_content_today"
	var count int
	var err error
	var query = func(c *mgo.Collection) error {
		var today = getTodayDate()
		var filter = bson.M{"meta.modifiedAt": bson.M{"$gte": today}}

		count, err = c.Find(filter).Count()

		return err
	}

	mongodb.Run("jAccounts", query)

	return identifier, count
}

func numberOfUsersWhoAreOnline() (string, int) {
	var identifier string = "number_of_users_who_are_online"
	var count int
	var err error
	var query = func(c *mgo.Collection) error {
		count, err = c.Find(bson.M{"onlineStatus": "online"}).Count()

		return err
	}

	mongodb.Run("jAccounts", query)

	return identifier, count
}

func numberOfUsersWhoLoggedInToday() (string, int) {
	var identifier string = "number_of_users_who_logged_in_today"
	var count int
	var err error
	var query = func(c *mgo.Collection) error {
		var today = getTodayDate()
		var filter = bson.M{"lastLoginDate": bson.M{"$gte": today}}

		count, err = c.Find(filter).Count()

		return err
	}

	mongodb.Run("jUsers", query)

	return identifier, count
}

func numberOfGuestAccounts() (string, int) {
	var identifier string = "number_of_guest_accounts"
	var count int
	var err error
	var query = func(c *mgo.Collection) error {
		var filter = bson.M{"type": "unregistered"}

		count, err = c.Find(filter).Count()

		return err
	}

	mongodb.Run("jAccounts", query)

	return identifier, count
}

func getTodayDate() time.Time {
	year, month, day := time.Now().Date()
	return time.Date(year, month, day, 0, 0, 0, 0, time.Local)
}
<|endoftext|>


<commit_before>package main

import (
	"koding/db/models"
	"koding/db/mongodb/modelhelper"
	"koding/db/mongodb/modelhelper/modeltesthelper"
	"testing"
	"time"

	"labix.org/v2/mgo/bson"

	. "github.com/smartystreets/goconvey/convey"
)

func TestReleaseUser(t *testing.T) {
	Convey("Given user", t, func() {
		username := "releaseuser"
		user, _, err := modeltesthelper.CreateUser(username)

		warning := &Warning{
			ID:     "warning",
			Select: []bson.M{bson.M{"username": user.Name}},
		}

		_, err = warning.FindAndLockUser()
		So(err, ShouldBeNil)

		Convey("When it releases user", func() {
			err = warning.ReleaseUser(user)
			So(err, ShouldBeNil)

			updatedUser, err := modelhelper.GetUser(user.Name)
			So(err, ShouldBeNil)

			Convey("Then it updates warning", func() {
				veryRecently := timeNow().Add(-1 * time.Second)
				So(updatedUser.Inactive.ModifiedAt.UTC(), ShouldHappenAfter,
					veryRecently)
			})
		})

		Reset(func() {
			deleteUserWithUsername(user)
		})
	})
}

func TestLockAndReleaseUser(t *testing.T) {
	Convey("Given user", t, func() {
		username := "releaseuser"
		user, _, err := modeltesthelper.CreateUser(username)
		So(err, ShouldBeNil)

		warning := &Warning{
			ID:     "warning",
			Select: []bson.M{bson.M{"username": user.Name}},
		}

		Convey("Then it shouldn't user if it was processed today", func() {
			err := updateUserModifiedAt(user, timeNow().Add(time.Hour*24))
			So(err, ShouldBeNil)

			_, err = warning.FindAndLockUser()
			So(err, ShouldNotBeNil)
		})

		Convey("Then it should find and lock user", func() {
			err := updateUserModifiedAt(user, timeNow().Add(-time.Hour*24))
			So(err, ShouldBeNil)

			newuser, err := warning.FindAndLockUser()
			So(err, ShouldBeNil)

			So(newuser.ObjectId, ShouldEqual, user.ObjectId)
			So(newuser.Inactive.Assigned, ShouldBeTrue)

			Convey("When it releases user", func() {
				err := warning.UpdateAndReleaseUser(user.ObjectId)
				So(err, ShouldBeNil)

				Convey("Then it should update user for next level", func() {
					updatedUser, err := modelhelper.GetUser(user.Name)
					So(err, ShouldBeNil)

					So(updatedUser.Inactive.Warning, ShouldEqual, warning.ID)
					So(updatedUser.Inactive.ModifiedAt.IsZero(), ShouldBeFalse)

					warnings := updatedUser.Inactive.Warnings

					So(warnings, ShouldNotBeNil)
					So(warnings[warning.ID], ShouldNotBeNil)
					So(warnings[warning.ID].IsZero(), ShouldBeFalse)

					Convey("Then it shouldn't get same user again", func() {
						_, err := warning.FindAndLockUser()
						So(err, ShouldNotBeNil)
					})
				})
			})
		})

		Reset(func() {
			deleteUserWithUsername(user)
		})
	})
}

func TestIsUserExempt(t *testing.T) {
	Convey("Given exempt conditions", t, func() {
		username := "exempt"
		user, _, err := modeltesthelper.CreateUser(username)
		So(err, ShouldBeNil)

		Convey("Then it should be exempt", func() {
			warning := &Warning{
				ExemptCheckers: []*ExemptChecker{
					{
						Name: "",
						IsExempt: func(_ *models.User, _ *Warning) (bool, error) {
							return true, nil
						},
					},
				},
			}

			isExempt, err := warning.IsUserExempt(user)

			So(err, ShouldBeNil)
			So(isExempt, ShouldBeTrue)
		})

		Convey("Then it should not be exempt", func() {
			warning := &Warning{
				ExemptCheckers: []*ExemptChecker{
					{
						Name: "",
						IsExempt: func(_ *models.User, _ *Warning) (bool, error) {
							return false, nil
						},
					},
				},
			}

			isExempt, err := warning.IsUserExempt(user)

			So(err, ShouldBeNil)
			So(isExempt, ShouldBeFalse)
		})

		Convey("Then it should be exempt even if one condition is exempt", func() {
			warning := &Warning{
				ExemptCheckers: []*ExemptChecker{
					{
						Name: "",
						IsExempt: func(_ *models.User, _ *Warning) (bool, error) {
							return false, nil
						},
					},
					{
						Name: "",
						IsExempt: func(_ *models.User, _ *Warning) (bool, error) {
							return true, nil
						},
					},
				},
			}

			isExempt, err := warning.IsUserExempt(user)

			So(err, ShouldBeNil)
			So(isExempt, ShouldBeTrue)
		})

		Reset(func() {
			deleteUserWithUsername(user)
		})
	})
}

func TestAct(t *testing.T) {
	Convey("Given action", t, func() {
		Convey("Then it should call action function", func() {
			username := "actionuser"
			user, _, err := modeltesthelper.CreateUser(username)
			So(err, ShouldBeNil)

			var called = false

			warning := &Warning{
				Action: func(user *models.User, _ string) error {
					called = true
					return nil
				},
			}

			err = warning.Act(user)
			So(err, ShouldBeNil)
			So(called, ShouldBeTrue)

			Reset(func() {
				deleteUserWithUsername(user)
			})
		})
	})
}

func TestGetCount(t *testing.T) {
	Convey("It should return count", t, func() {
		user1, err := createInactiveUser(21)
		So(err, ShouldBeNil)

		user2, err := createInactiveUser(21)
		So(err, ShouldBeNil)

		warning := VMDeletionWarning1

		count, err := modelhelper.CountUsersByQuery(warning.buildSelectQuery())
		So(err, ShouldBeNil)

		So(count, ShouldEqual, 2)

		Convey("It should return minimum sleep", func() {
			sleepTime, err := warning.getSleepTime()
			So(err, ShouldBeNil)

			So(sleepTime, ShouldEqual, 1*time.Second)

			Reset(func() {
				deleteUserWithUsername(user1)
				deleteUserWithUsername(user2)
			})
		})
	})
}
<commit_msg>janitor: fix tests<commit_after>package main

import (
	"koding/db/models"
	"koding/db/mongodb/modelhelper"
	"koding/db/mongodb/modelhelper/modeltesthelper"
	"testing"
	"time"

	"labix.org/v2/mgo/bson"

	. "github.com/smartystreets/goconvey/convey"
)

func TestReleaseUser(t *testing.T) {
	Convey("Given user", t, func() {
		username := "releaseuser"
		user, _, err := modeltesthelper.CreateUser(username)

		warning := &Warning{
			ID:     "warning",
			Select: []bson.M{bson.M{"username": user.Name}},
		}

		_, err = warning.FindAndLockUser()
		So(err, ShouldBeNil)

		Convey("When it releases user", func() {
			err = warning.ReleaseUser(user)
			So(err, ShouldBeNil)

			updatedUser, err := modelhelper.GetUser(user.Name)
			So(err, ShouldBeNil)

			Convey("Then it updates warning", func() {
				veryRecently := timeNow().Add(-1 * time.Second)
				So(updatedUser.Inactive.ModifiedAt.UTC(), ShouldHappenAfter,
					veryRecently)
			})
		})

		Reset(func() {
			deleteUserWithUsername(user)
		})
	})
}

func TestLockAndReleaseUser(t *testing.T) {
	Convey("Given user", t, func() {
		username := "releaseuser"
		user, _, err := modeltesthelper.CreateUser(username)
		So(err, ShouldBeNil)

		warning := &Warning{
			ID:     "warning",
			Select: []bson.M{bson.M{"username": user.Name}},
		}

		Convey("Then it shouldn't user if it was processed today", func() {
			err := updateUserModifiedAt(user, timeNow().Add(time.Hour*24))
			So(err, ShouldBeNil)

			_, err = warning.FindAndLockUser()
			So(err, ShouldNotBeNil)
		})

		Convey("Then it should find and lock user", func() {
			err := updateUserModifiedAt(user, timeNow().Add(-time.Hour*24))
			So(err, ShouldBeNil)

			newuser, err := warning.FindAndLockUser()
			So(err, ShouldBeNil)

			So(newuser.ObjectId, ShouldEqual, user.ObjectId)
			So(newuser.Inactive.Assigned, ShouldBeTrue)

			Convey("When it releases user", func() {
				err := warning.UpdateAndReleaseUser(user.ObjectId)
				So(err, ShouldBeNil)

				Convey("Then it should update user for next level", func() {
					updatedUser, err := modelhelper.GetUser(user.Name)
					So(err, ShouldBeNil)

					So(updatedUser.Inactive.Warning, ShouldEqual, warning.ID)
					So(updatedUser.Inactive.ModifiedAt.IsZero(), ShouldBeFalse)

					warnings := updatedUser.Inactive.Warnings

					So(warnings, ShouldNotBeNil)
					So(warnings[warning.ID], ShouldNotBeNil)
					So(warnings[warning.ID].IsZero(), ShouldBeFalse)

					Convey("Then it shouldn't get same user again", func() {
						_, err := warning.FindAndLockUser()
						So(err, ShouldNotBeNil)
					})
				})
			})
		})

		Reset(func() {
			deleteUserWithUsername(user)
		})
	})
}

func TestIsUserExempt(t *testing.T) {
	Convey("Given exempt conditions", t, func() {
		username := "exempt"
		user, _, err := modeltesthelper.CreateUser(username)
		So(err, ShouldBeNil)

		Convey("Then it should be exempt", func() {
			warning := &Warning{
				ExemptCheckers: []*ExemptChecker{
					{
						Name: "",
						IsExempt: func(_ *models.User, _ *Warning) (bool, error) {
							return true, nil
						},
					},
				},
			}

			isExempt, err := warning.IsUserExempt(user)

			So(err, ShouldBeNil)
			So(isExempt, ShouldBeTrue)
		})

		Convey("Then it should not be exempt", func() {
			warning := &Warning{
				ExemptCheckers: []*ExemptChecker{
					{
						Name: "",
						IsExempt: func(_ *models.User, _ *Warning) (bool, error) {
							return false, nil
						},
					},
				},
			}

			isExempt, err := warning.IsUserExempt(user)

			So(err, ShouldBeNil)
			So(isExempt, ShouldBeFalse)
		})

		Convey("Then it should be exempt even if one condition is exempt", func() {
			warning := &Warning{
				ExemptCheckers: []*ExemptChecker{
					{
						Name: "",
						IsExempt: func(_ *models.User, _ *Warning) (bool, error) {
							return false, nil
						},
					},
					{
						Name: "",
						IsExempt: func(_ *models.User, _ *Warning) (bool, error) {
							return true, nil
						},
					},
				},
			}

			isExempt, err := warning.IsUserExempt(user)

			So(err, ShouldBeNil)
			So(isExempt, ShouldBeTrue)
		})

		Reset(func() {
			deleteUserWithUsername(user)
		})
	})
}

func TestAct(t *testing.T) {
	Convey("Given action", t, func() {
		Convey("Then it should call action function", func() {
			username := "actionuser"
			user, _, err := modeltesthelper.CreateUser(username)
			So(err, ShouldBeNil)

			var called = false

			warning := &Warning{
				Action: func(user *models.User, _ string) error {
					called = true
					return nil
				},
			}

			err = warning.Act(user)
			So(err, ShouldBeNil)
			So(called, ShouldBeTrue)

			Reset(func() {
				deleteUserWithUsername(user)
			})
		})
	})
}

func TestGetCount(t *testing.T) {
	Convey("It should return count", t, func() {
		user1, err := createInactiveUser(21)
		So(err, ShouldBeNil)

		user2, err := createInactiveUser(21)
		So(err, ShouldBeNil)

		warning := VMDeletionWarning1

		count, err := modelhelper.CountUsersByQuery(warning.buildSelectQuery())
		So(err, ShouldBeNil)

		So(count, ShouldEqual, 2)

		Convey("It should return minimum sleep", func() {
			sleepTime, err := warning.getSleepTime()
			So(err, ShouldBeNil)

			So(sleepTime, ShouldEqual, 20*time.Second)

			Reset(func() {
				deleteUserWithUsername(user1)
				deleteUserWithUsername(user2)
			})
		})
	})
}
<|endoftext|>


<commit_before>/*
Copyright 2021 The Vitess 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 planbuilder

import (
	"vitess.io/vitess/go/sqltypes"
	querypb "vitess.io/vitess/go/vt/proto/query"
	vtrpcpb "vitess.io/vitess/go/vt/proto/vtrpc"
	"vitess.io/vitess/go/vt/sqlparser"
	"vitess.io/vitess/go/vt/vterrors"
	"vitess.io/vitess/go/vt/vtgate/engine"
)

func gen4SlowPlanner(query string) func(sqlparser.Statement, *sqlparser.ReservedVars, ContextVSchema) (engine.Primitive, error) {
	return func(statement sqlparser.Statement, vars *sqlparser.ReservedVars, schema ContextVSchema) (engine.Primitive, error) {
		defer schema.SetPlannerVersion(Gen4Slow)

		schema.SetPlannerVersion(V3)
		v3Primitive, v3Err := createInstructionFor(query, statement, sqlparser.NewReservedVars("", make(map[string]struct{})), schema, false, false)

		schema.SetPlannerVersion(Gen4)
		gen4Primitive, gen4Err := createInstructionFor(query, statement, sqlparser.NewReservedVars("", make(map[string]struct{})), schema, false, false)
		if v3Err != nil && gen4Err != nil {
			return nil, gen4Err
		}
		if v3Err == nil && gen4Err != nil {
			return nil, vterrors.Errorf(vtrpcpb.Code_INTERNAL, "Gen4 failed while v3 did not: %s", gen4Err.Error())
		}
		if v3Err != nil && gen4Err == nil {
			return nil, vterrors.Errorf(vtrpcpb.Code_INTERNAL, "v3 failed while Gen4 did not: %s", v3Err.Error())
		}
		return &comparer{v3: v3Primitive, gen4: gen4Primitive}, nil
	}
}

type comparer struct {
	v3, gen4 engine.Primitive
}

func (c *comparer) RouteType() string {
	return c.gen4.RouteType()
}

func (c *comparer) GetKeyspaceName() string {
	return c.gen4.GetKeyspaceName()
}

func (c *comparer) GetTableName() string {
	return c.gen4.GetTableName()
}

func (c *comparer) GetFields(vcursor engine.VCursor, bindVars map[string]*querypb.BindVariable) (*sqltypes.Result, error) {
	return c.gen4.GetFields(vcursor, bindVars)
}

func (c *comparer) NeedsTransaction() bool {
	return c.gen4.NeedsTransaction()
}

func (c *comparer) TryExecute(vcursor engine.VCursor, bindVars map[string]*querypb.BindVariable, wantfields bool) (*sqltypes.Result, error) {
	v3Result, v3Err := c.v3.TryExecute(vcursor, bindVars, wantfields)
	gen4Result, gen4Err := c.gen4.TryExecute(vcursor, bindVars, wantfields)
	if v3Err != nil && gen4Err != nil {
		return nil, gen4Err
	}
	if v3Err == nil && gen4Err != nil {
		return nil, vterrors.Errorf(vtrpcpb.Code_INTERNAL, "Gen4 failed while v3 did not: %s", gen4Err.Error())
	}
	if v3Err != nil && gen4Err == nil {
		return nil, vterrors.Errorf(vtrpcpb.Code_INTERNAL, "v3 failed while Gen4 did not: %s", v3Err.Error())
	}
	match := sqltypes.ResultsEqualUnordered([]sqltypes.Result{*v3Result}, []sqltypes.Result{*gen4Result})
	if !match {
		return nil, vterrors.Errorf(vtrpcpb.Code_INTERNAL, "results did not match")
	}
	return gen4Result, nil
}

func (c *comparer) TryStreamExecute(vcursor engine.VCursor, bindVars map[string]*querypb.BindVariable, wantfields bool, callback func(*sqltypes.Result) error) error {
	return c.v3.TryStreamExecute(vcursor, bindVars, wantfields, callback)
}

func (c *comparer) Inputs() []engine.Primitive {
	return c.gen4.Inputs()
}

func (c *comparer) Description() engine.PrimitiveDescription {
	return c.gen4.Description()
}

var _ engine.Primitive = (*comparer)(nil)
<commit_msg>Better logging for gen4slow planner<commit_after>/*
Copyright 2021 The Vitess 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 planbuilder

import (
	"vitess.io/vitess/go/sqltypes"
	querypb "vitess.io/vitess/go/vt/proto/query"
	vtrpcpb "vitess.io/vitess/go/vt/proto/vtrpc"
	"vitess.io/vitess/go/vt/sqlparser"
	"vitess.io/vitess/go/vt/vterrors"
	"vitess.io/vitess/go/vt/vtgate/engine"
)

func gen4SlowPlanner(query string) func(sqlparser.Statement, *sqlparser.ReservedVars, ContextVSchema) (engine.Primitive, error) {
	return func(statement sqlparser.Statement, vars *sqlparser.ReservedVars, schema ContextVSchema) (engine.Primitive, error) {
		defer schema.SetPlannerVersion(Gen4Slow)

		schema.SetPlannerVersion(Gen4)
		gen4Primitive, gen4Err := createInstructionFor(query, statement, vars, schema, false, false)

		// We insert data only once using the gen4 planner to avoid duplicated rows.
		if _, isInsert := statement.(*sqlparser.Insert); isInsert {
			return gen4Primitive, gen4Err
		}

		schema.SetPlannerVersion(V3)
		v3Primitive, v3Err := createInstructionFor(query, statement, vars, schema, false, false)

		err := treatV3AndGen4Errors(v3Err, gen4Err)
		if err != nil {
			return nil, err
		}
		return &comparer{v3: v3Primitive, gen4: gen4Primitive}, nil
	}
}

func treatV3AndGen4Errors(v3Err error, gen4Err error) error {
	if v3Err != nil && gen4Err != nil {
		return vterrors.Errorf(vtrpcpb.Code_INTERNAL, "v3 and Gen4 failed: v3: %s | Gen4: %s", v3Err.Error(), gen4Err.Error())
	}
	if v3Err == nil && gen4Err != nil {
		return vterrors.Errorf(vtrpcpb.Code_INTERNAL, "Gen4 failed while v3 did not: %s", gen4Err.Error())
	}
	if v3Err != nil && gen4Err == nil {
		return vterrors.Errorf(vtrpcpb.Code_INTERNAL, "v3 failed while Gen4 did not: %s", v3Err.Error())
	}
	return nil
}

type comparer struct {
	v3, gen4 engine.Primitive
}

func (c *comparer) RouteType() string {
	return c.gen4.RouteType()
}

func (c *comparer) GetKeyspaceName() string {
	return c.gen4.GetKeyspaceName()
}

func (c *comparer) GetTableName() string {
	return c.gen4.GetTableName()
}

func (c *comparer) GetFields(vcursor engine.VCursor, bindVars map[string]*querypb.BindVariable) (*sqltypes.Result, error) {
	return c.gen4.GetFields(vcursor, bindVars)
}

func (c *comparer) NeedsTransaction() bool {
	return c.gen4.NeedsTransaction()
}

func (c *comparer) TryExecute(vcursor engine.VCursor, bindVars map[string]*querypb.BindVariable, wantfields bool) (*sqltypes.Result, error) {
	gen4Result, gen4Err := c.gen4.TryExecute(vcursor, bindVars, wantfields)
	v3Result, v3Err := c.v3.TryExecute(vcursor, bindVars, wantfields)
	err := treatV3AndGen4Errors(v3Err, gen4Err)
	if err != nil {
		return nil, err
	}
	match := sqltypes.ResultsEqualUnordered([]sqltypes.Result{*v3Result}, []sqltypes.Result{*gen4Result})
	if !match {
		return nil, vterrors.Errorf(vtrpcpb.Code_INTERNAL, "results did not match")
	}
	return gen4Result, nil
}

func (c *comparer) TryStreamExecute(vcursor engine.VCursor, bindVars map[string]*querypb.BindVariable, wantfields bool, callback func(*sqltypes.Result) error) error {
	return c.v3.TryStreamExecute(vcursor, bindVars, wantfields, callback)
}

func (c *comparer) Inputs() []engine.Primitive {
	return c.gen4.Inputs()
}

func (c *comparer) Description() engine.PrimitiveDescription {
	return c.gen4.Description()
}

var _ engine.Primitive = (*comparer)(nil)
<|endoftext|>


<commit_before>// Copyright 2018 The Nakama 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 server

import (
	"bytes"
	"errors"
	"fmt"
	"sync"
	"time"

	"net"

	"github.com/golang/protobuf/jsonpb"
	"github.com/gorilla/websocket"
	"github.com/heroiclabs/nakama/rtapi"
	"github.com/satori/go.uuid"
	"go.uber.org/atomic"
	"go.uber.org/zap"
)

var ErrSessionQueueFull = errors.New("session outgoing queue full")

type sessionWS struct {
	sync.Mutex
	logger   *zap.Logger
	config   Config
	id       uuid.UUID
	userID   uuid.UUID
	username *atomic.String
	expiry   int64

	jsonpbMarshaler        *jsonpb.Marshaler
	jsonpbUnmarshaler      *jsonpb.Unmarshaler
	queuePriorityThreshold int

	sessionRegistry *SessionRegistry
	matchmaker      Matchmaker
	tracker         Tracker

	stopped        bool
	conn           *websocket.Conn
	pingTicker     *time.Ticker
	outgoingCh     chan []byte
	outgoingStopCh chan struct{}
}

func NewSessionWS(logger *zap.Logger, config Config, userID uuid.UUID, username string, expiry int64, jsonpbMarshaler *jsonpb.Marshaler, jsonpbUnmarshaler *jsonpb.Unmarshaler, conn *websocket.Conn, sessionRegistry *SessionRegistry, matchmaker Matchmaker, tracker Tracker) Session {
	sessionID := uuid.Must(uuid.NewV4())
	sessionLogger := logger.With(zap.String("uid", userID.String()), zap.String("sid", sessionID.String()))

	sessionLogger.Info("New WebSocket session connected")

	return &sessionWS{
		logger:   sessionLogger,
		config:   config,
		id:       sessionID,
		userID:   userID,
		username: atomic.NewString(username),
		expiry:   expiry,

		jsonpbMarshaler:        jsonpbMarshaler,
		jsonpbUnmarshaler:      jsonpbUnmarshaler,
		queuePriorityThreshold: (config.GetSocket().OutgoingQueueSize / 3) * 2,

		sessionRegistry: sessionRegistry,
		matchmaker:      matchmaker,
		tracker:         tracker,

		stopped:        false,
		conn:           conn,
		pingTicker:     time.NewTicker(time.Duration(config.GetSocket().PingPeriodMs) * time.Millisecond),
		outgoingCh:     make(chan []byte, config.GetSocket().OutgoingQueueSize),
		outgoingStopCh: make(chan struct{}),
	}
}

func (s *sessionWS) Logger() *zap.Logger {
	return s.logger
}

func (s *sessionWS) ID() uuid.UUID {
	return s.id
}

func (s *sessionWS) UserID() uuid.UUID {
	return s.userID
}

func (s *sessionWS) Username() string {
	return s.username.Load()
}

func (s *sessionWS) SetUsername(username string) {
	s.username.Store(username)
}

func (s *sessionWS) Expiry() int64 {
	return s.expiry
}

func (s *sessionWS) Consume(processRequest func(logger *zap.Logger, session Session, envelope *rtapi.Envelope) bool) {
	defer s.cleanupClosedConnection()
	s.conn.SetReadLimit(s.config.GetSocket().MaxMessageSizeBytes)
	s.conn.SetReadDeadline(time.Now().Add(time.Duration(s.config.GetSocket().PongWaitMs) * time.Millisecond))
	s.conn.SetPongHandler(func(string) error {
		s.conn.SetReadDeadline(time.Now().Add(time.Duration(s.config.GetSocket().PongWaitMs) * time.Millisecond))
		return nil
	})

	// Send an initial ping immediately.
	if !s.pingNow() {
		// If the first ping fails abort the rest of the consume sequence immediately.
		return
	}

	// Start a routine to process outbound messages.
	go s.processOutgoing()

	for {
		_, data, err := s.conn.ReadMessage()
		if err != nil {
			// Ignore "normal" WebSocket errors.
			if !websocket.IsCloseError(err, websocket.CloseNormalClosure, websocket.CloseGoingAway, websocket.CloseNoStatusReceived) {
				// Ignore underlying connection being shut down while read is waiting for data.
				if e, ok := err.(*net.OpError); !ok || e.Err.Error() != "use of closed network connection" {
					s.logger.Warn("Error reading message from client", zap.Error(err))
				}
			}
			break
		}

		request := &rtapi.Envelope{}
		if err = s.jsonpbUnmarshaler.Unmarshal(bytes.NewReader(data), request); err != nil {
			// If the payload is malformed the client is incompatible or misbehaving, either way disconnect it now.
			s.logger.Warn("Received malformed payload", zap.String("data", string(data)))
			break
		} else {
			// TODO Add session-global context here to cancel in-progress operations when the session is closed.
			requestLogger := s.logger.With(zap.String("cid", request.Cid))
			if !processRequest(requestLogger, s, request) {
				break
			}
		}
	}
}

func (s *sessionWS) processOutgoing() {
	for {
		select {
		case <-s.outgoingStopCh:
			// Session is closing, close the outgoing process routine.
			return
		case <-s.pingTicker.C:
			// Periodically send pings.
			if !s.pingNow() {
				// If ping fails the session will be stopped, clean up the loop.
				return
			}
		case payload := <-s.outgoingCh:
			s.Lock()
			if s.stopped {
				// The connection may have stopped between the payload being queued on the outgoing channel and reaching here.
				// If that's the case then abort outgoing processing at this point and exit.
				s.Unlock()
				return
			}
			// Process the outgoing message queue.
			s.conn.SetWriteDeadline(time.Now().Add(time.Duration(s.config.GetSocket().WriteWaitMs) * time.Millisecond))
			if err := s.conn.WriteMessage(websocket.TextMessage, payload); err != nil {
				s.Unlock()
				s.logger.Warn("Could not write message", zap.Error(err))
				return
			}
			s.Unlock()
		}
	}
}

func (s *sessionWS) pingNow() bool {
	s.Lock()
	if s.stopped {
		s.Unlock()
		return false
	}
	s.conn.SetWriteDeadline(time.Now().Add(time.Duration(s.config.GetSocket().WriteWaitMs) * time.Millisecond))
	err := s.conn.WriteMessage(websocket.PingMessage, []byte{})
	s.Unlock()
	if err != nil {
		s.logger.Warn("Could not send ping, closing channel", zap.String("remoteAddress", s.conn.RemoteAddr().String()), zap.Error(err))
		// The connection has already failed.
		s.cleanupClosedConnection()
		return false
	}
	return true
}

func (s *sessionWS) Format() SessionFormat {
	return SessionFormatJson
}

func (s *sessionWS) Send(isStream bool, mode uint8, envelope *rtapi.Envelope) error {
	payload, err := s.jsonpbMarshaler.MarshalToString(envelope)
	if err != nil {
		s.logger.Warn("Could not marshal to json", zap.Error(err))
		return err
	}

	if s.logger.Core().Enabled(zap.DebugLevel) {
		switch envelope.Message.(type) {
		case *rtapi.Envelope_Error:
			s.logger.Debug("Sending error message", zap.String("payload", payload))
		default:
			s.logger.Debug(fmt.Sprintf("Sending %T message", envelope.Message), zap.String("payload", payload))
		}
	}

	return s.SendBytes(isStream, mode, []byte(payload))
}

func (s *sessionWS) SendBytes(isStream bool, mode uint8, payload []byte) error {
	s.Lock()
	if s.stopped {
		s.Unlock()
		return nil
	}

	if isStream {
		switch mode {
		case StreamModeChannel:
			fallthrough
		case StreamModeGroup:
			fallthrough
		case StreamModeDM:
			// Chat messages are only allowed if the current outgoing queue size is below the threshold to
			// ensure there is always room to queue higher priority messages.
			if len(s.outgoingCh) < s.queuePriorityThreshold {
				s.outgoingCh <- payload
			}
			s.Unlock()
			return nil
		}
	}

	// By default attempt to queue messages and observe failures.
	select {
	case s.outgoingCh <- payload:
		s.Unlock()
		return nil
	default:
		// The outgoing queue is full, likely because the remote client can't keep up.
		// Terminate the connection immediately because the only alternative that doesn't block the server is
		// to start dropping messages, which might cause unexpected behaviour.
		s.Unlock()
		s.logger.Warn("Could not write message, session outgoing queue full")
		s.cleanupClosedConnection()
		return ErrSessionQueueFull
	}
}

func (s *sessionWS) cleanupClosedConnection() {
	s.Lock()
	if s.stopped {
		s.Unlock()
		return
	}
	s.stopped = true
	s.Unlock()

	if s.logger.Core().Enabled(zap.DebugLevel) {
		s.logger.Info("Cleaning up closed client connection", zap.String("remoteAddress", s.conn.RemoteAddr().String()))
	}

	// When connection close originates internally in the session, ensure cleanup of external resources and references.
	s.sessionRegistry.remove(s.id)
	s.matchmaker.RemoveAll(s.id)
	s.tracker.UntrackAll(s.id)

	// Clean up internals.
	s.pingTicker.Stop()
	close(s.outgoingStopCh)
	close(s.outgoingCh)

	// Close WebSocket.
	s.conn.Close()
	s.logger.Info("Closed client connection")
}

func (s *sessionWS) Close() {
	s.Lock()
	if s.stopped {
		s.Unlock()
		return
	}
	s.stopped = true
	s.Unlock()

	// Expect the caller of this session.Close() to clean up external resources (like presences) separately.

	// Clean up internals.
	s.pingTicker.Stop()
	close(s.outgoingStopCh)
	close(s.outgoingCh)

	// Send close message.
	err := s.conn.WriteControl(websocket.CloseMessage, []byte{}, time.Now().Add(time.Duration(s.config.GetSocket().WriteWaitMs)*time.Millisecond))
	if err != nil {
		s.logger.Warn("Could not send close message, closing prematurely", zap.String("remoteAddress", s.conn.RemoteAddr().String()), zap.Error(err))
	}

	// Close WebSocket.
	s.conn.Close()
	s.logger.Info("Closed client connection")
}
<commit_msg>Update log level for client disconnect message.<commit_after>// Copyright 2018 The Nakama 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 server

import (
	"bytes"
	"errors"
	"fmt"
	"sync"
	"time"

	"net"

	"github.com/golang/protobuf/jsonpb"
	"github.com/gorilla/websocket"
	"github.com/heroiclabs/nakama/rtapi"
	"github.com/satori/go.uuid"
	"go.uber.org/atomic"
	"go.uber.org/zap"
)

var ErrSessionQueueFull = errors.New("session outgoing queue full")

type sessionWS struct {
	sync.Mutex
	logger   *zap.Logger
	config   Config
	id       uuid.UUID
	userID   uuid.UUID
	username *atomic.String
	expiry   int64

	jsonpbMarshaler        *jsonpb.Marshaler
	jsonpbUnmarshaler      *jsonpb.Unmarshaler
	queuePriorityThreshold int

	sessionRegistry *SessionRegistry
	matchmaker      Matchmaker
	tracker         Tracker

	stopped        bool
	conn           *websocket.Conn
	pingTicker     *time.Ticker
	outgoingCh     chan []byte
	outgoingStopCh chan struct{}
}

func NewSessionWS(logger *zap.Logger, config Config, userID uuid.UUID, username string, expiry int64, jsonpbMarshaler *jsonpb.Marshaler, jsonpbUnmarshaler *jsonpb.Unmarshaler, conn *websocket.Conn, sessionRegistry *SessionRegistry, matchmaker Matchmaker, tracker Tracker) Session {
	sessionID := uuid.Must(uuid.NewV4())
	sessionLogger := logger.With(zap.String("uid", userID.String()), zap.String("sid", sessionID.String()))

	sessionLogger.Info("New WebSocket session connected")

	return &sessionWS{
		logger:   sessionLogger,
		config:   config,
		id:       sessionID,
		userID:   userID,
		username: atomic.NewString(username),
		expiry:   expiry,

		jsonpbMarshaler:        jsonpbMarshaler,
		jsonpbUnmarshaler:      jsonpbUnmarshaler,
		queuePriorityThreshold: (config.GetSocket().OutgoingQueueSize / 3) * 2,

		sessionRegistry: sessionRegistry,
		matchmaker:      matchmaker,
		tracker:         tracker,

		stopped:        false,
		conn:           conn,
		pingTicker:     time.NewTicker(time.Duration(config.GetSocket().PingPeriodMs) * time.Millisecond),
		outgoingCh:     make(chan []byte, config.GetSocket().OutgoingQueueSize),
		outgoingStopCh: make(chan struct{}),
	}
}

func (s *sessionWS) Logger() *zap.Logger {
	return s.logger
}

func (s *sessionWS) ID() uuid.UUID {
	return s.id
}

func (s *sessionWS) UserID() uuid.UUID {
	return s.userID
}

func (s *sessionWS) Username() string {
	return s.username.Load()
}

func (s *sessionWS) SetUsername(username string) {
	s.username.Store(username)
}

func (s *sessionWS) Expiry() int64 {
	return s.expiry
}

func (s *sessionWS) Consume(processRequest func(logger *zap.Logger, session Session, envelope *rtapi.Envelope) bool) {
	defer s.cleanupClosedConnection()
	s.conn.SetReadLimit(s.config.GetSocket().MaxMessageSizeBytes)
	s.conn.SetReadDeadline(time.Now().Add(time.Duration(s.config.GetSocket().PongWaitMs) * time.Millisecond))
	s.conn.SetPongHandler(func(string) error {
		s.conn.SetReadDeadline(time.Now().Add(time.Duration(s.config.GetSocket().PongWaitMs) * time.Millisecond))
		return nil
	})

	// Send an initial ping immediately.
	if !s.pingNow() {
		// If the first ping fails abort the rest of the consume sequence immediately.
		return
	}

	// Start a routine to process outbound messages.
	go s.processOutgoing()

	for {
		_, data, err := s.conn.ReadMessage()
		if err != nil {
			// Ignore "normal" WebSocket errors.
			if !websocket.IsCloseError(err, websocket.CloseNormalClosure, websocket.CloseGoingAway, websocket.CloseNoStatusReceived) {
				// Ignore underlying connection being shut down while read is waiting for data.
				if e, ok := err.(*net.OpError); !ok || e.Err.Error() != "use of closed network connection" {
					s.logger.Debug("Error reading message from client", zap.Error(err))
				}
			}
			break
		}

		request := &rtapi.Envelope{}
		if err = s.jsonpbUnmarshaler.Unmarshal(bytes.NewReader(data), request); err != nil {
			// If the payload is malformed the client is incompatible or misbehaving, either way disconnect it now.
			s.logger.Warn("Received malformed payload", zap.String("data", string(data)))
			break
		} else {
			// TODO Add session-global context here to cancel in-progress operations when the session is closed.
			requestLogger := s.logger.With(zap.String("cid", request.Cid))
			if !processRequest(requestLogger, s, request) {
				break
			}
		}
	}
}

func (s *sessionWS) processOutgoing() {
	for {
		select {
		case <-s.outgoingStopCh:
			// Session is closing, close the outgoing process routine.
			return
		case <-s.pingTicker.C:
			// Periodically send pings.
			if !s.pingNow() {
				// If ping fails the session will be stopped, clean up the loop.
				return
			}
		case payload := <-s.outgoingCh:
			s.Lock()
			if s.stopped {
				// The connection may have stopped between the payload being queued on the outgoing channel and reaching here.
				// If that's the case then abort outgoing processing at this point and exit.
				s.Unlock()
				return
			}
			// Process the outgoing message queue.
			s.conn.SetWriteDeadline(time.Now().Add(time.Duration(s.config.GetSocket().WriteWaitMs) * time.Millisecond))
			if err := s.conn.WriteMessage(websocket.TextMessage, payload); err != nil {
				s.Unlock()
				s.logger.Warn("Could not write message", zap.Error(err))
				return
			}
			s.Unlock()
		}
	}
}

func (s *sessionWS) pingNow() bool {
	s.Lock()
	if s.stopped {
		s.Unlock()
		return false
	}
	s.conn.SetWriteDeadline(time.Now().Add(time.Duration(s.config.GetSocket().WriteWaitMs) * time.Millisecond))
	err := s.conn.WriteMessage(websocket.PingMessage, []byte{})
	s.Unlock()
	if err != nil {
		s.logger.Warn("Could not send ping, closing channel", zap.String("remoteAddress", s.conn.RemoteAddr().String()), zap.Error(err))
		// The connection has already failed.
		s.cleanupClosedConnection()
		return false
	}
	return true
}

func (s *sessionWS) Format() SessionFormat {
	return SessionFormatJson
}

func (s *sessionWS) Send(isStream bool, mode uint8, envelope *rtapi.Envelope) error {
	payload, err := s.jsonpbMarshaler.MarshalToString(envelope)
	if err != nil {
		s.logger.Warn("Could not marshal to json", zap.Error(err))
		return err
	}

	if s.logger.Core().Enabled(zap.DebugLevel) {
		switch envelope.Message.(type) {
		case *rtapi.Envelope_Error:
			s.logger.Debug("Sending error message", zap.String("payload", payload))
		default:
			s.logger.Debug(fmt.Sprintf("Sending %T message", envelope.Message), zap.String("payload", payload))
		}
	}

	return s.SendBytes(isStream, mode, []byte(payload))
}

func (s *sessionWS) SendBytes(isStream bool, mode uint8, payload []byte) error {
	s.Lock()
	if s.stopped {
		s.Unlock()
		return nil
	}

	if isStream {
		switch mode {
		case StreamModeChannel:
			fallthrough
		case StreamModeGroup:
			fallthrough
		case StreamModeDM:
			// Chat messages are only allowed if the current outgoing queue size is below the threshold to
			// ensure there is always room to queue higher priority messages.
			if len(s.outgoingCh) < s.queuePriorityThreshold {
				s.outgoingCh <- payload
			}
			s.Unlock()
			return nil
		}
	}

	// By default attempt to queue messages and observe failures.
	select {
	case s.outgoingCh <- payload:
		s.Unlock()
		return nil
	default:
		// The outgoing queue is full, likely because the remote client can't keep up.
		// Terminate the connection immediately because the only alternative that doesn't block the server is
		// to start dropping messages, which might cause unexpected behaviour.
		s.Unlock()
		s.logger.Warn("Could not write message, session outgoing queue full")
		s.cleanupClosedConnection()
		return ErrSessionQueueFull
	}
}

func (s *sessionWS) cleanupClosedConnection() {
	s.Lock()
	if s.stopped {
		s.Unlock()
		return
	}
	s.stopped = true
	s.Unlock()

	if s.logger.Core().Enabled(zap.DebugLevel) {
		s.logger.Info("Cleaning up closed client connection", zap.String("remoteAddress", s.conn.RemoteAddr().String()))
	}

	// When connection close originates internally in the session, ensure cleanup of external resources and references.
	s.sessionRegistry.remove(s.id)
	s.matchmaker.RemoveAll(s.id)
	s.tracker.UntrackAll(s.id)

	// Clean up internals.
	s.pingTicker.Stop()
	close(s.outgoingStopCh)
	close(s.outgoingCh)

	// Close WebSocket.
	s.conn.Close()
	s.logger.Info("Closed client connection")
}

func (s *sessionWS) Close() {
	s.Lock()
	if s.stopped {
		s.Unlock()
		return
	}
	s.stopped = true
	s.Unlock()

	// Expect the caller of this session.Close() to clean up external resources (like presences) separately.

	// Clean up internals.
	s.pingTicker.Stop()
	close(s.outgoingStopCh)
	close(s.outgoingCh)

	// Send close message.
	err := s.conn.WriteControl(websocket.CloseMessage, []byte{}, time.Now().Add(time.Duration(s.config.GetSocket().WriteWaitMs)*time.Millisecond))
	if err != nil {
		s.logger.Warn("Could not send close message, closing prematurely", zap.String("remoteAddress", s.conn.RemoteAddr().String()), zap.Error(err))
	}

	// Close WebSocket.
	s.conn.Close()
	s.logger.Info("Closed client connection")
}
<|endoftext|>


<commit_before>package pin

import (
	"bytes"
	"crypto/rand"
	"encoding/binary"
	"errors"
	"fmt"
	"hash/fnv"
	"sort"
	"unsafe"

	"context"
	"github.com/ipfs/go-ipfs/merkledag"
	"github.com/ipfs/go-ipfs/pin/internal/pb"
	"gx/ipfs/QmYEoKZXHoAToWfhGF3vryhMn3WWhE1o2MasQ8uzY5iDi9/go-key"
	"gx/ipfs/QmZ4Qi3GaRbjcx28Sme5eMH7RQjGkt8wHxt2a65oLaeFEV/gogo-protobuf/proto"
	cid "gx/ipfs/QmakyCk6Vnn16WEKjbkxieZmM2YLTzkFWizbmGowoYPjro/go-cid"
)

const (
	defaultFanout = 256
	maxItems      = 8192
)

func randomSeed() (uint32, error) {
	var buf [4]byte
	if _, err := rand.Read(buf[:]); err != nil {
		return 0, err
	}
	return binary.LittleEndian.Uint32(buf[:]), nil
}

func hash(seed uint32, c *cid.Cid) uint32 {
	var buf [4]byte
	binary.LittleEndian.PutUint32(buf[:], seed)
	h := fnv.New32a()
	_, _ = h.Write(buf[:])
	_, _ = h.Write(c.Bytes())
	return h.Sum32()
}

type itemIterator func() (c *cid.Cid, data []byte, ok bool)

type keyObserver func(*cid.Cid)

// refcount is the marshaled format of refcounts. It may change
// between versions; this is valid for version 1. Changing it may
// become desirable if there are many links with refcount > 255.
//
// There are two guarantees that need to be preserved, if this is
// changed:
//
//     - the marshaled format is of fixed size, matching
//       unsafe.Sizeof(refcount(0))
//     - methods of refcount handle endianness, and may
//       in later versions need encoding/binary.
type refcount uint8

func (r refcount) Bytes() []byte {
	return []byte{byte(r)}
}

// readRefcount returns the idx'th refcount in []byte, which is
// assumed to be a sequence of refcount.Bytes results.
func (r *refcount) ReadFromIdx(buf []byte, idx int) {
	*r = refcount(buf[idx])
}

type sortByHash struct {
	links []*merkledag.Link
	data  []byte
}

func (s sortByHash) Len() int {
	return len(s.links)
}

func (s sortByHash) Less(a, b int) bool {
	return bytes.Compare(s.links[a].Hash, s.links[b].Hash) == -1
}

func (s sortByHash) Swap(a, b int) {
	s.links[a], s.links[b] = s.links[b], s.links[a]
	if len(s.data) != 0 {
		const n = int(unsafe.Sizeof(refcount(0)))
		tmp := make([]byte, n)
		copy(tmp, s.data[a*n:a*n+n])
		copy(s.data[a*n:a*n+n], s.data[b*n:b*n+n])
		copy(s.data[b*n:b*n+n], tmp)
	}
}

func storeItems(ctx context.Context, dag merkledag.DAGService, estimatedLen uint64, iter itemIterator, internalKeys keyObserver) (*merkledag.Node, error) {
	seed, err := randomSeed()
	if err != nil {
		return nil, err
	}
	n := &merkledag.Node{
		Links: make([]*merkledag.Link, 0, defaultFanout+maxItems),
	}
	for i := 0; i < defaultFanout; i++ {
		n.Links = append(n.Links, &merkledag.Link{Hash: emptyKey.Hash()})
	}
	internalKeys(emptyKey)
	hdr := &pb.Set{
		Version: proto.Uint32(1),
		Fanout:  proto.Uint32(defaultFanout),
		Seed:    proto.Uint32(seed),
	}
	if err := writeHdr(n, hdr); err != nil {
		return nil, err
	}
	hdrLen := len(n.Data())

	if estimatedLen < maxItems {
		// it'll probably fit
		for i := 0; i < maxItems; i++ {
			k, data, ok := iter()
			if !ok {
				// all done
				break
			}
			n.Links = append(n.Links, &merkledag.Link{Hash: k.Hash()})
			n.SetData(append(n.Data(), data...))
		}
		// sort by hash, also swap item Data
		s := sortByHash{
			links: n.Links[defaultFanout:],
			data:  n.Data()[hdrLen:],
		}
		sort.Stable(s)
	}

	// wasteful but simple
	type item struct {
		c    *cid.Cid
		data []byte
	}
	hashed := make(map[uint32][]item)
	for {
		k, data, ok := iter()
		if !ok {
			break
		}
		h := hash(seed, k) % defaultFanout
		hashed[h] = append(hashed[h], item{k, data})
	}
	for h, items := range hashed {
		childIter := func() (c *cid.Cid, data []byte, ok bool) {
			if len(items) == 0 {
				return nil, nil, false
			}
			first := items[0]
			items = items[1:]
			return first.c, first.data, true
		}
		child, err := storeItems(ctx, dag, uint64(len(items)), childIter, internalKeys)
		if err != nil {
			return nil, err
		}
		size, err := child.Size()
		if err != nil {
			return nil, err
		}
		childKey, err := dag.Add(child)
		if err != nil {
			return nil, err
		}
		internalKeys(childKey)
		l := &merkledag.Link{
			Name: "",
			Hash: childKey.Hash(),
			Size: size,
		}
		n.Links[int(h%defaultFanout)] = l
	}
	return n, nil
}

func readHdr(n *merkledag.Node) (*pb.Set, []byte, error) {
	hdrLenRaw, consumed := binary.Uvarint(n.Data())
	if consumed <= 0 {
		return nil, nil, errors.New("invalid Set header length")
	}
	buf := n.Data()[consumed:]
	if hdrLenRaw > uint64(len(buf)) {
		return nil, nil, errors.New("impossibly large Set header length")
	}
	// as hdrLenRaw was <= an int, we now know it fits in an int
	hdrLen := int(hdrLenRaw)
	var hdr pb.Set
	if err := proto.Unmarshal(buf[:hdrLen], &hdr); err != nil {
		return nil, nil, err
	}
	buf = buf[hdrLen:]

	if v := hdr.GetVersion(); v != 1 {
		return nil, nil, fmt.Errorf("unsupported Set version: %d", v)
	}
	if uint64(hdr.GetFanout()) > uint64(len(n.Links)) {
		return nil, nil, errors.New("impossibly large Fanout")
	}
	return &hdr, buf, nil
}

func writeHdr(n *merkledag.Node, hdr *pb.Set) error {
	hdrData, err := proto.Marshal(hdr)
	if err != nil {
		return err
	}
	n.SetData(make([]byte, binary.MaxVarintLen64, binary.MaxVarintLen64+len(hdrData)))
	written := binary.PutUvarint(n.Data(), uint64(len(hdrData)))
	n.SetData(n.Data()[:written])
	n.SetData(append(n.Data(), hdrData...))
	return nil
}

type walkerFunc func(buf []byte, idx int, link *merkledag.Link) error

func walkItems(ctx context.Context, dag merkledag.DAGService, n *merkledag.Node, fn walkerFunc, children keyObserver) error {
	hdr, buf, err := readHdr(n)
	if err != nil {
		return err
	}
	// readHdr guarantees fanout is a safe value
	fanout := hdr.GetFanout()
	for i, l := range n.Links[fanout:] {
		if err := fn(buf, i, l); err != nil {
			return err
		}
	}
	for _, l := range n.Links[:fanout] {
		c := cid.NewCidV0(l.Hash)
		children(c)
		if c.Equals(emptyKey) {
			continue
		}
		subtree, err := l.GetNode(ctx, dag)
		if err != nil {
			return err
		}
		if err := walkItems(ctx, dag, subtree, fn, children); err != nil {
			return err
		}
	}
	return nil
}

func loadSet(ctx context.Context, dag merkledag.DAGService, root *merkledag.Node, name string, internalKeys keyObserver) ([]*cid.Cid, error) {
	l, err := root.GetNodeLink(name)
	if err != nil {
		return nil, err
	}

	lnkc := cid.NewCidV0(l.Hash)
	internalKeys(lnkc)

	n, err := l.GetNode(ctx, dag)
	if err != nil {
		return nil, err
	}

	var res []*cid.Cid
	walk := func(buf []byte, idx int, link *merkledag.Link) error {
		res = append(res, cid.NewCidV0(link.Hash))
		return nil
	}
	if err := walkItems(ctx, dag, n, walk, internalKeys); err != nil {
		return nil, err
	}
	return res, nil
}

func loadMultiset(ctx context.Context, dag merkledag.DAGService, root *merkledag.Node, name string, internalKeys keyObserver) (map[key.Key]uint64, error) {
	l, err := root.GetNodeLink(name)
	if err != nil {
		return nil, fmt.Errorf("Failed to get link %s: %v", name, err)
	}
	c := cid.NewCidV0(l.Hash)
	internalKeys(c)
	n, err := l.GetNode(ctx, dag)
	if err != nil {
		return nil, fmt.Errorf("Failed to get node from link %s: %v", name, err)
	}

	refcounts := make(map[key.Key]uint64)
	walk := func(buf []byte, idx int, link *merkledag.Link) error {
		var r refcount
		r.ReadFromIdx(buf, idx)
		refcounts[key.Key(link.Hash)] += uint64(r)
		return nil
	}
	if err := walkItems(ctx, dag, n, walk, internalKeys); err != nil {
		return nil, err
	}
	return refcounts, nil
}

func storeSet(ctx context.Context, dag merkledag.DAGService, cids []*cid.Cid, internalKeys keyObserver) (*merkledag.Node, error) {
	iter := func() (c *cid.Cid, data []byte, ok bool) {
		if len(cids) == 0 {
			return nil, nil, false
		}
		first := cids[0]
		cids = cids[1:]
		return first, nil, true
	}
	n, err := storeItems(ctx, dag, uint64(len(cids)), iter, internalKeys)
	if err != nil {
		return nil, err
	}
	c, err := dag.Add(n)
	if err != nil {
		return nil, err
	}
	internalKeys(c)
	return n, nil
}

func copyRefcounts(orig map[key.Key]uint64) map[key.Key]uint64 {
	r := make(map[key.Key]uint64, len(orig))
	for k, v := range orig {
		r[k] = v
	}
	return r
}
<commit_msg>add comment detailing the algorithm and fix<commit_after>package pin

import (
	"bytes"
	"crypto/rand"
	"encoding/binary"
	"errors"
	"fmt"
	"hash/fnv"
	"sort"
	"unsafe"

	"context"
	"github.com/ipfs/go-ipfs/merkledag"
	"github.com/ipfs/go-ipfs/pin/internal/pb"
	"gx/ipfs/QmYEoKZXHoAToWfhGF3vryhMn3WWhE1o2MasQ8uzY5iDi9/go-key"
	"gx/ipfs/QmZ4Qi3GaRbjcx28Sme5eMH7RQjGkt8wHxt2a65oLaeFEV/gogo-protobuf/proto"
	cid "gx/ipfs/QmakyCk6Vnn16WEKjbkxieZmM2YLTzkFWizbmGowoYPjro/go-cid"
)

const (
	defaultFanout = 256
	maxItems      = 8192
)

func randomSeed() (uint32, error) {
	var buf [4]byte
	if _, err := rand.Read(buf[:]); err != nil {
		return 0, err
	}
	return binary.LittleEndian.Uint32(buf[:]), nil
}

func hash(seed uint32, c *cid.Cid) uint32 {
	var buf [4]byte
	binary.LittleEndian.PutUint32(buf[:], seed)
	h := fnv.New32a()
	_, _ = h.Write(buf[:])
	_, _ = h.Write(c.Bytes())
	return h.Sum32()
}

type itemIterator func() (c *cid.Cid, data []byte, ok bool)

type keyObserver func(*cid.Cid)

// refcount is the marshaled format of refcounts. It may change
// between versions; this is valid for version 1. Changing it may
// become desirable if there are many links with refcount > 255.
//
// There are two guarantees that need to be preserved, if this is
// changed:
//
//     - the marshaled format is of fixed size, matching
//       unsafe.Sizeof(refcount(0))
//     - methods of refcount handle endianness, and may
//       in later versions need encoding/binary.
type refcount uint8

func (r refcount) Bytes() []byte {
	return []byte{byte(r)}
}

// readRefcount returns the idx'th refcount in []byte, which is
// assumed to be a sequence of refcount.Bytes results.
func (r *refcount) ReadFromIdx(buf []byte, idx int) {
	*r = refcount(buf[idx])
}

type sortByHash struct {
	links []*merkledag.Link
	data  []byte
}

func (s sortByHash) Len() int {
	return len(s.links)
}

func (s sortByHash) Less(a, b int) bool {
	return bytes.Compare(s.links[a].Hash, s.links[b].Hash) == -1
}

func (s sortByHash) Swap(a, b int) {
	s.links[a], s.links[b] = s.links[b], s.links[a]
	if len(s.data) != 0 {
		const n = int(unsafe.Sizeof(refcount(0)))
		tmp := make([]byte, n)
		copy(tmp, s.data[a*n:a*n+n])
		copy(s.data[a*n:a*n+n], s.data[b*n:b*n+n])
		copy(s.data[b*n:b*n+n], tmp)
	}
}

func storeItems(ctx context.Context, dag merkledag.DAGService, estimatedLen uint64, iter itemIterator, internalKeys keyObserver) (*merkledag.Node, error) {
	seed, err := randomSeed()
	if err != nil {
		return nil, err
	}
	n := &merkledag.Node{
		Links: make([]*merkledag.Link, 0, defaultFanout+maxItems),
	}
	for i := 0; i < defaultFanout; i++ {
		n.Links = append(n.Links, &merkledag.Link{Hash: emptyKey.Hash()})
	}
	internalKeys(emptyKey)
	hdr := &pb.Set{
		Version: proto.Uint32(1),
		Fanout:  proto.Uint32(defaultFanout),
		Seed:    proto.Uint32(seed),
	}
	if err := writeHdr(n, hdr); err != nil {
		return nil, err
	}
	hdrLen := len(n.Data())

	if estimatedLen < maxItems {
		// it'll probably fit
		for i := 0; i < maxItems; i++ {
			k, data, ok := iter()
			if !ok {
				// all done
				break
			}
			n.Links = append(n.Links, &merkledag.Link{Hash: k.Hash()})
			n.SetData(append(n.Data(), data...))
		}
		// sort by hash, also swap item Data
		s := sortByHash{
			links: n.Links[defaultFanout:],
			data:  n.Data()[hdrLen:],
		}
		sort.Stable(s)
	}

	// wasteful but simple
	type item struct {
		c    *cid.Cid
		data []byte
	}
	hashed := make(map[uint32][]item)
	for {
		// This loop essentially enumerates every single item in the set
		// and maps them all into a set of buckets. Each bucket will be recursively
		// turned into its own sub-set, and so on down the chain. Each sub-set
		// gets added to the dagservice, and put into its place in a set nodes
		// links array.
		//
		// Previously, the bucket was selected by taking an int32 from the hash of
		// the input key + seed. This was erroneous as we would later be assigning
		// the created sub-sets into an array of length 256 by the modulus of the
		// int32 hash value with 256. This resulted in overwriting existing sub-sets
		// and losing pins. The fix (a few lines down from this comment), is to
		// map the hash value down to the 8 bit keyspace here while creating the
		// buckets. This way, we avoid any overlapping later on.
		k, data, ok := iter()
		if !ok {
			break
		}
		h := hash(seed, k) % defaultFanout
		hashed[h] = append(hashed[h], item{k, data})
	}
	for h, items := range hashed {
		childIter := func() (c *cid.Cid, data []byte, ok bool) {
			if len(items) == 0 {
				return nil, nil, false
			}
			first := items[0]
			items = items[1:]
			return first.c, first.data, true
		}
		child, err := storeItems(ctx, dag, uint64(len(items)), childIter, internalKeys)
		if err != nil {
			return nil, err
		}
		size, err := child.Size()
		if err != nil {
			return nil, err
		}
		childKey, err := dag.Add(child)
		if err != nil {
			return nil, err
		}
		internalKeys(childKey)
		l := &merkledag.Link{
			Name: "",
			Hash: childKey.Hash(),
			Size: size,
		}
		n.Links[int(h%defaultFanout)] = l
	}
	return n, nil
}

func readHdr(n *merkledag.Node) (*pb.Set, []byte, error) {
	hdrLenRaw, consumed := binary.Uvarint(n.Data())
	if consumed <= 0 {
		return nil, nil, errors.New("invalid Set header length")
	}
	buf := n.Data()[consumed:]
	if hdrLenRaw > uint64(len(buf)) {
		return nil, nil, errors.New("impossibly large Set header length")
	}
	// as hdrLenRaw was <= an int, we now know it fits in an int
	hdrLen := int(hdrLenRaw)
	var hdr pb.Set
	if err := proto.Unmarshal(buf[:hdrLen], &hdr); err != nil {
		return nil, nil, err
	}
	buf = buf[hdrLen:]

	if v := hdr.GetVersion(); v != 1 {
		return nil, nil, fmt.Errorf("unsupported Set version: %d", v)
	}
	if uint64(hdr.GetFanout()) > uint64(len(n.Links)) {
		return nil, nil, errors.New("impossibly large Fanout")
	}
	return &hdr, buf, nil
}

func writeHdr(n *merkledag.Node, hdr *pb.Set) error {
	hdrData, err := proto.Marshal(hdr)
	if err != nil {
		return err
	}
	n.SetData(make([]byte, binary.MaxVarintLen64, binary.MaxVarintLen64+len(hdrData)))
	written := binary.PutUvarint(n.Data(), uint64(len(hdrData)))
	n.SetData(n.Data()[:written])
	n.SetData(append(n.Data(), hdrData...))
	return nil
}

type walkerFunc func(buf []byte, idx int, link *merkledag.Link) error

func walkItems(ctx context.Context, dag merkledag.DAGService, n *merkledag.Node, fn walkerFunc, children keyObserver) error {
	hdr, buf, err := readHdr(n)
	if err != nil {
		return err
	}
	// readHdr guarantees fanout is a safe value
	fanout := hdr.GetFanout()
	for i, l := range n.Links[fanout:] {
		if err := fn(buf, i, l); err != nil {
			return err
		}
	}
	for _, l := range n.Links[:fanout] {
		c := cid.NewCidV0(l.Hash)
		children(c)
		if c.Equals(emptyKey) {
			continue
		}
		subtree, err := l.GetNode(ctx, dag)
		if err != nil {
			return err
		}
		if err := walkItems(ctx, dag, subtree, fn, children); err != nil {
			return err
		}
	}
	return nil
}

func loadSet(ctx context.Context, dag merkledag.DAGService, root *merkledag.Node, name string, internalKeys keyObserver) ([]*cid.Cid, error) {
	l, err := root.GetNodeLink(name)
	if err != nil {
		return nil, err
	}

	lnkc := cid.NewCidV0(l.Hash)
	internalKeys(lnkc)

	n, err := l.GetNode(ctx, dag)
	if err != nil {
		return nil, err
	}

	var res []*cid.Cid
	walk := func(buf []byte, idx int, link *merkledag.Link) error {
		res = append(res, cid.NewCidV0(link.Hash))
		return nil
	}
	if err := walkItems(ctx, dag, n, walk, internalKeys); err != nil {
		return nil, err
	}
	return res, nil
}

func loadMultiset(ctx context.Context, dag merkledag.DAGService, root *merkledag.Node, name string, internalKeys keyObserver) (map[key.Key]uint64, error) {
	l, err := root.GetNodeLink(name)
	if err != nil {
		return nil, fmt.Errorf("Failed to get link %s: %v", name, err)
	}
	c := cid.NewCidV0(l.Hash)
	internalKeys(c)
	n, err := l.GetNode(ctx, dag)
	if err != nil {
		return nil, fmt.Errorf("Failed to get node from link %s: %v", name, err)
	}

	refcounts := make(map[key.Key]uint64)
	walk := func(buf []byte, idx int, link *merkledag.Link) error {
		var r refcount
		r.ReadFromIdx(buf, idx)
		refcounts[key.Key(link.Hash)] += uint64(r)
		return nil
	}
	if err := walkItems(ctx, dag, n, walk, internalKeys); err != nil {
		return nil, err
	}
	return refcounts, nil
}

func storeSet(ctx context.Context, dag merkledag.DAGService, cids []*cid.Cid, internalKeys keyObserver) (*merkledag.Node, error) {
	iter := func() (c *cid.Cid, data []byte, ok bool) {
		if len(cids) == 0 {
			return nil, nil, false
		}
		first := cids[0]
		cids = cids[1:]
		return first, nil, true
	}
	n, err := storeItems(ctx, dag, uint64(len(cids)), iter, internalKeys)
	if err != nil {
		return nil, err
	}
	c, err := dag.Add(n)
	if err != nil {
		return nil, err
	}
	internalKeys(c)
	return n, nil
}

func copyRefcounts(orig map[key.Key]uint64) map[key.Key]uint64 {
	r := make(map[key.Key]uint64, len(orig))
	for k, v := range orig {
		r[k] = v
	}
	return r
}
<|endoftext|>


<commit_before>package main

import (
	"bufio"
	"fmt"
	"log"
	"os"
	"sync"

	"github.com/andrew-d/go-termutil"
	"github.com/codegangsta/cli"
	"github.com/streadway/amqp"
)

func failOnError(err error, msg string) {
	if err != nil {
		log.Fatalf("%s: %s", msg, err)
		panic(fmt.Sprintf("%s: %s", msg, err))
	}
}

func main() {
	app := cli.NewApp()
	app.Name = "pipecat"
	app.Usage = "Connect unix pipes and message queues"
	app.Flags = []cli.Flag{
		cli.StringFlag{
			Name:   "amqpUri",
			Value:  "amqp://guest:guest@localhost:5672/",
			Usage:  "AMQP URI",
			EnvVar: "AMQP_URI",
		},
	}

	app.Action = func(c *cli.Context) {
		list := c.Args().First()
		if list == "" {
			fmt.Println("Please provide name of the queue")
			os.Exit(1)
		}

		conn, err := amqp.Dial(c.String("amqpUri"))
		failOnError(err, "Failed to connect to AMPQ broker")
		defer conn.Close()

		channel, err := conn.Channel()
		failOnError(err, "Failed to open a channel")
		defer channel.Close()

		q, err := channel.QueueDeclare(
			list,  // name
			true,  // durable
			false, // delete when unused
			false, // exclusive
			false, // no-wait
			nil,   // arguments
		)
		failOnError(err, "Failed to declare a queue")

		var mutex sync.Mutex
		unackedMessages := make([]amqp.Delivery, 100)
		if termutil.Isatty(os.Stdin.Fd()) {
			msgs, err := channel.Consume(
				q.Name, // queue
				"",     // consumer
				false,  // auto-ack
				false,  // exclusive
				false,  // no-local
				false,  // no-wait
				nil,    // args
			)
			failOnError(err, "Failed to register a consumer")

			go func() {
				for msg := range msgs {
					mutex.Lock()
					unackedMessages = append(unackedMessages, msg)
					mutex.Unlock()

					line := fmt.Sprintf("%s", msg.Body)
					fmt.Println(line)
					fmt.Fprintln(os.Stderr, fmt.Sprintf("%d", len(unackedMessages)))
				}
			}()
		} else {
			scanner := bufio.NewScanner(os.Stdin)
			for scanner.Scan() {
				line := scanner.Text()
				err = channel.Publish(
					"",     // exchange
					q.Name, // routing key
					false,  // mandatory
					false,  // immediate
					amqp.Publishing{
						ContentType: "text/plain",
						Body:        []byte(line),
					})
				fmt.Println(line)
				failOnError(err, "Failed to publish a message")
			}
			if err := scanner.Err(); err != nil {
				fmt.Fprintln(os.Stderr, "Reading standard input:", err)
			}
		}
	}

	app.Run(os.Args)
}
<commit_msg>Works with acking<commit_after>package main

import (
	"bufio"
	"fmt"
	"log"
	"os"
	"sync"
	"time"

	"github.com/andrew-d/go-termutil"
	"github.com/codegangsta/cli"
	"github.com/streadway/amqp"
)

func failOnError(err error, msg string) {
	if err != nil {
		log.Fatalf("%s: %s", msg, err)
		panic(fmt.Sprintf("%s: %s", msg, err))
	}
}

func main() {
	app := cli.NewApp()
	app.Name = "pipecat"
	app.Usage = "Connect unix pipes and message queues"
	app.Flags = []cli.Flag{
		cli.StringFlag{
			Name:   "amqpUri",
			Value:  "amqp://guest:guest@localhost:5672/",
			Usage:  "AMQP URI",
			EnvVar: "AMQP_URI",
		},
		cli.StringFlag{
			Name:   "fack",
			Value:  os.DevNull,
			Usage:  "Writes the acknowledged lines to file handle",
			EnvVar: "FACK",
		},
	}

	app.Action = func(c *cli.Context) {
		queueName := c.Args().First()
		if queueName == "" {
			fmt.Println("Please provide name of the queue")
			os.Exit(1)
		}

		conn, err := amqp.Dial(c.String("amqpUri"))
		failOnError(err, "Failed to connect to AMPQ broker")
		defer conn.Close()

		channel, err := conn.Channel()
		failOnError(err, "Failed to open a channel")
		defer channel.Close()

		q, err := channel.QueueDeclare(
			queueName, // name
			true,      // durable
			false,     // delete when unused
			false,     // exclusive
			false,     // no-wait
			nil,       // arguments
		)
		failOnError(err, "Failed to declare a queue")

		var mutex sync.Mutex
		unackedMessages := make([]amqp.Delivery, 100)

		msgs, err := channel.Consume(
			q.Name, // queue
			"",     // consumer
			false,  // auto-ack
			false,  // exclusive
			false,  // no-local
			false,  // no-wait
			nil,    // args
		)
		failOnError(err, "Failed to register a consumer")

		stdack, err := os.OpenFile(c.String("fack"), os.O_APPEND, 0660)
		failOnError(err, "Could not open ack file")
		defer stdack.Close()

		if termutil.Isatty(os.Stdin.Fd()) {
			// only read
			forever := make(chan bool)

			go func() {

				scanner := bufio.NewScanner(os.Stdin)
				for scanner.Scan() {
					ackedLine := scanner.Text()

					mutex.Lock()
					for _, msg := range unackedMessages {
						unackedLine := fmt.Sprintf("%s", msg.Body)
						if unackedLine == ackedLine {
							fmt.Fprintln(stdack, ackedLine)
							msg.Ack(false)
						}
					}
					mutex.Unlock()
				}
				if err := scanner.Err(); err != nil {
					fmt.Fprintln(os.Stderr, "Reading standard input:", err)
				}

			}()

			go func() {
				for msg := range msgs {
					mutex.Lock()
					unackedMessages = append(unackedMessages, msg)
					mutex.Unlock()

					time.Sleep(100 * time.Millisecond)
					line := fmt.Sprintf("%s", msg.Body)
					fmt.Println(line)
				}
			}()
			<-forever
		} else {

			scanner := bufio.NewScanner(os.Stdin)
			for scanner.Scan() {
				line := scanner.Text()
				err = channel.Publish(
					"",     // exchange
					q.Name, // routing key
					false,  // mandatory
					false,  // immediate
					amqp.Publishing{
						ContentType: "text/plain",
						Body:        []byte(line),
					})

				failOnError(err, "Failed to publish a message")
				fmt.Println(line)
				fmt.Fprintln(stdack, line)
			}
			if err := scanner.Err(); err != nil {
				fmt.Fprintln(os.Stderr, "Reading standard input:", err)
			}

		}

		/*} else {
			scanner := bufio.NewScanner(os.Stdin)
			for scanner.Scan() {
				line := scanner.Text()
				err = channel.Publish(
					"",     // exchange
					q.Name, // routing key
					false,  // mandatory
					false,  // immediate
					amqp.Publishing{
						ContentType: "text/plain",
						Body:        []byte(line),
					})
				fmt.Println(line)
				failOnError(err, "Failed to publish a message")
			}
			if err := scanner.Err(); err != nil {
				fmt.Fprintln(os.Stderr, "Reading standard input:", err)
			}
		}*/
	}

	app.Run(os.Args)
}
<|endoftext|>


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

package common

import (
	"fmt"

	"github.com/juju/names"

	"github.com/juju/juju/environs"
)

// EnvFullName returns a string based on the provided environment
// that is suitable for identifying the env on a provider. The resuling
// string clearly associates the value with juju, whereas the
// environment's UUID alone isn't very distinctive for humans. This
// benefits users by helping them quickly identify in their hosting
// management tools which instances are juju related.
func EnvFullName(env environs.Environ) string {
	envUUID, _ := env.Config().UUID() // Env should have validated this.
	return fmt.Sprintf("juju-%s", envUUID)
}

// MachineFullName returns a string based on the provided environment
// and machine ID that is suitable for identifying instances on a
// provider. See EnvFullName for an explanation on how this function
// helps juju users.
func MachineFullName(env environs.Environ, machineId string) string {
	envstr := EnvFullName(env)
	machineTag := names.NewMachineTag(machineId)
	return fmt.Sprintf("%s-%s", envstr, machineTag)
}
<commit_msg>provider/gce: Do not use names.NewMachineTag in MachineFullName.<commit_after>// Copyright 2014 Canonical Ltd.
// Licensed under the AGPLv3, see LICENCE file for details.

package common

import (
	"fmt"

	"github.com/juju/juju/environs"
)

// EnvFullName returns a string based on the provided environment
// that is suitable for identifying the env on a provider. The resuling
// string clearly associates the value with juju, whereas the
// environment's UUID alone isn't very distinctive for humans. This
// benefits users by helping them quickly identify in their hosting
// management tools which instances are juju related.
func EnvFullName(env environs.Environ) string {
	envUUID, _ := env.Config().UUID() // Env should have validated this.
	return fmt.Sprintf("juju-%s", envUUID)
}

// MachineFullName returns a string based on the provided environment
// and machine ID that is suitable for identifying instances on a
// provider. See EnvFullName for an explanation on how this function
// helps juju users.
func MachineFullName(env environs.Environ, machineID string) string {
	envstr := EnvFullName(env)
	return fmt.Sprintf("%s-machine-%s", envstr, machineID)
}
<|endoftext|>


<commit_before>package main

import (
  "flag"
  "fmt"
  "io"
  "log"
  "os"
  "path/filepath"
  "sort"
  "time"
)

var (
  output_name string
)

func init() {
  flag.StringVar(&output_name, "output", output_name, "Output filename (- for STDOUT, do not specify for default, multifasta.txt)")
  flag.Parse()
}

func defaultOutput(inputPath string) string {
  dir := filepath.Dir(inputPath)
  t := time.Now()
  outfile := fmt.Sprintf("multifasta_output_%d%02d%02d-%02d%02d.txt",
    t.Year(),
    t.Month(),
    t.Day(),
    t.Hour(),
    t.Minute())
  return filepath.Join(dir, outfile)
}

func openOutput(outputPath string) *os.File {
  switch outputPath {
  case "-":
    return os.Stdout
  default:
    out, err := os.Create(outputPath)
    if nil != err {
      log.Fatal(err)
    }
    return out
  }
}

func main() {
  if len(flag.Args()) <= 0 {
    fmt.Fprintf(os.Stderr, "Missing one or more input file(s)\n")
    flag.Usage()
    os.Exit(1)
  }
  input_paths := flag.Args()
  sort.Strings(input_paths)
  if "" == output_name {
    output_name = defaultOutput(input_paths[0])
  }
  out := openOutput(output_name)

  defer out.Close()

  for _, infile := range input_paths {
    basefile := filepath.Base(infile)
    extension := filepath.Ext(basefile)
    fmt.Fprintf(out, ">%s exported from %s\r\n", basefile[:len(basefile)-len(extension)], basefile)
    in, err := os.Open(infile)
    if err != nil {
      log.Fatal(err)
    }
    io.Copy(out, in)
    in.Close()
    fmt.Fprintf(out, "\r\n")
  }
}
<commit_msg>Skip new line on final file<commit_after>package main

import (
  "flag"
  "fmt"
  "io"
  "log"
  "os"
  "path/filepath"
  "sort"
  "time"
)

var (
  output_name string
)

func init() {
  flag.StringVar(&output_name, "output", output_name, "Output filename (- for STDOUT, do not specify for default, multifasta.txt)")
  flag.Parse()
}

func defaultOutput(inputPath string) string {
  dir := filepath.Dir(inputPath)
  t := time.Now()
  outfile := fmt.Sprintf("multifasta_output_%d%02d%02d-%02d%02d.txt",
    t.Year(),
    t.Month(),
    t.Day(),
    t.Hour(),
    t.Minute())
  return filepath.Join(dir, outfile)
}

func openOutput(outputPath string) *os.File {
  switch outputPath {
  case "-":
    return os.Stdout
  default:
    out, err := os.Create(outputPath)
    if nil != err {
      log.Fatal(err)
    }
    return out
  }
}

func main() {
  if len(flag.Args()) <= 0 {
    fmt.Fprintf(os.Stderr, "Missing one or more input file(s)\n")
    flag.Usage()
    os.Exit(1)
  }
  input_paths := flag.Args()
  sort.Strings(input_paths)
  if "" == output_name {
    output_name = defaultOutput(input_paths[0])
  }
  out := openOutput(output_name)

  defer out.Close()

  for index, infile := range input_paths {
    basefile := filepath.Base(infile)
    extension := filepath.Ext(basefile)
    fmt.Fprintf(out, ">%s exported from %s\r\n", basefile[:len(basefile)-len(extension)], basefile)
    in, err := os.Open(infile)
    if err != nil {
      log.Fatal(err)
    }
    io.Copy(out, in)
    in.Close()
    if index+1 < len(input_paths) {
      fmt.Fprintf(out, "\r\n")
    }
  }
}
<|endoftext|>