<commit_before>// Copyright 2013 Unknown
//
// 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 models

import (
	"errors"
	"strconv"
	"strings"
	"time"

	"github.com/Unknwon/gowalker/utils"
	"github.com/astaxie/beego"
	"github.com/coocood/qbs"
)

/*
	GetPopulars returns <num>
		1. Recent viewed
		2. Top rank
		3. Top viewed
		4. Rock this week
	projects and recent updated examples.
*/
func GetPopulars(proNum, exNum int) (error, []*PkgExam, []*PkgInfo, []*PkgInfo, []*PkgInfo, []*PkgInfo) {
	q := connDb()
	defer q.Close()

	var ruExs []*PkgExam
	err := q.OmitFields("examples", "created").
		Limit(exNum).OrderByDesc("created").FindAll(&ruExs)
	if err != nil {
		return err, nil, nil, nil, nil, nil
	}

	var rvPros, trPros, tvPros, rtwPros []*PkgInfo
	var procks []*PkgRock
	err = q.OmitFields("ProName", "IsCmd", "Tags", "Views", "Created",
		"Etag", "Labels", "ImportedNum", "ImportPid", "Note").
		Limit(proNum).OrderByDesc("viewed_time").FindAll(&rvPros)
	if err != nil {
		return err, nil, nil, nil, nil, nil
	}
	err = q.OmitFields("ProName", "IsCmd", "Tags", "Views", "ViewedTime", "Created",
		"Etag", "Labels", "ImportedNum", "ImportPid", "Note").
		Limit(proNum).OrderByDesc("rank").FindAll(&trPros)
	if err != nil {
		return err, nil, nil, nil, nil, nil
	}
	err = q.OmitFields("ProName", "IsCmd", "Tags", "ViewedTime", "Created",
		"Etag", "Labels", "ImportedNum", "ImportPid", "Note").
		Limit(proNum).OrderByDesc("views").FindAll(&tvPros)
	if err != nil {
		return err, nil, nil, nil, nil, nil
	}
	err = q.OmitFields("ProName", "IsCmd", "Tags", "Views", "ViewedTime", "Created",
		"Etag", "Labels", "ImportedNum", "ImportPid", "Note").
		Limit(proNum).OrderByDesc("delta").FindAll(&procks)
	if err != nil {
		return err, nil, nil, nil, nil, nil
	}
	for _, pr := range procks {
		rtwPros = append(rtwPros, &PkgInfo{
			Id:   pr.Pid,
			Path: pr.Path,
			Rank: pr.Rank,
		})
	}
	return nil, ruExs, rvPros, trPros, tvPros, rtwPros
}

// SaveProject saves package information, declaration and functions;
// update import information.
func SaveProject(pinfo *PkgInfo, pdecl *PkgDecl, pfuncs []*PkgFunc, imports []string) error {
	q := connDb()
	defer q.Close()

	// Load package information(save after checked import information).
	info := new(PkgInfo)
	err := q.WhereEqual("path", pinfo.Path).Find(info)
	if err == nil {
		pinfo.Id = info.Id
	}

	// Save package declaration.
	decl := new(PkgDecl)
	if pdecl != nil {
		cond := qbs.NewCondition("pid = ?", pinfo.Id).And("tag = ?", pdecl.Tag)
		err = q.Condition(cond).Find(decl)
		if err == nil {
			pdecl.Id = decl.Id
		}

		pdecl.Pid = pinfo.Id
		_, err = q.Save(pdecl)
		if err != nil {
			beego.Error("models.SaveProject(", pinfo.Path, ") -> Declaration:", err)
		}
	}

	// ------------------------------
	// Save package functions.
	// ------------------------------

	if pfuncs != nil {
		// Old package need to clean old data.
		if decl.Id > 0 {
			// Update all old functions' 'IsOle' to be true.
			type pkgFunc struct {
				IsOld bool
			}
			pfunc := new(pkgFunc)
			pfunc.IsOld = true
			_, err = q.WhereEqual("pid", pdecl.Id).Update(pfunc)
		}

		isMaster := pdecl.Tag == ""
		// Save new ones.
		for _, pf := range pfuncs {
			f := new(PkgFunc)
			cond := qbs.NewCondition("pid = ?", pdecl.Id).And("name = ?", pf.Name)
			err = q.Condition(cond).Find(f)
			if err == nil {
				pf.Id = f.Id
			}

			pf.Pid = pdecl.Id
			pf.IsMaster = isMaster
			_, err = q.Save(pf)
			if err != nil {
				beego.Error("models.SaveProject(", pinfo.Path, ") -> Update function(", pf.Name, "):", err)
			}
		}

		if decl.Id > 0 {
			// Delete old ones if exist.
			cond := qbs.NewCondition("pid = ?", pdecl.Id).And("is_old = ?", true)
			_, err = q.Condition(cond).Delete(new(PkgFunc))
			if err != nil {
				beego.Error("models.SaveProject(", pinfo.Path, ") -> Delete functions:", err)
			}
		}
	}

	// ------------- END ------------

	// ------------------------------
	// Update imported information.
	// ------------------------------

	if info.Id > 0 {
		// Current package.
		importeds := strings.Split(
			strings.Replace(info.ImportPid, "$", "", -1), "|")
		importPids := make([]string, 0, len(importeds))
		for _, v := range importeds {
			pid, _ := strconv.ParseInt(v, 10, 64)
			if checkImport(q, info.Path, pid) {
				importPids = append(importPids, "$"+v)
			}
		}
		pinfo.ImportPid = strings.Join(importPids, "|")
		pinfo.ImportedNum = len(importPids)
	}

	_, err = q.Save(pinfo)
	if err != nil {
		beego.Error("models.SaveProject(", pinfo.Path, ") -> Information:", err)
	}

	// Don't need to check standard library.
	if imports != nil && !utils.IsGoRepoPath(pinfo.Path) {
		// Other packages.
		for _, v := range imports {
			if !utils.IsGoRepoPath(v) {
				// Only count non-standard library.
				updateImportInfo(q, v, int(pinfo.Id), true)
			}
		}
	}
	// ------------- END ------------

	return nil
}

// checkImport returns true if the package(id) imports given package(path).
func checkImport(q *qbs.Qbs, path string, id int64) bool {
	pinfo := &PkgInfo{
		Id: id,
	}
	err := q.Find(pinfo)
	if err != nil {
		return false
	}

	decl := new(PkgDecl)
	cond := qbs.NewCondition("pid = ?", pinfo.Id).And("tag = ?", "")
	err = q.Condition(cond).Find(decl)
	if err != nil {
		return false
	}

	if strings.Index(decl.Imports, path) == -1 {
		return false
	}

	return true
}

// LoadProject returns package declaration.
func LoadProject(pid int64, tag string) (*PkgDecl, error) {
	// Check path length to reduce connect times.
	if pid == 0 {
		return nil, errors.New("models.LoadProject -> Zero id.")
	}

	q := connDb()
	defer q.Close()

	pdecl := new(PkgDecl)
	cond := qbs.NewCondition("pid = ?", pid).And("tag = ?", tag)
	err := q.Condition(cond).Find(pdecl)
	return pdecl, err
}

// DeleteProject deletes everything about the path in database, and update import information.
func DeleteProject(path string) error {
	// Check path length to reduce connect times. (except launchpad.net)
	if path[0] != 'l' && len(strings.Split(path, "/")) <= 2 {
		beego.Error("models.DeleteProject(", path, ") -> Short path as not needed")
		return nil
	}

	q := connDb()
	defer q.Close()

	var i1, i2, i3, i4, i5 int64
	// Delete package information.
	info := new(PkgInfo)
	err := q.WhereEqual("path", path).Find(info)
	if err == nil {
		i1, err = q.WhereEqual("path", path).Delete(info)
		if err != nil {
			beego.Error("models.DeleteProject(", path, ") -> Information:", err)
		}
	}

	// Delete package declaration.
	if info.Id > 0 {
		// Find.
		var pdecls []*PkgDecl
		err = q.WhereEqual("pid", info.Id).FindAll(&pdecls)
		if err != nil {
			beego.Error("models.DeleteProject(", path, ") -> Find declaration:", err)
		}

		// Update.
		if !utils.IsGoRepoPath(path) {
			for _, pd := range pdecls {
				// Don't need to check standard library.
				// Update import information.
				imports := strings.Split(pd.Imports, "|")
				imports = imports[:len(imports)-1]
				for _, v := range imports {
					if !utils.IsGoRepoPath(v) {
						// Only count non-standard library.
						updateImportInfo(q, v, int(info.Id), false)
					}
				}
			}
		}

		// Delete.
		i2, err = q.WhereEqual("pid", info.Id).Delete(new(PkgDecl))
		if err != nil {
			beego.Error("models.DeleteProject(", path, ") -> Delete declaration:", err)
		}
	}

	// Delete package documentation.
	i3, err = q.WhereEqual("path", path).Delete(new(PkgDoc))
	if err != nil {
		beego.Error("models.DeleteProject(", path, ") -> Documentation:", err)
	}

	// Delete package examples.
	i4, err = q.WhereEqual("path", path).Delete(new(PkgExam))
	if err != nil {
		beego.Error("models.DeleteProject(", path, ") -> Examples:", err)
	}

	// Delete package functions.
	if info.Id > 0 {
		i5, err = q.WhereEqual("path", path).Delete(new(PkgExam))
		if err != nil {
			beego.Error("models.DeleteProject(", path, ") -> Functions:", err)
		}
	}

	if i1+i2+i3+i4+i5 > 0 {
		beego.Info("models.DeleteProject(", path, i1, i2, i3, i4, i5, ")")
	}

	return nil
}

func updateImportInfo(q *qbs.Qbs, path string, pid int, add bool) {
	// Save package information.
	info := new(PkgInfo)
	err := q.WhereEqual("path", path).Find(info)
	if err == nil {
		// Check if pid exists in this project.
		i := strings.Index(info.ImportPid, "$"+strconv.Itoa(pid)+"|")
		switch {
		case i == -1 && add: // Add operation and does not contain.
			info.ImportPid += "$" + strconv.Itoa(pid) + "|"
			info.ImportedNum = len(strings.Split(info.ImportPid, "|")) - 1
			_, err = q.Save(info)
			if err != nil {
				beego.Error("models.updateImportInfo -> add:", path, err)
			}
		case i > -1 && !add: // Delete operation and contains.
			info.ImportPid = strings.Replace(info.ImportPid, "$"+strconv.Itoa(pid)+"|", "", 1)
			info.ImportedNum = len(strings.Split(info.ImportPid, "|")) - 1
			_, err = q.Save(info)
			if err != nil {
				beego.Error("models.updateImportInfo -> delete:", path, err)
			}
		}
	}

	// Error means this project does not exist, simply skip.
}

// FlushCacheProjects saves cache data to database.
func FlushCacheProjects(pinfos []*PkgInfo, procks []*PkgRock) {
	q := connDb()
	defer q.Close()

	// Update project data.
	for _, p := range pinfos {
		info := new(PkgInfo)
		err := q.WhereEqual("path", p.Path).Find(info)
		if err == nil {
			// Shoule always be nil, just in case not exist.
			p.Id = info.Id
			// Limit 10 views each period.
			if p.Views-info.Views > 10 {
				p.Views = info.Views + 10
			}
		}
		_, err = q.Save(p)
		if err != nil {
			beego.Error("models.FlushCacheProjects(", p.Path, ") ->", err)
		}
	}

	// Update rock this week.
	if time.Now().Weekday() == time.Monday && utils.Cfg.MustBool("task", "rock_reset") {
		utils.Cfg.SetValue("task", "rock_reset", "0")
		utils.SaveConfig()
		// Reset rock table.
		_, err := q.Where("id > ?", int64(0)).Delete(new(PkgRock))
		if err != nil {
			beego.Error("models.FlushCacheProjects -> Reset rock table:", err)
		}
	} else if time.Now().Weekday() != time.Monday && !utils.Cfg.MustBool("task", "rock_reset") {
		utils.Cfg.SetValue("task", "rock_reset", "1")
		utils.SaveConfig()
	}

	for _, pr := range procks {
		r := new(PkgRock)
		err := q.WhereEqual("pid", pr.Pid).Find(r)
		if err == nil {
			pr.Id = r.Id
			r.Delta += pr.Rank - r.Rank
			pr.Delta = r.Delta
		}
		q.Save(pr)
	}
}
<commit_msg>Update<commit_after>// Copyright 2013 Unknown
//
// 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 models

import (
	"errors"
	"strconv"
	"strings"
	"time"

	"github.com/Unknwon/gowalker/utils"
	"github.com/astaxie/beego"
	"github.com/coocood/qbs"
)

/*
	GetPopulars returns <num>
		1. Recent viewed
		2. Top rank
		3. Top viewed
		4. Rock this week
	projects and recent updated examples.
*/
func GetPopulars(proNum, exNum int) (error, []*PkgExam, []*PkgInfo, []*PkgInfo, []*PkgInfo, []*PkgInfo) {
	q := connDb()
	defer q.Close()

	var ruExs []*PkgExam
	err := q.OmitFields("examples", "created").
		Limit(exNum).OrderByDesc("created").FindAll(&ruExs)
	if err != nil {
		return err, nil, nil, nil, nil, nil
	}

	var rvPros, trPros, tvPros, rtwPros []*PkgInfo
	var procks []*PkgRock
	err = q.OmitFields("ProName", "IsCmd", "Tags", "Views", "Created",
		"Etag", "Labels", "ImportedNum", "ImportPid", "Note").
		Limit(proNum).OrderByDesc("viewed_time").FindAll(&rvPros)
	if err != nil {
		return err, nil, nil, nil, nil, nil
	}
	err = q.OmitFields("ProName", "IsCmd", "Tags", "Views", "ViewedTime", "Created",
		"Etag", "Labels", "ImportedNum", "ImportPid", "Note").
		Limit(proNum).OrderByDesc("rank").FindAll(&trPros)
	if err != nil {
		return err, nil, nil, nil, nil, nil
	}
	err = q.OmitFields("ProName", "IsCmd", "Tags", "ViewedTime", "Created",
		"Etag", "Labels", "ImportedNum", "ImportPid", "Note").
		Limit(proNum).OrderByDesc("views").FindAll(&tvPros)
	if err != nil {
		return err, nil, nil, nil, nil, nil
	}
	err = q.OmitFields("ProName", "IsCmd", "Tags", "Views", "ViewedTime", "Created",
		"Etag", "Labels", "ImportedNum", "ImportPid", "Note").
		Limit(proNum).OrderByDesc("delta").FindAll(&procks)
	if err != nil {
		return err, nil, nil, nil, nil, nil
	}
	for _, pr := range procks {
		rtwPros = append(rtwPros, &PkgInfo{
			Id:   pr.Pid,
			Path: pr.Path,
			Rank: pr.Rank,
		})
	}
	return nil, ruExs, rvPros, trPros, tvPros, rtwPros
}

// SaveProject saves package information, declaration and functions;
// update import information.
func SaveProject(pinfo *PkgInfo, pdecl *PkgDecl, pfuncs []*PkgFunc, imports []string) error {
	q := connDb()
	defer q.Close()

	// Load package information(save after checked import information).
	info := new(PkgInfo)
	err := q.WhereEqual("path", pinfo.Path).Find(info)
	if err == nil {
		pinfo.Id = info.Id
	}

	// Save package declaration.
	decl := new(PkgDecl)
	if pdecl != nil {
		cond := qbs.NewCondition("pid = ?", pinfo.Id).And("tag = ?", pdecl.Tag)
		err = q.Condition(cond).Find(decl)
		if err == nil {
			pdecl.Id = decl.Id
		}

		pdecl.Pid = pinfo.Id
		_, err = q.Save(pdecl)
		if err != nil {
			beego.Error("models.SaveProject(", pinfo.Path, ") -> Declaration:", err)
		}
	}

	// ------------------------------
	// Save package functions.
	// ------------------------------

	if pfuncs != nil {
		// Old package need to clean old data.
		if decl.Id > 0 {
			// Update all old functions' 'IsOle' to be true.
			type pkgFunc struct {
				IsOld bool
			}
			pfunc := new(pkgFunc)
			pfunc.IsOld = true
			_, err = q.WhereEqual("pid", pdecl.Id).Update(pfunc)
		}

		isMaster := pdecl.Tag == ""
		// Save new ones.
		for _, pf := range pfuncs {
			f := new(PkgFunc)
			cond := qbs.NewCondition("pid = ?", pdecl.Id).And("name = ?", pf.Name)
			err = q.Condition(cond).Find(f)
			if err == nil {
				pf.Id = f.Id
			}

			pf.Pid = pdecl.Id
			pf.IsMaster = isMaster
			_, err = q.Save(pf)
			if err != nil {
				beego.Error("models.SaveProject(", pinfo.Path, ") -> Update function(", pf.Name, "):", err)
			}
		}

		if decl.Id > 0 {
			// Delete old ones if exist.
			cond := qbs.NewCondition("pid = ?", pdecl.Id).And("is_old = ?", true)
			_, err = q.Condition(cond).Delete(new(PkgFunc))
			if err != nil {
				beego.Error("models.SaveProject(", pinfo.Path, ") -> Delete functions:", err)
			}
		}
	}

	// ------------- END ------------

	// ------------------------------
	// Update imported information.
	// ------------------------------

	if info.Id > 0 {
		// Current package.
		importeds := strings.Split(
			strings.Replace(info.ImportPid, "$", "", -1), "|")
		importPids := make([]string, 0, len(importeds))
		for _, v := range importeds {
			pid, _ := strconv.ParseInt(v, 10, 64)
			if checkImport(q, info.Path, pid) {
				importPids = append(importPids, "$"+v)
			}
		}
		pinfo.ImportPid = strings.Join(importPids, "|")
		pinfo.ImportedNum = len(importPids)
	}

	_, err = q.Save(pinfo)
	if err != nil {
		beego.Error("models.SaveProject(", pinfo.Path, ") -> Information:", err)
	}

	// Don't need to check standard library.
	if imports != nil && !utils.IsGoRepoPath(pinfo.Path) {
		// Other packages.
		for _, v := range imports {
			if !utils.IsGoRepoPath(v) {
				// Only count non-standard library.
				updateImportInfo(q, v, int(pinfo.Id), true)
			}
		}
	}
	// ------------- END ------------

	return nil
}

// checkImport returns true if the package(id) imports given package(path).
func checkImport(q *qbs.Qbs, path string, id int64) bool {
	pinfo := &PkgInfo{
		Id: id,
	}
	err := q.Find(pinfo)
	if err != nil {
		return false
	}

	decl := new(PkgDecl)
	cond := qbs.NewCondition("pid = ?", pinfo.Id).And("tag = ?", "")
	err = q.Condition(cond).Find(decl)
	if err != nil {
		return false
	}

	if strings.Index(decl.Imports, path) == -1 {
		return false
	}

	return true
}

// LoadProject returns package declaration.
func LoadProject(pid int64, tag string) (*PkgDecl, error) {
	// Check path length to reduce connect times.
	if pid == 0 {
		return nil, errors.New("models.LoadProject -> Zero id.")
	}

	q := connDb()
	defer q.Close()

	pdecl := new(PkgDecl)
	cond := qbs.NewCondition("pid = ?", pid).And("tag = ?", tag)
	err := q.Condition(cond).Find(pdecl)
	return pdecl, err
}

// DeleteProject deletes everything about the path in database, and update import information.
func DeleteProject(path string) error {
	// Check path length to reduce connect times. (except launchpad.net)
	if path[0] != 'l' && len(strings.Split(path, "/")) <= 2 {
		beego.Error("models.DeleteProject(", path, ") -> Short path as not needed")
		return nil
	}

	q := connDb()
	defer q.Close()

	var i1, i2, i3, i4, i5 int64
	// Delete package information.
	info := new(PkgInfo)
	err := q.WhereEqual("path", path).Find(info)
	if err == nil {
		i1, err = q.WhereEqual("path", path).Delete(info)
		if err != nil {
			beego.Error("models.DeleteProject(", path, ") -> Information:", err)
		}
	}

	// Delete package declaration.
	if info.Id > 0 {
		// Find.
		var pdecls []*PkgDecl
		err = q.WhereEqual("pid", info.Id).FindAll(&pdecls)
		if err != nil {
			beego.Error("models.DeleteProject(", path, ") -> Find declaration:", err)
		}

		// Update.
		if !utils.IsGoRepoPath(path) {
			for _, pd := range pdecls {
				// Don't need to check standard library.
				// Update import information.
				imports := strings.Split(pd.Imports, "|")
				imports = imports[:len(imports)-1]
				for _, v := range imports {
					if !utils.IsGoRepoPath(v) {
						// Only count non-standard library.
						updateImportInfo(q, v, int(info.Id), false)
					}
				}
			}
		}

		// Delete.
		i2, err = q.WhereEqual("pid", info.Id).Delete(new(PkgDecl))
		if err != nil {
			beego.Error("models.DeleteProject(", path, ") -> Delete declaration:", err)
		}
	}

	// Delete package documentation.
	i3, err = q.WhereEqual("path", path).Delete(new(PkgDoc))
	if err != nil {
		beego.Error("models.DeleteProject(", path, ") -> Documentation:", err)
	}

	// Delete package examples.
	i4, err = q.WhereEqual("path", path).Delete(new(PkgExam))
	if err != nil {
		beego.Error("models.DeleteProject(", path, ") -> Examples:", err)
	}

	// Delete package functions.
	if info.Id > 0 {
		i5, err = q.WhereEqual("path", path).Delete(new(PkgExam))
		if err != nil {
			beego.Error("models.DeleteProject(", path, ") -> Functions:", err)
		}
	}

	if i1+i2+i3+i4+i5 > 0 {
		beego.Info("models.DeleteProject(", path, i1, i2, i3, i4, i5, ")")
	}

	return nil
}

func updateImportInfo(q *qbs.Qbs, path string, pid int, add bool) {
	// Save package information.
	info := new(PkgInfo)
	err := q.WhereEqual("path", path).Find(info)
	if err == nil {
		// Check if pid exists in this project.
		i := strings.Index(info.ImportPid, "$"+strconv.Itoa(pid)+"|")
		switch {
		case i == -1 && add: // Add operation and does not contain.
			info.ImportPid += "$" + strconv.Itoa(pid) + "|"
			info.ImportedNum = len(strings.Split(info.ImportPid, "|")) - 1
			_, err = q.Save(info)
			if err != nil {
				beego.Error("models.updateImportInfo -> add:", path, err)
			}
		case i > -1 && !add: // Delete operation and contains.
			info.ImportPid = strings.Replace(info.ImportPid, "$"+strconv.Itoa(pid)+"|", "", 1)
			info.ImportedNum = len(strings.Split(info.ImportPid, "|")) - 1
			_, err = q.Save(info)
			if err != nil {
				beego.Error("models.updateImportInfo -> delete:", path, err)
			}
		}
	}

	// Error means this project does not exist, simply skip.
}

// FlushCacheProjects saves cache data to database.
func FlushCacheProjects(pinfos []*PkgInfo, procks []*PkgRock) {
	q := connDb()
	defer q.Close()

	// Update project data.
	for _, p := range pinfos {
		info := new(PkgInfo)
		err := q.WhereEqual("path", p.Path).Find(info)
		if err == nil {
			// Shoule always be nil, just in case not exist.
			p.Id = info.Id
			// Limit 10 views each period.
			if p.Views-info.Views > 10 {
				p.Views = info.Views + 10
			}
		}
		_, err = q.Save(p)
		if err != nil {
			beego.Error("models.FlushCacheProjects(", p.Path, ") ->", err)
		}
	}

	// Update rock this week.
	if time.Now().UTC().Weekday() == time.Monday && utils.Cfg.MustBool("task", "rock_reset") {
		utils.Cfg.SetValue("task", "rock_reset", "0")
		utils.SaveConfig()
		// Reset rock table.
		_, err := q.Where("id > ?", int64(0)).Delete(new(PkgRock))
		if err != nil {
			beego.Error("models.FlushCacheProjects -> Reset rock table:", err)
		}
	} else if time.Now().UTC().Weekday() != time.Monday && !utils.Cfg.MustBool("task", "rock_reset") {
		utils.Cfg.SetValue("task", "rock_reset", "1")
		utils.SaveConfig()
	}

	for _, pr := range procks {
		r := new(PkgRock)
		err := q.WhereEqual("pid", pr.Pid).Find(r)
		if err == nil {
			pr.Id = r.Id
			r.Delta += pr.Rank - r.Rank
			pr.Delta = r.Delta
		}
		q.Save(pr)
	}
}
<|endoftext|>


<commit_before>package kcp

import (
	"crypto/sha1"
	"fmt"
	"log"
	"net"
	"net/http"
	_ "net/http/pprof"
	"sync"
	"testing"
	"time"

	"golang.org/x/crypto/pbkdf2"
)

const portEcho = "127.0.0.1:9999"
const portSink = "127.0.0.1:19999"
const salt = "kcptest"

var key = []byte("testkey")
var fec = 4
var pass = pbkdf2.Key(key, []byte(portSink), 4096, 32, sha1.New)

func init() {
	go func() {
		log.Println(http.ListenAndServe("localhost:6060", nil))
	}()

	go echoServer()
	go sinkServer()
	println("beginning tests, encryption:salsa20, fec:10/3")
}

func dialEcho() (*UDPSession, error) {
	//block, _ := NewNoneBlockCrypt(pass)
	//block, _ := NewSimpleXORBlockCrypt(pass)
	//block, _ := NewTEABlockCrypt(pass[:16])
	//block, _ := NewAESBlockCrypt(pass)
	block, _ := NewSalsa20BlockCrypt(pass)
	sess, err := DialWithOptions(portEcho, block, 10, 3)
	if err != nil {
		panic(err)
	}

	sess.SetStreamMode(true)
	sess.SetWindowSize(4096, 4096)
	sess.SetReadBuffer(4 * 1024 * 1024)
	sess.SetWriteBuffer(4 * 1024 * 1024)
	sess.SetStreamMode(true)
	sess.SetNoDelay(1, 10, 2, 1)
	sess.SetMtu(1400)
	sess.SetACKNoDelay(true)
	sess.SetDeadline(time.Now().Add(time.Minute))
	return sess, err
}

func dialSink() (*UDPSession, error) {
	block, _ := NewNoneBlockCrypt(pass)
	//block, _ := NewSimpleXORBlockCrypt(pass)
	//block, _ := NewTEABlockCrypt(pass[:16])
	//block, _ := NewAESBlockCrypt(pass)
	//block, _ := NewSalsa20BlockCrypt(pass)
	sess, err := DialWithOptions(portSink, block, 10, 3)
	if err != nil {
		panic(err)
	}

	sess.SetStreamMode(true)
	sess.SetWindowSize(4096, 4096)
	sess.SetReadBuffer(4 * 1024 * 1024)
	sess.SetWriteBuffer(4 * 1024 * 1024)
	sess.SetStreamMode(true)
	sess.SetNoDelay(1, 10, 2, 1)
	sess.SetMtu(1400)
	sess.SetACKNoDelay(true)
	sess.SetDeadline(time.Now().Add(time.Minute))
	return sess, err
}

//////////////////////////

func listenEcho() (net.Listener, error) {
	//block, _ := NewNoneBlockCrypt(pass)
	//block, _ := NewSimpleXORBlockCrypt(pass)
	//block, _ := NewTEABlockCrypt(pass[:16])
	//block, _ := NewAESBlockCrypt(pass)
	block, _ := NewSalsa20BlockCrypt(pass)
	return ListenWithOptions(portEcho, block, 10, 3)
}

func listenSink() (net.Listener, error) {
	block, _ := NewNoneBlockCrypt(pass)
	//block, _ := NewSimpleXORBlockCrypt(pass)
	//block, _ := NewTEABlockCrypt(pass[:16])
	//block, _ := NewAESBlockCrypt(pass)
	//block, _ := NewSalsa20BlockCrypt(pass)
	return ListenWithOptions(portSink, block, 10, 3)
}

func echoServer() {
	l, err := listenEcho()
	if err != nil {
		panic(err)
	}

	go func() {
		kcplistener := l.(*Listener)
		kcplistener.SetReadBuffer(4 * 1024 * 1024)
		kcplistener.SetWriteBuffer(4 * 1024 * 1024)
		kcplistener.SetDSCP(46)
		for {
			s, err := l.Accept()
			if err != nil {
				return
			}

			// coverage test
			s.(*UDPSession).SetReadBuffer(4 * 1024 * 1024)
			s.(*UDPSession).SetWriteBuffer(4 * 1024 * 1024)
			go handleEcho(s.(*UDPSession))
		}
	}()
}

func sinkServer() {
	l, err := listenSink()
	if err != nil {
		panic(err)
	}

	go func() {
		kcplistener := l.(*Listener)
		kcplistener.SetReadBuffer(4 * 1024 * 1024)
		kcplistener.SetWriteBuffer(4 * 1024 * 1024)
		kcplistener.SetDSCP(46)
		for {
			s, err := l.Accept()
			if err != nil {
				return
			}

			go handleSink(s.(*UDPSession))
		}
	}()
}

///////////////////////////

func handleEcho(conn *UDPSession) {
	conn.SetStreamMode(true)
	conn.SetWindowSize(4096, 4096)
	conn.SetNoDelay(1, 10, 2, 1)
	conn.SetDSCP(46)
	conn.SetMtu(1400)
	conn.SetACKNoDelay(false)
	conn.SetReadDeadline(time.Now().Add(time.Hour))
	conn.SetWriteDeadline(time.Now().Add(time.Hour))
	buf := make([]byte, 65536)
	for {
		n, err := conn.Read(buf)
		if err != nil {
			panic(err)
		}
		conn.Write(buf[:n])
	}
}

func handleSink(conn *UDPSession) {
	conn.SetStreamMode(true)
	conn.SetWindowSize(4096, 4096)
	conn.SetNoDelay(1, 10, 2, 1)
	conn.SetDSCP(46)
	conn.SetMtu(1400)
	conn.SetACKNoDelay(false)
	conn.SetReadDeadline(time.Now().Add(time.Hour))
	conn.SetWriteDeadline(time.Now().Add(time.Hour))
	buf := make([]byte, 65536)
	for {
		_, err := conn.Read(buf)
		if err != nil {
			panic(err)
		}
	}
}

///////////////////////////

func TestTimeout(t *testing.T) {
	cli, err := dialEcho()
	if err != nil {
		panic(err)
	}
	buf := make([]byte, 10)

	//timeout
	cli.SetDeadline(time.Now().Add(time.Second))
	<-time.After(2 * time.Second)
	n, err := cli.Read(buf)
	if n != 0 || err == nil {
		t.Fail()
	}
	n, err = cli.Write(buf)
	if n != 0 || err == nil {
		t.Fail()
	}
	cli.Close()
}

func TestSendRecv(t *testing.T) {
	cli, err := dialEcho()
	if err != nil {
		panic(err)
	}
	const N = 100
	buf := make([]byte, 10)
	for i := 0; i < N; i++ {
		msg := fmt.Sprintf("hello%v", i)
		cli.Write([]byte(msg))
		if n, err := cli.Read(buf); err == nil {
			if string(buf[:n]) != msg {
				t.Fail()
			}
		} else {
			panic(err)
		}
	}
	cli.Close()
}

func TestClose(t *testing.T) {
	cli, err := dialEcho()
	if err != nil {
		panic(err)
	}
	buf := make([]byte, 10)

	cli.Close()
	if cli.Close() == nil {
		t.Fail()
	}
	n, err := cli.Write(buf)
	if n != 0 || err == nil {
		t.Fail()
	}
	n, err = cli.Read(buf)
	if n != 0 || err == nil {
		t.Fail()
	}
	cli.Close()
}

func TestParallel1024CLIENT_64BMSG_64CNT(t *testing.T) {
	var wg sync.WaitGroup
	wg.Add(1024)
	for i := 0; i < 1024; i++ {
		go parallel_client(&wg)
	}
	wg.Wait()
}

func parallel_client(wg *sync.WaitGroup) (err error) {
	cli, err := dialEcho()
	if err != nil {
		panic(err)
	}

	err = echo_tester(cli, 64, 64)
	wg.Done()
	return
}

func BenchmarkEchoSpeed4K(b *testing.B) {
	speedclient(b, 4096)
}

func BenchmarkEchoSpeed64K(b *testing.B) {
	speedclient(b, 65536)
}

func BenchmarkEchoSpeed512K(b *testing.B) {
	speedclient(b, 524288)
}

func BenchmarkEchoSpeed1M(b *testing.B) {
	speedclient(b, 1048576)
}

func speedclient(b *testing.B, nbytes int) {
	cli, err := dialEcho()
	if err != nil {
		panic(err)
	}

	if err := echo_tester(cli, nbytes, b.N); err != nil {
		b.Fail()
	}
	b.SetBytes(int64(nbytes))
}

func BenchmarkSinkSpeed4K(b *testing.B) {
	sinkclient(b, 4096)
}

func BenchmarkSinkSpeed64K(b *testing.B) {
	sinkclient(b, 65536)
}

func BenchmarkSinkSpeed256K(b *testing.B) {
	sinkclient(b, 524288)
}

func BenchmarkSinkSpeed1M(b *testing.B) {
	sinkclient(b, 1048576)
}

func sinkclient(b *testing.B, nbytes int) {
	cli, err := dialSink()
	if err != nil {
		panic(err)
	}

	sink_tester(cli, nbytes, b.N)
	b.SetBytes(int64(nbytes))
}

func echo_tester(cli net.Conn, msglen, msgcount int) error {
	for i := 0; i < msgcount; i++ {
		// send packet
		buf := make([]byte, msglen)
		if _, err := cli.Write(buf); err != nil {
			return err
		}

		// receive packet
		nrecv := 0
		for {
			n, err := cli.Read(buf)
			if err != nil {
				return err
			} else {
				nrecv += n
				if nrecv == msglen {
					break
				}
			}
		}
	}
	return nil
}

func sink_tester(cli *UDPSession, msglen, msgcount int) error {
	// sender
	buf := make([]byte, msglen)
	for i := 0; i < msgcount; i++ {
		if _, err := cli.Write(buf); err != nil {
			return err
		}
	}

	// window checker
	for {
		cli.mu.Lock()
		waitsnd := cli.kcp.WaitSnd()
		cli.mu.Unlock()
		if waitsnd != 0 {
			<-time.After(10 * time.Millisecond)
		} else {
			return nil
		}
	}
}

func TestSNMP(t *testing.T) {
	t.Log(DefaultSnmp.Copy())
	t.Log(DefaultSnmp.Header())
	t.Log(DefaultSnmp.ToSlice())
	DefaultSnmp.Reset()
	t.Log(DefaultSnmp.ToSlice())
}
<commit_msg>add tests to memory optimized Read<commit_after>package kcp

import (
	"crypto/sha1"
	"fmt"
	"io"
	"log"
	"net"
	"net/http"
	_ "net/http/pprof"
	"sync"
	"testing"
	"time"

	"golang.org/x/crypto/pbkdf2"
)

const portEcho = "127.0.0.1:9999"
const portSink = "127.0.0.1:19999"
const portTinyBufferEcho = "127.0.0.1:29999"
const salt = "kcptest"

var key = []byte("testkey")
var fec = 4
var pass = pbkdf2.Key(key, []byte(portSink), 4096, 32, sha1.New)

func init() {
	go func() {
		log.Println(http.ListenAndServe("localhost:6060", nil))
	}()

	go echoServer()
	go sinkServer()
	go tinyBufferEchoServer()
	println("beginning tests, encryption:salsa20, fec:10/3")
}

func dialEcho() (*UDPSession, error) {
	//block, _ := NewNoneBlockCrypt(pass)
	//block, _ := NewSimpleXORBlockCrypt(pass)
	//block, _ := NewTEABlockCrypt(pass[:16])
	//block, _ := NewAESBlockCrypt(pass)
	block, _ := NewSalsa20BlockCrypt(pass)
	sess, err := DialWithOptions(portEcho, block, 10, 3)
	if err != nil {
		panic(err)
	}

	sess.SetStreamMode(true)
	sess.SetStreamMode(false)
	sess.SetStreamMode(true)
	sess.SetWindowSize(4096, 4096)
	sess.SetReadBuffer(4 * 1024 * 1024)
	sess.SetWriteBuffer(4 * 1024 * 1024)
	sess.SetStreamMode(true)
	sess.SetNoDelay(1, 10, 2, 1)
	sess.SetMtu(1400)
	sess.SetMtu(1600)
	sess.SetMtu(1400)
	sess.SetACKNoDelay(true)
	sess.SetDeadline(time.Now().Add(time.Minute))
	return sess, err
}

func dialSink() (*UDPSession, error) {
	block, _ := NewNoneBlockCrypt(pass)
	//block, _ := NewSimpleXORBlockCrypt(pass)
	//block, _ := NewTEABlockCrypt(pass[:16])
	//block, _ := NewAESBlockCrypt(pass)
	//block, _ := NewSalsa20BlockCrypt(pass)
	sess, err := DialWithOptions(portSink, block, 10, 3)
	if err != nil {
		panic(err)
	}

	sess.SetStreamMode(true)
	sess.SetWindowSize(4096, 4096)
	sess.SetReadBuffer(4 * 1024 * 1024)
	sess.SetWriteBuffer(4 * 1024 * 1024)
	sess.SetStreamMode(true)
	sess.SetNoDelay(1, 10, 2, 1)
	sess.SetMtu(1400)
	sess.SetACKNoDelay(true)
	sess.SetDeadline(time.Now().Add(time.Minute))
	return sess, err
}

func dialTinyBufferEcho() (*UDPSession, error) {
	//block, _ := NewNoneBlockCrypt(pass)
	//block, _ := NewSimpleXORBlockCrypt(pass)
	//block, _ := NewTEABlockCrypt(pass[:16])
	//block, _ := NewAESBlockCrypt(pass)
	block, _ := NewSalsa20BlockCrypt(pass)
	sess, err := DialWithOptions(portTinyBufferEcho, block, 10, 3)
	if err != nil {
		panic(err)
	}
	return sess, err
}

//////////////////////////

func listenEcho() (net.Listener, error) {
	//block, _ := NewNoneBlockCrypt(pass)
	//block, _ := NewSimpleXORBlockCrypt(pass)
	//block, _ := NewTEABlockCrypt(pass[:16])
	//block, _ := NewAESBlockCrypt(pass)
	block, _ := NewSalsa20BlockCrypt(pass)
	return ListenWithOptions(portEcho, block, 10, 3)
}

func listenTinyBufferEcho() (net.Listener, error) {
	//block, _ := NewNoneBlockCrypt(pass)
	//block, _ := NewSimpleXORBlockCrypt(pass)
	//block, _ := NewTEABlockCrypt(pass[:16])
	//block, _ := NewAESBlockCrypt(pass)
	block, _ := NewSalsa20BlockCrypt(pass)
	return ListenWithOptions(portTinyBufferEcho, block, 10, 3)
}

func listenSink() (net.Listener, error) {
	block, _ := NewNoneBlockCrypt(pass)
	//block, _ := NewSimpleXORBlockCrypt(pass)
	//block, _ := NewTEABlockCrypt(pass[:16])
	//block, _ := NewAESBlockCrypt(pass)
	//block, _ := NewSalsa20BlockCrypt(pass)
	return ListenWithOptions(portSink, block, 10, 3)
}

func echoServer() {
	l, err := listenEcho()
	if err != nil {
		panic(err)
	}

	go func() {
		kcplistener := l.(*Listener)
		kcplistener.SetReadBuffer(4 * 1024 * 1024)
		kcplistener.SetWriteBuffer(4 * 1024 * 1024)
		kcplistener.SetDSCP(46)
		for {
			s, err := l.Accept()
			if err != nil {
				return
			}

			// coverage test
			s.(*UDPSession).SetReadBuffer(4 * 1024 * 1024)
			s.(*UDPSession).SetWriteBuffer(4 * 1024 * 1024)
			go handleEcho(s.(*UDPSession))
		}
	}()
}

func sinkServer() {
	l, err := listenSink()
	if err != nil {
		panic(err)
	}

	go func() {
		kcplistener := l.(*Listener)
		kcplistener.SetReadBuffer(4 * 1024 * 1024)
		kcplistener.SetWriteBuffer(4 * 1024 * 1024)
		kcplistener.SetDSCP(46)
		for {
			s, err := l.Accept()
			if err != nil {
				return
			}

			go handleSink(s.(*UDPSession))
		}
	}()
}

func tinyBufferEchoServer() {
	l, err := listenTinyBufferEcho()
	if err != nil {
		panic(err)
	}

	go func() {
		for {
			s, err := l.Accept()
			if err != nil {
				return
			}
			go handleTinyBufferEcho(s.(*UDPSession))
		}
	}()
}

///////////////////////////

func handleEcho(conn *UDPSession) {
	conn.SetStreamMode(true)
	conn.SetWindowSize(4096, 4096)
	conn.SetNoDelay(1, 10, 2, 1)
	conn.SetDSCP(46)
	conn.SetMtu(1400)
	conn.SetACKNoDelay(false)
	conn.SetReadDeadline(time.Now().Add(time.Hour))
	conn.SetWriteDeadline(time.Now().Add(time.Hour))
	buf := make([]byte, 65536)
	for {
		n, err := conn.Read(buf)
		if err != nil {
			panic(err)
		}
		conn.Write(buf[:n])
	}
}

func handleSink(conn *UDPSession) {
	conn.SetStreamMode(true)
	conn.SetWindowSize(4096, 4096)
	conn.SetNoDelay(1, 10, 2, 1)
	conn.SetDSCP(46)
	conn.SetMtu(1400)
	conn.SetACKNoDelay(false)
	conn.SetReadDeadline(time.Now().Add(time.Hour))
	conn.SetWriteDeadline(time.Now().Add(time.Hour))
	buf := make([]byte, 65536)
	for {
		_, err := conn.Read(buf)
		if err != nil {
			panic(err)
		}
	}
}

func handleTinyBufferEcho(conn *UDPSession) {
	conn.SetStreamMode(true)
	buf := make([]byte, 2)
	for {
		n, err := conn.Read(buf)
		if err != nil {
			panic(err)
		}
		conn.Write(buf[:n])
	}
}

///////////////////////////

func TestTimeout(t *testing.T) {
	cli, err := dialEcho()
	if err != nil {
		panic(err)
	}
	buf := make([]byte, 10)

	//timeout
	cli.SetDeadline(time.Now().Add(time.Second))
	<-time.After(2 * time.Second)
	n, err := cli.Read(buf)
	if n != 0 || err == nil {
		t.Fail()
	}
	n, err = cli.Write(buf)
	if n != 0 || err == nil {
		t.Fail()
	}
	cli.Close()
}

func TestSendRecv(t *testing.T) {
	cli, err := dialEcho()
	if err != nil {
		panic(err)
	}
	const N = 100
	buf := make([]byte, 10)
	for i := 0; i < N; i++ {
		msg := fmt.Sprintf("hello%v", i)
		cli.Write([]byte(msg))
		if n, err := cli.Read(buf); err == nil {
			if string(buf[:n]) != msg {
				t.Fail()
			}
		} else {
			panic(err)
		}
	}
	cli.Close()
}

func TestTinyBufferReceiver(t *testing.T) {
	cli, err := dialTinyBufferEcho()
	if err != nil {
		panic(err)
	}
	const N = 100
	snd := byte(0)
	fillBuffer := func(buf []byte) {
		for i := 0; i < len(buf); i++ {
			buf[i] = snd
			snd++
		}
	}

	rcv := byte(0)
	check := func(buf []byte) bool {
		for i := 0; i < len(buf); i++ {
			if buf[i] != rcv {
				return false
			}
			rcv++
		}
		return true
	}
	sndbuf := make([]byte, 7)
	rcvbuf := make([]byte, 7)
	for i := 0; i < N; i++ {
		fillBuffer(sndbuf)
		cli.Write(sndbuf)
		if n, err := io.ReadFull(cli, rcvbuf); err == nil {
			if !check(rcvbuf[:n]) {
				t.Fail()
			}
		} else {
			panic(err)
		}
	}
	cli.Close()
}

func TestClose(t *testing.T) {
	cli, err := dialEcho()
	if err != nil {
		panic(err)
	}
	buf := make([]byte, 10)

	cli.Close()
	if cli.Close() == nil {
		t.Fail()
	}
	n, err := cli.Write(buf)
	if n != 0 || err == nil {
		t.Fail()
	}
	n, err = cli.Read(buf)
	if n != 0 || err == nil {
		t.Fail()
	}
	cli.Close()
}

func TestParallel1024CLIENT_64BMSG_64CNT(t *testing.T) {
	var wg sync.WaitGroup
	wg.Add(1024)
	for i := 0; i < 1024; i++ {
		go parallel_client(&wg)
	}
	wg.Wait()
}

func parallel_client(wg *sync.WaitGroup) (err error) {
	cli, err := dialEcho()
	if err != nil {
		panic(err)
	}

	err = echo_tester(cli, 64, 64)
	wg.Done()
	return
}

func BenchmarkEchoSpeed4K(b *testing.B) {
	speedclient(b, 4096)
}

func BenchmarkEchoSpeed64K(b *testing.B) {
	speedclient(b, 65536)
}

func BenchmarkEchoSpeed512K(b *testing.B) {
	speedclient(b, 524288)
}

func BenchmarkEchoSpeed1M(b *testing.B) {
	speedclient(b, 1048576)
}

func speedclient(b *testing.B, nbytes int) {
	cli, err := dialEcho()
	if err != nil {
		panic(err)
	}

	if err := echo_tester(cli, nbytes, b.N); err != nil {
		b.Fail()
	}
	b.SetBytes(int64(nbytes))
}

func BenchmarkSinkSpeed4K(b *testing.B) {
	sinkclient(b, 4096)
}

func BenchmarkSinkSpeed64K(b *testing.B) {
	sinkclient(b, 65536)
}

func BenchmarkSinkSpeed256K(b *testing.B) {
	sinkclient(b, 524288)
}

func BenchmarkSinkSpeed1M(b *testing.B) {
	sinkclient(b, 1048576)
}

func sinkclient(b *testing.B, nbytes int) {
	cli, err := dialSink()
	if err != nil {
		panic(err)
	}

	sink_tester(cli, nbytes, b.N)
	b.SetBytes(int64(nbytes))
}

func echo_tester(cli net.Conn, msglen, msgcount int) error {
	for i := 0; i < msgcount; i++ {
		// send packet
		buf := make([]byte, msglen)
		if _, err := cli.Write(buf); err != nil {
			return err
		}

		// receive packet
		nrecv := 0
		for {
			n, err := cli.Read(buf)
			if err != nil {
				return err
			} else {
				nrecv += n
				if nrecv == msglen {
					break
				}
			}
		}
	}
	return nil
}

func sink_tester(cli *UDPSession, msglen, msgcount int) error {
	// sender
	buf := make([]byte, msglen)
	for i := 0; i < msgcount; i++ {
		if _, err := cli.Write(buf); err != nil {
			return err
		}
	}

	// window checker
	for {
		cli.mu.Lock()
		waitsnd := cli.kcp.WaitSnd()
		cli.mu.Unlock()
		if waitsnd != 0 {
			<-time.After(10 * time.Millisecond)
		} else {
			return nil
		}
	}
}

func TestSNMP(t *testing.T) {
	t.Log(DefaultSnmp.Copy())
	t.Log(DefaultSnmp.Header())
	t.Log(DefaultSnmp.ToSlice())
	DefaultSnmp.Reset()
	t.Log(DefaultSnmp.ToSlice())
}
<|endoftext|>


<commit_before>package exec

import (
	"context"

	"code.cloudfoundry.org/lager"

	"github.com/concourse/atc/worker"
)

//go:generate counterfeiter . Action

type Action interface {
	Run(context.Context, lager.Logger, *worker.ArtifactRepository) error
	ExitStatus() ExitStatus
}

func NewActionsStep(
	logger lager.Logger, // TODO: can we move that to method? need to change all steps though
	actions []Action,
	buildEventsDelegate ActionsBuildEventsDelegate,
) ActionsStep {
	return ActionsStep{
		logger:              logger,
		actions:             actions,
		buildEventsDelegate: buildEventsDelegate,
	}
}

//go:generate counterfeiter . ActionsBuildEventsDelegate

type ActionsBuildEventsDelegate interface {
	ActionCompleted(lager.Logger, Action)
	Failed(lager.Logger, error)
}

// ActionsStep will execute actions in specified order and notify build events
// delegate about different execution events.
type ActionsStep struct {
	actions             []Action
	buildEventsDelegate ActionsBuildEventsDelegate

	logger lager.Logger // TODO: can we move that to method? need to change all steps though

	repository *worker.ArtifactRepository
	succeeded  bool
}

func (s ActionsStep) Using(repo *worker.ArtifactRepository) Step {
	s.repository = repo
	return &s
}

// Run will first call Initializing on build events delegate. Then it will call
// Run on every action. If any action fails it will notify delegate with Failed.
// It will call ActionCompleted after each action run that succeeds.
func (s *ActionsStep) Run(ctx context.Context) error {
	succeeded := true
	for _, action := range s.actions {
		err := action.Run(ctx, s.logger, s.repository)
		if err != nil {
			s.logger.Error("failed-to-run-action", err)
			s.buildEventsDelegate.Failed(s.logger, err)
			return err
		}

		s.buildEventsDelegate.ActionCompleted(s.logger, action)

		if action.ExitStatus() != 0 {
			succeeded = false
		}
	}

	s.succeeded = succeeded

	return nil
}

// Succeeded will return true if all actions exited with exit status 0.
func (s *ActionsStep) Succeeded() bool {
	return s.succeeded
}
<commit_msg>restore error messaging for aborting/interrupting<commit_after>package exec

import (
	"context"
	"errors"

	"code.cloudfoundry.org/lager"

	"github.com/concourse/atc/worker"
)

//go:generate counterfeiter . Action

type Action interface {
	Run(context.Context, lager.Logger, *worker.ArtifactRepository) error
	ExitStatus() ExitStatus
}

func NewActionsStep(
	logger lager.Logger, // TODO: can we move that to method? need to change all steps though
	actions []Action,
	buildEventsDelegate ActionsBuildEventsDelegate,
) ActionsStep {
	return ActionsStep{
		logger:              logger,
		actions:             actions,
		buildEventsDelegate: buildEventsDelegate,
	}
}

//go:generate counterfeiter . ActionsBuildEventsDelegate

type ActionsBuildEventsDelegate interface {
	ActionCompleted(lager.Logger, Action)
	Failed(lager.Logger, error)
}

// ActionsStep will execute actions in specified order and notify build events
// delegate about different execution events.
type ActionsStep struct {
	actions             []Action
	buildEventsDelegate ActionsBuildEventsDelegate

	logger lager.Logger // TODO: can we move that to method? need to change all steps though

	repository *worker.ArtifactRepository
	succeeded  bool
}

func (s ActionsStep) Using(repo *worker.ArtifactRepository) Step {
	s.repository = repo
	return &s
}

// Run will first call Initializing on build events delegate. Then it will call
// Run on every action. If any action fails it will notify delegate with Failed.
// It will call ActionCompleted after each action run that succeeds.
func (s *ActionsStep) Run(ctx context.Context) error {
	succeeded := true
	for _, action := range s.actions {
		err := action.Run(ctx, s.logger, s.repository)
		if err != nil {
			if err == context.DeadlineExceeded || err == context.Canceled {
				// XXX: here for backwards-compatibility. this shouldn't be so coupled.
				s.buildEventsDelegate.Failed(s.logger, errors.New("interrupted"))
				return err
			}

			s.logger.Error("failed-to-run-action", err)
			s.buildEventsDelegate.Failed(s.logger, err)
			return err
		}

		s.buildEventsDelegate.ActionCompleted(s.logger, action)

		if action.ExitStatus() != 0 {
			succeeded = false
		}
	}

	s.succeeded = succeeded

	return nil
}

// Succeeded will return true if all actions exited with exit status 0.
func (s *ActionsStep) Succeeded() bool {
	return s.succeeded
}
<|endoftext|>


<commit_before>package qr

import (
	"github.com/boombuler/barcode"
	"image"
	"image/color"
	"math"
)

type qrcode struct {
	dimension int
	data      *barcode.BitList
	content   string
}

func (qr *qrcode) Content() string {
	return qr.content
}

func (qr *qrcode) Metadata() barcode.Metadata {
	return barcode.Metadata{"QR Code", 2}
}

func (qr *qrcode) ColorModel() color.Model {
	return color.Gray16Model
}

func (qr *qrcode) Bounds() image.Rectangle {
	return image.Rect(0, 0, qr.dimension, qr.dimension)
}

func (qr *qrcode) At(x, y int) color.Color {
	if qr.Get(x, y) {
		return color.Black
	}
	return color.White
}

func (qr *qrcode) Get(x, y int) bool {
	return qr.data.GetBit(x*qr.dimension + y)
}

func (qr *qrcode) Set(x, y int, val bool) {
	qr.data.SetBit(x*qr.dimension+y, val)
}

func (qr *qrcode) calcPenalty() uint {
	return qr.calcPenaltyRule1() + qr.calcPenaltyRule2() + qr.calcPenaltyRule3() + qr.calcPenaltyRule4()
}

func (qr *qrcode) calcPenaltyRule1() uint {
	var result uint = 0
	for x := 0; x < qr.dimension; x++ {
		checkForX := false
		var cntX uint = 0
		checkForY := false
		var cntY uint = 0

		for y := 0; y < qr.dimension; y++ {
			if qr.Get(x, y) == checkForX {
				cntX += 1
			} else {
				checkForX = !checkForX
				if cntX >= 5 {
					result += cntX - 2
				}
				cntX = 0
			}

			if qr.Get(y, x) == checkForY {
				cntY += 1
			} else {
				checkForY = !checkForY
				if cntY >= 5 {
					result += cntY - 2
				}
				cntY = 0
			}
		}

		if cntX >= 5 {
			result += cntX - 2
		}
		if cntY >= 5 {
			result += cntY - 2
		}
	}

	return result
}

func (qr *qrcode) calcPenaltyRule2() uint {
	var result uint = 0
	for x := 0; x < qr.dimension-1; x++ {
		for y := 0; y < qr.dimension-1; y++ {
			check := qr.Get(x, y)
			if qr.Get(x, y+1) == check && qr.Get(x+1, y) == check && qr.Get(x+1, y+1) == check {
				result += 3
			}
		}
	}
	return result
}

func (qr *qrcode) calcPenaltyRule3() uint {
	var pattern1 []bool = []bool{true, false, true, true, true, false, true, false, false, false, false}
	var pattern2 []bool = []bool{false, false, false, false, true, false, true, true, true, false, true}

	var result uint = 0
	for x := 0; x <= qr.dimension-len(pattern1); x++ {
		for y := 0; y < qr.dimension; y++ {
			pattern1XFound := true
			pattern2XFound := true
			pattern1YFound := true
			pattern2YFound := true

			for i := 0; i < len(pattern1); i++ {
				iv := qr.Get(x+i, y)
				if iv != pattern1[i] {
					pattern1XFound = false
				}
				if iv != pattern2[i] {
					pattern2XFound = false
				}
				iv = qr.Get(y, x+i)
				if iv != pattern1[i] {
					pattern1YFound = false
				}
				if iv != pattern2[i] {
					pattern2YFound = false
				}
			}
			if pattern1XFound || pattern2XFound {
				result += 40
			}
			if pattern1YFound || pattern2YFound {
				result += 40
			}
		}
	}

	return result
}

func (qr *qrcode) calcPenaltyRule4() uint {
	totalNum := qr.data.Len()
	trueCnt := 0
	for i := 0; i < totalNum; i++ {
		if qr.data.GetBit(i) {
			trueCnt += 1
		}
	}
	percDark := float64(trueCnt) * 100 / float64(totalNum)
	floor := math.Abs(math.Floor(percDark/5) - 10)
	ceil := math.Abs(math.Ceil(percDark/5) - 10)
	return uint(math.Min(floor, ceil) * 10)
}

func newBarcode(dim int) *qrcode {
	res := new(qrcode)
	res.dimension = dim
	res.data = barcode.NewBitList(dim * dim)
	return res
}
<commit_msg>fixed bug in calcPenaltyRule1<commit_after>package qr

import (
	"github.com/boombuler/barcode"
	"image"
	"image/color"
	"math"
)

type qrcode struct {
	dimension int
	data      *barcode.BitList
	content   string
}

func (qr *qrcode) Content() string {
	return qr.content
}

func (qr *qrcode) Metadata() barcode.Metadata {
	return barcode.Metadata{"QR Code", 2}
}

func (qr *qrcode) ColorModel() color.Model {
	return color.Gray16Model
}

func (qr *qrcode) Bounds() image.Rectangle {
	return image.Rect(0, 0, qr.dimension, qr.dimension)
}

func (qr *qrcode) At(x, y int) color.Color {
	if qr.Get(x, y) {
		return color.Black
	}
	return color.White
}

func (qr *qrcode) Get(x, y int) bool {
	return qr.data.GetBit(x*qr.dimension + y)
}

func (qr *qrcode) Set(x, y int, val bool) {
	qr.data.SetBit(x*qr.dimension+y, val)
}

func (qr *qrcode) calcPenalty() uint {
	return qr.calcPenaltyRule1() + qr.calcPenaltyRule2() + qr.calcPenaltyRule3() + qr.calcPenaltyRule4()
}

func (qr *qrcode) calcPenaltyRule1() uint {
	var result uint = 0
	for x := 0; x < qr.dimension; x++ {
		checkForX := false
		var cntX uint = 0
		checkForY := false
		var cntY uint = 0

		for y := 0; y < qr.dimension; y++ {
			if qr.Get(x, y) == checkForX {
				cntX += 1
			} else {
				checkForX = !checkForX
				if cntX >= 5 {
					result += cntX - 2
				}
				cntX = 1
			}

			if qr.Get(y, x) == checkForY {
				cntY += 1
			} else {
				checkForY = !checkForY
				if cntY >= 5 {
					result += cntY - 2
				}
				cntY = 1
			}
		}

		if cntX >= 5 {
			result += cntX - 2
		}
		if cntY >= 5 {
			result += cntY - 2
		}
	}

	return result
}

func (qr *qrcode) calcPenaltyRule2() uint {
	var result uint = 0
	for x := 0; x < qr.dimension-1; x++ {
		for y := 0; y < qr.dimension-1; y++ {
			check := qr.Get(x, y)
			if qr.Get(x, y+1) == check && qr.Get(x+1, y) == check && qr.Get(x+1, y+1) == check {
				result += 3
			}
		}
	}
	return result
}

func (qr *qrcode) calcPenaltyRule3() uint {
	var pattern1 []bool = []bool{true, false, true, true, true, false, true, false, false, false, false}
	var pattern2 []bool = []bool{false, false, false, false, true, false, true, true, true, false, true}

	var result uint = 0
	for x := 0; x <= qr.dimension-len(pattern1); x++ {
		for y := 0; y < qr.dimension; y++ {
			pattern1XFound := true
			pattern2XFound := true
			pattern1YFound := true
			pattern2YFound := true

			for i := 0; i < len(pattern1); i++ {
				iv := qr.Get(x+i, y)
				if iv != pattern1[i] {
					pattern1XFound = false
				}
				if iv != pattern2[i] {
					pattern2XFound = false
				}
				iv = qr.Get(y, x+i)
				if iv != pattern1[i] {
					pattern1YFound = false
				}
				if iv != pattern2[i] {
					pattern2YFound = false
				}
			}
			if pattern1XFound || pattern2XFound {
				result += 40
			}
			if pattern1YFound || pattern2YFound {
				result += 40
			}
		}
	}

	return result
}

func (qr *qrcode) calcPenaltyRule4() uint {
	totalNum := qr.data.Len()
	trueCnt := 0
	for i := 0; i < totalNum; i++ {
		if qr.data.GetBit(i) {
			trueCnt += 1
		}
	}
	percDark := float64(trueCnt) * 100 / float64(totalNum)
	floor := math.Abs(math.Floor(percDark/5) - 10)
	ceil := math.Abs(math.Ceil(percDark/5) - 10)
	return uint(math.Min(floor, ceil) * 10)
}

func newBarcode(dim int) *qrcode {
	res := new(qrcode)
	res.dimension = dim
	res.data = barcode.NewBitList(dim * dim)
	return res
}
<|endoftext|>


<commit_before>package executor

import (
	"encoding/json"
	"sync"
	"time"

	"code.google.com/p/goprotobuf/proto"
	"github.com/GoogleCloudPlatform/kubernetes/pkg/api"
	"github.com/GoogleCloudPlatform/kubernetes/pkg/kubelet"
	log "github.com/golang/glog"
	"github.com/mesos/mesos-go/mesos"
	"gopkg.in/v1/yaml"
)

const (
	containerPollTime = 300 * time.Millisecond
	launchGracePeriod = 5 * time.Minute
)

type kuberTask struct {
	mesosTaskInfo *mesos.TaskInfo
	podName       string
}

// KubernetesExecutor is an mesos executor that runs pods
// in a minion machine.
type KubernetesExecutor struct {
	kl         *kubelet.Kubelet // the kubelet instance.
	updateChan chan<- interface{}
	driver     mesos.ExecutorDriver
	registered bool
	tasks      map[string]*kuberTask
	pods       map[string]*api.BoundPod
	lock       sync.RWMutex
	sourcename string
}

// New creates a new kubernetes executor.
func New(driver mesos.ExecutorDriver, kl *kubelet.Kubelet, ch chan<- interface{}, ns string) *KubernetesExecutor {
	return &KubernetesExecutor{
		kl:         kl,
		updateChan: ch,
		driver:     driver,
		registered: false,
		tasks:      make(map[string]*kuberTask),
		pods:       make(map[string]*api.BoundPod),
		sourcename: ns,
	}
}

// Registered is called when the executor is successfully registered with the slave.
func (k *KubernetesExecutor) Registered(driver mesos.ExecutorDriver,
	executorInfo *mesos.ExecutorInfo, frameworkInfo *mesos.FrameworkInfo, slaveInfo *mesos.SlaveInfo) {
	log.Infof("Executor %v of framework %v registered with slave %v\n",
		executorInfo, frameworkInfo, slaveInfo)
	k.registered = true
}

// Reregistered is called when the executor is successfully re-registered with the slave.
// This can happen when the slave fails over.
func (k *KubernetesExecutor) Reregistered(driver mesos.ExecutorDriver, slaveInfo *mesos.SlaveInfo) {
	log.Infof("Reregistered with slave %v\n", slaveInfo)
	k.registered = true
}

// Disconnected is called when the executor is disconnected with the slave.
func (k *KubernetesExecutor) Disconnected(driver mesos.ExecutorDriver) {
	log.Infof("Slave is disconnected\n")
	k.registered = false
}

// LaunchTask is called when the executor receives a request to launch a task.
func (k *KubernetesExecutor) LaunchTask(driver mesos.ExecutorDriver, taskInfo *mesos.TaskInfo) {
	log.Infof("Launch task %v\n", taskInfo)

	if !k.registered {
		log.Warningf("Ignore launch task because the executor is disconnected\n")
		k.sendStatusUpdate(taskInfo.GetTaskId(),
			mesos.TaskState_TASK_FAILED, "Executor not registered yet")
		return
	}

	k.lock.Lock()
	defer k.lock.Unlock()

	taskId := taskInfo.GetTaskId().GetValue()
	if _, found := k.tasks[taskId]; found {
		log.Warningf("Task already launched\n")
		// Not to send back TASK_RUNNING here, because
		// may be duplicated messages or duplicated task id.
		return
	}
	// Get the bound pod spec from the taskInfo.
	var pod api.BoundPod
	if err := yaml.Unmarshal(taskInfo.GetData(), &pod); err != nil {
		log.Warningf("Failed to extract yaml data from the taskInfo.data %v\n", err)
		k.sendStatusUpdate(taskInfo.GetTaskId(),
			mesos.TaskState_TASK_FAILED, "Failed to extract yaml data")
		return
	}

	podFullName := kubelet.GetPodFullName(&api.BoundPod{
		ObjectMeta: api.ObjectMeta{
			Name:        pod.Name,
			Namespace:   pod.Namespace,
			Annotations: map[string]string{kubelet.ConfigSourceAnnotationKey: k.sourcename},
		},
	})

	// Add the task.
	k.tasks[taskId] = &kuberTask{
		mesosTaskInfo: taskInfo,
		podName:       podFullName,
	}
	k.pods[podFullName] = &pod

	getPidInfo := func(name string) (api.PodInfo, error) {
		return k.kl.GetPodInfo(name, "")
	}

	// TODO(nnielsen): Fail if container is already running.
	// TODO(nnielsen) Checkpoint pods.

	// Send the pod updates to the channel.
	update := kubelet.PodUpdate{Op: kubelet.SET}
	for _, p := range k.pods {
		update.Pods = append(update.Pods, *p)
	}
	k.updateChan <- update

	// Delay reporting 'task running' until container is up.
	go func() {
		expires := time.Now().Add(launchGracePeriod)
		for {
			now := time.Now()
			if now.After(expires) {
				log.Warningf("Launch expired grace period of '%v'", launchGracePeriod)
				break
			}

			// We need to poll the kublet for the pod state, as
			// there is no existing event / push model for this.
			time.Sleep(containerPollTime)

			info, err := getPidInfo(podFullName)
			if err != nil {
				continue
			}

			log.V(2).Infof("Found pod info: '%v'", info)
			data, err := json.Marshal(info)

			statusUpdate := &mesos.TaskStatus{
				TaskId:  taskInfo.GetTaskId(),
				State:   mesos.NewTaskState(mesos.TaskState_TASK_RUNNING),
				Message: proto.String("Pod '" + podFullName + "' is running"),
				Data:    data,
			}

			if err := k.driver.SendStatusUpdate(statusUpdate); err != nil {
				log.Warningf("Failed to send status update%v, %v", err)
			}

			// TODO(nnielsen): Monitor health of container and report if lost.
			// Should we also allow this to fail a couple of times before reporting lost?
			// What if the docker daemon is restarting and we can't connect, but it's
			// going to bring the pods back online as soon as it restarts?
			knownPod := func() bool {
				_, err := getPidInfo(podFullName)
				return err == nil
			}
			go func() {
				// Wait for the pod to go away and stop monitoring once it does
				// TODO (jdefelice) replace with an /events watch?
				for {
					time.Sleep(containerPollTime)
					if k.checkForLostPodTask(taskInfo, knownPod) {
						return
					}
				}
			}()

			return
		}

		k.lock.Lock()
		defer k.lock.Unlock()
		k.reportLostTask(taskId, "Task lost: launch failed")
	}()
}

// Intended to be executed as part of the pod monitoring loop, this fn (ultimately) checks with Docker
// whether the pod is running. It will only return false if the task is still registered and the pod is
// registered in Docker. Otherwise it returns true. If there's still a task record on file, but no pod
// in Docker, then we'll also send a TASK_LOST event.
func (k *KubernetesExecutor) checkForLostPodTask(taskInfo *mesos.TaskInfo, isKnownPod func() bool) bool {
	// TODO (jdefelice) don't send false alarms for deleted pods (KILLED tasks)
	k.lock.Lock()
	defer k.lock.Unlock()

	// TODO(jdef) we should really consider k.pods here, along with what docker is reporting, since the kubelet
	// may constantly attempt to instantiate a pod as long as it's in the pod state that we're handing to it.
	// otherwise, we're probably reporting a TASK_LOST prematurely. Should probably consult RestartPolicy to
	// determine appropriate behavior. Should probably also gracefully handle docker daemon restarts.
	taskId := taskInfo.GetTaskId().GetValue()
	if _, ok := k.tasks[taskId]; ok {
		if isKnownPod() {
			return false
		} else {
			log.Warningf("Detected lost pod, reporting lost task %v", taskId)
			k.reportLostTask(taskId, "Task lost: container disappeared")
		}
	} else {
		log.V(2).Infof("Task %v no longer registered, stop monitoring for lost pods", taskId)
	}
	return true
}

// KillTask is called when the executor receives a request to kill a task.
func (k *KubernetesExecutor) KillTask(driver mesos.ExecutorDriver, taskId *mesos.TaskID) {
	k.lock.Lock()
	defer k.lock.Unlock()

	log.Infof("Kill task %v\n", taskId)

	if !k.registered {
		//TODO(jdefelice) sent TASK_LOST here?
		log.Warningf("Ignore kill task because the executor is disconnected\n")
		return
	}

	k.killPodForTask(taskId.GetValue(), "Task killed")
}

// Kills the pod associated with the given task. Assumes that the caller is locking around
// pod and task storage.
func (k *KubernetesExecutor) killPodForTask(tid, reason string) {
	task, ok := k.tasks[tid]
	if !ok {
		log.Infof("Failed to kill task, unknown task %v\n", tid)
		return
	}
	delete(k.tasks, tid)

	pid := task.podName
	if _, found := k.pods[pid]; !found {
		log.Warningf("Cannot remove Unknown pod %v for task %v", pid, tid)
	} else {
		log.V(2).Infof("Deleting pod %v for task %v", pid, tid)
		delete(k.pods, pid)

		// Send the pod updates to the channel.
		update := kubelet.PodUpdate{Op: kubelet.SET}
		for _, p := range k.pods {
			update.Pods = append(update.Pods, *p)
		}
		k.updateChan <- update
	}
	// TODO(yifan): Check the result of the kill event.

	k.sendStatusUpdate(task.mesosTaskInfo.GetTaskId(), mesos.TaskState_TASK_KILLED, reason)
}

// Reports a lost task to the slave and updates internal task and pod tracking state.
// Assumes that the caller is locking around pod and task state.
func (k *KubernetesExecutor) reportLostTask(tid, reason string) {
	task, ok := k.tasks[tid]
	if !ok {
		log.Infof("Failed to report lost task, unknown task %v\n", tid)
		return
	}
	delete(k.tasks, tid)

	pid := task.podName
	if _, found := k.pods[pid]; !found {
		log.Warningf("Cannot remove Unknown pod %v for lost task %v", pid, tid)
	} else {
		log.V(2).Infof("Deleting pod %v for lost task %v", pid, tid)
		delete(k.pods, pid)

		// Send the pod updates to the channel.
		update := kubelet.PodUpdate{Op: kubelet.SET}
		for _, p := range k.pods {
			update.Pods = append(update.Pods, *p)
		}
		k.updateChan <- update
	}
	// TODO(yifan): Check the result of the kill event.

	k.sendStatusUpdate(task.mesosTaskInfo.GetTaskId(), mesos.TaskState_TASK_LOST, reason)
}

// FrameworkMessage is called when the framework sends some message to the executor
func (k *KubernetesExecutor) FrameworkMessage(driver mesos.ExecutorDriver, message string) {
	log.Infof("Receives message from framework %v\n", message)
	// TODO(yifan): Check for update message.
}

// Shutdown is called when the executor receives a shutdown request.
func (k *KubernetesExecutor) Shutdown(driver mesos.ExecutorDriver) {
	log.Infof("Shutdown the executor\n")

	k.lock.Lock()
	defer k.lock.Unlock()

	for tid, _ := range k.tasks {
		k.killPodForTask(tid, "Executor shutdown")
	}
}

// Error is called when some error happens.
func (k *KubernetesExecutor) Error(driver mesos.ExecutorDriver, message string) {
	log.Errorf("Executor error: %v\n", message)
}

func (k *KubernetesExecutor) sendStatusUpdate(taskId *mesos.TaskID, state mesos.TaskState, message string) {
	statusUpdate := &mesos.TaskStatus{
		TaskId:  taskId,
		State:   &state,
		Message: proto.String(message),
	}
	// TODO(yifan): Maybe try to resend again in the future.
	if err := k.driver.SendStatusUpdate(statusUpdate); err != nil {
		log.Warningf("Failed to send status update%v, %v", err)
	}
}
<commit_msg>do not communicate TASK_RUNNING status prematurely ("No network settings" error)<commit_after>package executor

import (
	"encoding/json"
	"sync"
	"time"

	"code.google.com/p/goprotobuf/proto"
	"github.com/GoogleCloudPlatform/kubernetes/pkg/api"
	"github.com/GoogleCloudPlatform/kubernetes/pkg/kubelet"
	log "github.com/golang/glog"
	"github.com/mesos/mesos-go/mesos"
	"gopkg.in/v1/yaml"
)

const (
	containerPollTime = 300 * time.Millisecond
	launchGracePeriod = 5 * time.Minute
)

type kuberTask struct {
	mesosTaskInfo *mesos.TaskInfo
	podName       string
}

// KubernetesExecutor is an mesos executor that runs pods
// in a minion machine.
type KubernetesExecutor struct {
	kl         *kubelet.Kubelet // the kubelet instance.
	updateChan chan<- interface{}
	driver     mesos.ExecutorDriver
	registered bool
	tasks      map[string]*kuberTask
	pods       map[string]*api.BoundPod
	lock       sync.RWMutex
	sourcename string
}

// New creates a new kubernetes executor.
func New(driver mesos.ExecutorDriver, kl *kubelet.Kubelet, ch chan<- interface{}, ns string) *KubernetesExecutor {
	return &KubernetesExecutor{
		kl:         kl,
		updateChan: ch,
		driver:     driver,
		registered: false,
		tasks:      make(map[string]*kuberTask),
		pods:       make(map[string]*api.BoundPod),
		sourcename: ns,
	}
}

// Registered is called when the executor is successfully registered with the slave.
func (k *KubernetesExecutor) Registered(driver mesos.ExecutorDriver,
	executorInfo *mesos.ExecutorInfo, frameworkInfo *mesos.FrameworkInfo, slaveInfo *mesos.SlaveInfo) {
	log.Infof("Executor %v of framework %v registered with slave %v\n",
		executorInfo, frameworkInfo, slaveInfo)
	k.registered = true
}

// Reregistered is called when the executor is successfully re-registered with the slave.
// This can happen when the slave fails over.
func (k *KubernetesExecutor) Reregistered(driver mesos.ExecutorDriver, slaveInfo *mesos.SlaveInfo) {
	log.Infof("Reregistered with slave %v\n", slaveInfo)
	k.registered = true
}

// Disconnected is called when the executor is disconnected with the slave.
func (k *KubernetesExecutor) Disconnected(driver mesos.ExecutorDriver) {
	log.Infof("Slave is disconnected\n")
	k.registered = false
}

// LaunchTask is called when the executor receives a request to launch a task.
func (k *KubernetesExecutor) LaunchTask(driver mesos.ExecutorDriver, taskInfo *mesos.TaskInfo) {
	log.Infof("Launch task %v\n", taskInfo)

	if !k.registered {
		log.Warningf("Ignore launch task because the executor is disconnected\n")
		k.sendStatusUpdate(taskInfo.GetTaskId(),
			mesos.TaskState_TASK_FAILED, "Executor not registered yet")
		return
	}

	k.lock.Lock()
	defer k.lock.Unlock()

	taskId := taskInfo.GetTaskId().GetValue()
	if _, found := k.tasks[taskId]; found {
		log.Warningf("Task already launched\n")
		// Not to send back TASK_RUNNING here, because
		// may be duplicated messages or duplicated task id.
		return
	}
	// Get the bound pod spec from the taskInfo.
	var pod api.BoundPod
	if err := yaml.Unmarshal(taskInfo.GetData(), &pod); err != nil {
		log.Warningf("Failed to extract yaml data from the taskInfo.data %v\n", err)
		k.sendStatusUpdate(taskInfo.GetTaskId(),
			mesos.TaskState_TASK_FAILED, "Failed to extract yaml data")
		return
	}

	podFullName := kubelet.GetPodFullName(&api.BoundPod{
		ObjectMeta: api.ObjectMeta{
			Name:        pod.Name,
			Namespace:   pod.Namespace,
			Annotations: map[string]string{kubelet.ConfigSourceAnnotationKey: k.sourcename},
		},
	})

	// Add the task.
	k.tasks[taskId] = &kuberTask{
		mesosTaskInfo: taskInfo,
		podName:       podFullName,
	}
	k.pods[podFullName] = &pod

	getPidInfo := func(name string) (api.PodInfo, error) {
		return k.kl.GetPodInfo(name, "")
	}

	// TODO(nnielsen): Fail if container is already running.
	// TODO(nnielsen) Checkpoint pods.

	// Send the pod updates to the channel.
	update := kubelet.PodUpdate{Op: kubelet.SET}
	for _, p := range k.pods {
		update.Pods = append(update.Pods, *p)
	}
	k.updateChan <- update

	// Delay reporting 'task running' until container is up.
	go func() {
		expires := time.Now().Add(launchGracePeriod)
		for {
			now := time.Now()
			if now.After(expires) {
				log.Warningf("Launch expired grace period of '%v'", launchGracePeriod)
				break
			}

			// We need to poll the kublet for the pod state, as
			// there is no existing event / push model for this.
			time.Sleep(containerPollTime)

			info, err := getPidInfo(podFullName)
			if err != nil {
				continue
			}

			// avoid sending back a running status while pod networking is down
			if podnet, ok := info["net"]; !ok || podnet.State.Running == nil {
				continue
			}

			log.V(2).Infof("Found pod info: '%v'", info)
			data, err := json.Marshal(info)

			statusUpdate := &mesos.TaskStatus{
				TaskId:  taskInfo.GetTaskId(),
				State:   mesos.NewTaskState(mesos.TaskState_TASK_RUNNING),
				Message: proto.String("Pod '" + podFullName + "' is running"),
				Data:    data,
			}

			if err := k.driver.SendStatusUpdate(statusUpdate); err != nil {
				log.Warningf("Failed to send status update%v, %v", err)
			}

			// TODO(nnielsen): Monitor health of container and report if lost.
			// Should we also allow this to fail a couple of times before reporting lost?
			// What if the docker daemon is restarting and we can't connect, but it's
			// going to bring the pods back online as soon as it restarts?
			knownPod := func() bool {
				_, err := getPidInfo(podFullName)
				return err == nil
			}
			go func() {
				// Wait for the pod to go away and stop monitoring once it does
				// TODO (jdefelice) replace with an /events watch?
				for {
					time.Sleep(containerPollTime)
					if k.checkForLostPodTask(taskInfo, knownPod) {
						return
					}
				}
			}()

			return
		}

		k.lock.Lock()
		defer k.lock.Unlock()
		k.reportLostTask(taskId, "Task lost: launch failed")
	}()
}

// Intended to be executed as part of the pod monitoring loop, this fn (ultimately) checks with Docker
// whether the pod is running. It will only return false if the task is still registered and the pod is
// registered in Docker. Otherwise it returns true. If there's still a task record on file, but no pod
// in Docker, then we'll also send a TASK_LOST event.
func (k *KubernetesExecutor) checkForLostPodTask(taskInfo *mesos.TaskInfo, isKnownPod func() bool) bool {
	// TODO (jdefelice) don't send false alarms for deleted pods (KILLED tasks)
	k.lock.Lock()
	defer k.lock.Unlock()

	// TODO(jdef) we should really consider k.pods here, along with what docker is reporting, since the kubelet
	// may constantly attempt to instantiate a pod as long as it's in the pod state that we're handing to it.
	// otherwise, we're probably reporting a TASK_LOST prematurely. Should probably consult RestartPolicy to
	// determine appropriate behavior. Should probably also gracefully handle docker daemon restarts.
	taskId := taskInfo.GetTaskId().GetValue()
	if _, ok := k.tasks[taskId]; ok {
		if isKnownPod() {
			return false
		} else {
			log.Warningf("Detected lost pod, reporting lost task %v", taskId)
			k.reportLostTask(taskId, "Task lost: container disappeared")
		}
	} else {
		log.V(2).Infof("Task %v no longer registered, stop monitoring for lost pods", taskId)
	}
	return true
}

// KillTask is called when the executor receives a request to kill a task.
func (k *KubernetesExecutor) KillTask(driver mesos.ExecutorDriver, taskId *mesos.TaskID) {
	k.lock.Lock()
	defer k.lock.Unlock()

	log.Infof("Kill task %v\n", taskId)

	if !k.registered {
		//TODO(jdefelice) sent TASK_LOST here?
		log.Warningf("Ignore kill task because the executor is disconnected\n")
		return
	}

	k.killPodForTask(taskId.GetValue(), "Task killed")
}

// Kills the pod associated with the given task. Assumes that the caller is locking around
// pod and task storage.
func (k *KubernetesExecutor) killPodForTask(tid, reason string) {
	task, ok := k.tasks[tid]
	if !ok {
		log.Infof("Failed to kill task, unknown task %v\n", tid)
		return
	}
	delete(k.tasks, tid)

	pid := task.podName
	if _, found := k.pods[pid]; !found {
		log.Warningf("Cannot remove Unknown pod %v for task %v", pid, tid)
	} else {
		log.V(2).Infof("Deleting pod %v for task %v", pid, tid)
		delete(k.pods, pid)

		// Send the pod updates to the channel.
		update := kubelet.PodUpdate{Op: kubelet.SET}
		for _, p := range k.pods {
			update.Pods = append(update.Pods, *p)
		}
		k.updateChan <- update
	}
	// TODO(yifan): Check the result of the kill event.

	k.sendStatusUpdate(task.mesosTaskInfo.GetTaskId(), mesos.TaskState_TASK_KILLED, reason)
}

// Reports a lost task to the slave and updates internal task and pod tracking state.
// Assumes that the caller is locking around pod and task state.
func (k *KubernetesExecutor) reportLostTask(tid, reason string) {
	task, ok := k.tasks[tid]
	if !ok {
		log.Infof("Failed to report lost task, unknown task %v\n", tid)
		return
	}
	delete(k.tasks, tid)

	pid := task.podName
	if _, found := k.pods[pid]; !found {
		log.Warningf("Cannot remove Unknown pod %v for lost task %v", pid, tid)
	} else {
		log.V(2).Infof("Deleting pod %v for lost task %v", pid, tid)
		delete(k.pods, pid)

		// Send the pod updates to the channel.
		update := kubelet.PodUpdate{Op: kubelet.SET}
		for _, p := range k.pods {
			update.Pods = append(update.Pods, *p)
		}
		k.updateChan <- update
	}
	// TODO(yifan): Check the result of the kill event.

	k.sendStatusUpdate(task.mesosTaskInfo.GetTaskId(), mesos.TaskState_TASK_LOST, reason)
}

// FrameworkMessage is called when the framework sends some message to the executor
func (k *KubernetesExecutor) FrameworkMessage(driver mesos.ExecutorDriver, message string) {
	log.Infof("Receives message from framework %v\n", message)
	// TODO(yifan): Check for update message.
}

// Shutdown is called when the executor receives a shutdown request.
func (k *KubernetesExecutor) Shutdown(driver mesos.ExecutorDriver) {
	log.Infof("Shutdown the executor\n")

	k.lock.Lock()
	defer k.lock.Unlock()

	for tid, _ := range k.tasks {
		k.killPodForTask(tid, "Executor shutdown")
	}
}

// Error is called when some error happens.
func (k *KubernetesExecutor) Error(driver mesos.ExecutorDriver, message string) {
	log.Errorf("Executor error: %v\n", message)
}

func (k *KubernetesExecutor) sendStatusUpdate(taskId *mesos.TaskID, state mesos.TaskState, message string) {
	statusUpdate := &mesos.TaskStatus{
		TaskId:  taskId,
		State:   &state,
		Message: proto.String(message),
	}
	// TODO(yifan): Maybe try to resend again in the future.
	if err := k.driver.SendStatusUpdate(statusUpdate); err != nil {
		log.Warningf("Failed to send status update%v, %v", err)
	}
}
<|endoftext|>


<commit_before>package main

import (
	"encoding/json"
	"fmt"
	"regexp"
	"sync"
	"time"

	pubsub "google.golang.org/api/pubsub/v1"

	log "github.com/Sirupsen/logrus"
)

type (
	JobSustainerConfig struct {
		Delay    float64 `json:"delay,omitempty"`
		Interval float64 `json:"interval,omitempty"`
	}

	JobMessageStatus uint8

	JobMessage struct {
		sub    string
		raw    *pubsub.ReceivedMessage
		config *JobSustainerConfig
		puller Puller
		status JobMessageStatus
		mux    sync.Mutex
	}
)

const (
	running JobMessageStatus = iota
	done
	acked
)

func (m *JobMessage) Validate() error {
	if m.MessageId() == "" {
		return &InvalidJobError{msg: "no MessageId is given"}
	}
	return nil
}

func (m *JobMessage) MessageId() string {
	return m.raw.Message.MessageId
}

func (m *JobMessage) DownloadFiles() interface{} {
	str, ok := m.raw.Message.Attributes["download_files"]
	if !ok {
		return nil
	}
	return m.parseJson(str)
}

func (m *JobMessage) parseJson(str string) interface{} {
	matched, err := regexp.MatchString(`\A\[.*\]\z|\A\{.*\}\z|`, str)
	if err != nil {
		return str
	}
	if !matched {
		return str
	}
	var dest interface{}
	err = json.Unmarshal([]byte(str), &dest)
	if err != nil {
		return str
	}
	return dest
}

func (m *JobMessage) Ack() error {
	m.mux.Lock()
	defer m.mux.Unlock()

	logAttrs := log.Fields{"job_message_id": m.MessageId(), "ack_id": m.raw.AckId}
	log.WithFields(logAttrs).Debugln("Sending ACK")

	_, err := m.puller.Acknowledge(m.sub, m.raw.AckId)
	if err != nil {
		logAttrs["raw"] = fmt.Sprintf("%v", m.raw)
		logAttrs["error"] = err
		log.WithFields(logAttrs).Errorln("Failed to acknowledge")
		return err
	}

	logAttrs["status"] = m.status
	log.WithFields(logAttrs).Debugln("Updating status to acked")

	m.status = acked
	return nil
}

func (m *JobMessage) Nack() error {
	m.mux.Lock()
	defer m.mux.Unlock()

	logAttrs := log.Fields{"job_message_id": m.MessageId(), "ack_id": m.raw.AckId}
	log.WithFields(logAttrs).Debugln("Sending NACK")

	_, err := m.puller.ModifyAckDeadline(m.sub, []string{m.raw.AckId}, 0)
	if err != nil {
		logAttrs["raw"] = fmt.Sprintf("%v", m.raw)
		logAttrs["error"] = err
		log.WithFields(logAttrs).Errorln("Failed to send ModifyAckDeadline as a NACK")
		return err
	}

	logAttrs["status"] = m.status
	log.WithFields(logAttrs).Debugln("Updating status to done")

	m.status = done
	return nil
}

func (m *JobMessage) Done() {
	logAttrs := log.Fields{"job_message_id": m.MessageId(), "status": m.status}
	log.WithFields(logAttrs).Debugln("Done()")
	if m.status == running {
		m.status = done
	}
}

func (m *JobMessage) running() bool {
	return m.status == running
}

func (m *JobMessage) sendMADPeriodically(notification *ProgressNotification) error {
	log.Printf("sendMADPeriodically start\n")
	for {
		nextLimit := time.Now().Add(time.Duration(m.config.Interval) * time.Second)
		err := m.waitAndSendMAD(notification, nextLimit)
		if err != nil {
			log.WithFields(log.Fields{"error": err}).Errorln("Error in sendMADPeriodically")
			return err
		}
		if !m.running() {
			log.Debugln("sendMADPeriodically return")
			return nil
		}
	}
	// return nil
}

func (m *JobMessage) waitAndSendMAD(notification *ProgressNotification, nextLimit time.Time) error {
	log.Debugln("waitAndSendMAD starting")
	ticker := time.NewTicker(100 * time.Millisecond)
	for now := range ticker.C {
		if !m.running() {
			log.Debugln("waitAndSendMAD ticker stopping")
			ticker.Stop()
			log.Debugln("waitAndSendMAD ticker stopped")
			return nil
		}
		if now.After(nextLimit) {
			log.Debugln("waitAndSendMAD nextLimit passed")
			ticker.Stop()
			log.Debugln("waitAndSendMAD ticker stopped")
			break
		}
	}

	log.Debugln("waitAndSendMAD m.mux locking")

	m.mux.Lock()
	defer m.mux.Unlock()

	logAttrs := log.Fields{"status": m.status}
	log.WithFields(logAttrs).Debugln("waitAndSendMAD")

	// Don't send MAD after sending ACK
	if m.status == acked {
		log.WithFields(logAttrs).Infoln("waitAndSendMAD already acked")
		return nil
	}

	log.WithFields(logAttrs).Debugln("waitAndSendMAD sending ModifyAckDeadline")
	_, err := m.puller.ModifyAckDeadline(m.sub, []string{m.raw.AckId}, int64(m.config.Delay))
	if err != nil {
		logAttrs["error"] = err
		logAttrs["subscription"] = m.sub
		logAttrs["AckId"] = m.raw.AckId
		logAttrs["Delay"] = m.config.Delay
		log.WithFields(logAttrs).Errorln("waitAndSendMAD ModifyAckDeadline")
		msg := fmt.Sprintf("Failed modifyAckDeadline %v, %v, %v cause of %v\n", m.sub, m.raw.AckId, m.config.Delay, err)
		log.WithFields(logAttrs).Fatalf(msg)
		notification.notifyProgress(m.MessageId(), WORKING, false, "error", msg)
	}
	return nil
}
<commit_msg>:+1: Add String method to JobMessageStatus<commit_after>package main

import (
	"encoding/json"
	"fmt"
	"regexp"
	"sync"
	"time"

	pubsub "google.golang.org/api/pubsub/v1"

	log "github.com/Sirupsen/logrus"
)

type JobMessageStatus uint8

const (
	running JobMessageStatus = iota
	done
	acked
)

func (s JobMessageStatus) String() string {
	switch s {
	case running:
		return "running"
	case done:
		return "done"
	case acked:
		return "acked"
	default:
		return fmt.Sprintf("JobMessageStatus[Unknown:%v]", uint8(s))
	}
}

type (
	JobSustainerConfig struct {
		Delay    float64 `json:"delay,omitempty"`
		Interval float64 `json:"interval,omitempty"`
	}

	JobMessage struct {
		sub    string
		raw    *pubsub.ReceivedMessage
		config *JobSustainerConfig
		puller Puller
		status JobMessageStatus
		mux    sync.Mutex
	}
)

func (m *JobMessage) Validate() error {
	if m.MessageId() == "" {
		return &InvalidJobError{msg: "no MessageId is given"}
	}
	return nil
}

func (m *JobMessage) MessageId() string {
	return m.raw.Message.MessageId
}

func (m *JobMessage) DownloadFiles() interface{} {
	str, ok := m.raw.Message.Attributes["download_files"]
	if !ok {
		return nil
	}
	return m.parseJson(str)
}

func (m *JobMessage) parseJson(str string) interface{} {
	matched, err := regexp.MatchString(`\A\[.*\]\z|\A\{.*\}\z|`, str)
	if err != nil {
		return str
	}
	if !matched {
		return str
	}
	var dest interface{}
	err = json.Unmarshal([]byte(str), &dest)
	if err != nil {
		return str
	}
	return dest
}

func (m *JobMessage) Ack() error {
	m.mux.Lock()
	defer m.mux.Unlock()

	logAttrs := log.Fields{"job_message_id": m.MessageId(), "ack_id": m.raw.AckId}
	log.WithFields(logAttrs).Debugln("Sending ACK")

	_, err := m.puller.Acknowledge(m.sub, m.raw.AckId)
	if err != nil {
		logAttrs["raw"] = fmt.Sprintf("%v", m.raw)
		logAttrs["error"] = err
		log.WithFields(logAttrs).Errorln("Failed to acknowledge")
		return err
	}

	logAttrs["status"] = m.status
	log.WithFields(logAttrs).Debugln("Updating status to acked")

	m.status = acked
	return nil
}

func (m *JobMessage) Nack() error {
	m.mux.Lock()
	defer m.mux.Unlock()

	logAttrs := log.Fields{"job_message_id": m.MessageId(), "ack_id": m.raw.AckId}
	log.WithFields(logAttrs).Debugln("Sending NACK")

	_, err := m.puller.ModifyAckDeadline(m.sub, []string{m.raw.AckId}, 0)
	if err != nil {
		logAttrs["raw"] = fmt.Sprintf("%v", m.raw)
		logAttrs["error"] = err
		log.WithFields(logAttrs).Errorln("Failed to send ModifyAckDeadline as a NACK")
		return err
	}

	logAttrs["status"] = m.status
	log.WithFields(logAttrs).Debugln("Updating status to done")

	m.status = done
	return nil
}

func (m *JobMessage) Done() {
	logAttrs := log.Fields{"job_message_id": m.MessageId(), "status": m.status}
	log.WithFields(logAttrs).Debugln("Done()")
	if m.status == running {
		m.status = done
	}
}

func (m *JobMessage) running() bool {
	return m.status == running
}

func (m *JobMessage) sendMADPeriodically(notification *ProgressNotification) error {
	log.Printf("sendMADPeriodically start\n")
	for {
		nextLimit := time.Now().Add(time.Duration(m.config.Interval) * time.Second)
		err := m.waitAndSendMAD(notification, nextLimit)
		if err != nil {
			log.WithFields(log.Fields{"error": err}).Errorln("Error in sendMADPeriodically")
			return err
		}
		if !m.running() {
			log.Debugln("sendMADPeriodically return")
			return nil
		}
	}
	// return nil
}

func (m *JobMessage) waitAndSendMAD(notification *ProgressNotification, nextLimit time.Time) error {
	log.Debugln("waitAndSendMAD starting")
	ticker := time.NewTicker(100 * time.Millisecond)
	for now := range ticker.C {
		if !m.running() {
			log.Debugln("waitAndSendMAD ticker stopping")
			ticker.Stop()
			log.Debugln("waitAndSendMAD ticker stopped")
			return nil
		}
		if now.After(nextLimit) {
			log.Debugln("waitAndSendMAD nextLimit passed")
			ticker.Stop()
			log.Debugln("waitAndSendMAD ticker stopped")
			break
		}
	}

	log.Debugln("waitAndSendMAD m.mux locking")

	m.mux.Lock()
	defer m.mux.Unlock()

	logAttrs := log.Fields{"status": m.status}
	log.WithFields(logAttrs).Debugln("waitAndSendMAD")

	// Don't send MAD after sending ACK
	if m.status == acked {
		log.WithFields(logAttrs).Infoln("waitAndSendMAD already acked")
		return nil
	}

	log.WithFields(logAttrs).Debugln("waitAndSendMAD sending ModifyAckDeadline")
	_, err := m.puller.ModifyAckDeadline(m.sub, []string{m.raw.AckId}, int64(m.config.Delay))
	if err != nil {
		logAttrs["error"] = err
		logAttrs["subscription"] = m.sub
		logAttrs["AckId"] = m.raw.AckId
		logAttrs["Delay"] = m.config.Delay
		log.WithFields(logAttrs).Errorln("waitAndSendMAD ModifyAckDeadline")
		msg := fmt.Sprintf("Failed modifyAckDeadline %v, %v, %v cause of %v\n", m.sub, m.raw.AckId, m.config.Delay, err)
		log.WithFields(logAttrs).Fatalf(msg)
		notification.notifyProgress(m.MessageId(), WORKING, false, "error", msg)
	}
	return nil
}
<|endoftext|>


<commit_before>package nscatools

import (
	"bytes"
	"encoding/binary"
	"fmt"
	"hash/crc32"
	"io"
	"time"
)

// DataPacket stores the data received for the client-server communication
type DataPacket struct {
	Version      int16
	Crc          uint32
	Timestamp    uint32
	State        int16
	HostName     string
	Service      string
	PluginOutput string
	Ipkt         *InitPacket
	Password     []byte
	Encryption   int
}

// NewDataPacket initializes a new blank data packet
func NewDataPacket(encryption int, password []byte, ipkt *InitPacket) *DataPacket {
	packet := DataPacket{
		Version:      3,
		Crc:          0,
		Timestamp:    uint32(time.Now().Unix()),
		State:        StateUnknown,
		HostName:     "",
		Service:      "",
		PluginOutput: "",
		Ipkt:         ipkt,
		Password:     password,
		Encryption:   encryption,
	}
	return &packet
}

// CalculateCrc returns the Crc of a packet ready to be sent over the network,
// ignoring the Crc data part of it as it's done in the original nsca code
func (p *DataPacket) CalculateCrc(buffer []byte) uint32 {
	crcdPacket := make([]byte, 4304)
	copy(crcdPacket, buffer[0:4])
	copy(crcdPacket[8:], buffer[8:])
	return crc32.ChecksumIEEE(crcdPacket)
}

// Read gets the data packet and populates the attributes of the DataPacket
// according. When encountering an error, it returns the error and don't process
// further.
func (p *DataPacket) Read(conn io.Reader) error {
	// We need to read the full packet 1st to check the crc and decrypt it too
	fullPacket := make([]byte, 4304)
	if _, err := io.ReadFull(conn, fullPacket); err != nil {
		return err
	}

	if err := p.Decrypt(fullPacket); err != nil {
		return err
	}

	p.Crc = binary.BigEndian.Uint32(fullPacket[4:8])
	if crc32 := p.CalculateCrc(fullPacket); p.Crc != crc32 {
		return fmt.Errorf("Dropping packet with invalid CRC32 - possibly due to client using wrong password or crypto algorithm?")
	}

	p.Timestamp = binary.BigEndian.Uint32(fullPacket[8:12])
	// MaxPacketAge <= 0 means that we don't check it
	if MaxPacketAge > 0 {
		if p.Timestamp > (p.Ipkt.Timestamp+MaxPacketAge) || p.Timestamp < (p.Ipkt.Timestamp-MaxPacketAge) {
			return fmt.Errorf("Dropping packet with stale timestamp - Max age difference is %d seconds", MaxPacketAge)
		}
	}

	sep := []byte("\x00") // sep is used to extract only the useful string
	p.Version = int16(binary.BigEndian.Uint16(fullPacket[0:2]))
	p.State = int16(binary.BigEndian.Uint16(fullPacket[12:14]))
	p.HostName = string(bytes.Split(fullPacket[14:78], sep)[0])
	p.Service = string(bytes.Split(fullPacket[78:206], sep)[0])
	p.PluginOutput = string(bytes.Split(fullPacket[206:], sep)[0])

	return nil
}

// Write generates the buffer to write to the writer based on the populated
// fields of the DataPacket instance encrypts it if needed and send it to the
// writer.
// When encountering an error, it returns the error and don't process
// further.
func (p *DataPacket) Write(w io.Writer) error {

	// Build network packet
	packet := new(bytes.Buffer)
	binary.Write(packet, binary.BigEndian, p.Version)
	binary.Write(packet, binary.BigEndian, make([]byte, 6))
	binary.Write(packet, binary.BigEndian, p.Timestamp)
	binary.Write(packet, binary.BigEndian, p.State)
	h := make([]byte, 64)
	copy(h, p.HostName)
	binary.Write(packet, binary.BigEndian, h)
	s := make([]byte, 128)
	copy(s, p.Service)
	binary.Write(packet, binary.BigEndian, s)
	o := make([]byte, 64)
	copy(o, p.PluginOutput)
	binary.Write(packet, binary.BigEndian, o)
	binary.Write(packet, binary.BigEndian, make([]byte, 2))

	buf := make([]byte, 4304)
	copy(buf, packet.Bytes())

	// Calculate the Crc
	binary.BigEndian.PutUint32(buf[4:8], p.CalculateCrc(buf))

	// Encrypt
	if err := p.Encrypt(buf); err != nil {
		return err
	}

	// Write + consistency check
	if n, err := w.Write(buf); err != nil || n != len(buf) {
		return fmt.Errorf("%d bytes written, expecting %d. Error: %s", n, len(buf), err)
	}

	return nil
}

// Performs a XOR operation on a buffer using the initialization vector and the
// password.
func (p *DataPacket) xor(buffer []byte) {
	bufferSize := len(buffer)
	ivSize := len(p.Ipkt.Iv)
	pwdSize := len(p.Password)
	// Rotating over the initialization vector of the connection and the password
	for y := 0; y < bufferSize; y++ {
		buffer[y] ^= p.Ipkt.Iv[y%ivSize] ^ p.Password[y%pwdSize]
	}
}

// setAlgo translates the encryption to an mcrypt-understandable from algorithm name
func (p *DataPacket) setAlgo() string {
	var algo string
	switch p.Encryption {
	case EncryptDES:
		algo = "des"
	case Encrypt3DES:
		algo = "tripledes"
	case EncryptCAST128:
		algo = "cast-128"
	case EncryptCAST256:
		algo = "cast-256"
	case EncryptXTEA:
		algo = "xtea"
	case Encrypt3WAY:
		algo = "threeway"
	case EncryptBLOWFISH:
		algo = "blowfish"
	case EncryptTWOFISH:
		algo = "twofish"
	case EncryptLOKI97:
		algo = "loki97"
	case EncryptRC2:
		algo = "rc2"
	case EncryptARCFOUR:
		algo = "arcfour"
	case EncryptRIJNDAEL128:
		algo = "rijndael-128"
	case EncryptRIJNDAEL192:
		algo = "rijndael-192"
	case EncryptRIJNDAEL256:
		algo = "rijndael-256"
	case EncryptWAKE:
		algo = "wake"
	case EncryptSERPENT:
		algo = "serpent"
	case EncryptENIGMA:
		algo = "enigma"
	case EncryptGOST:
		algo = "gost"
	case EncryptSAFER64:
		algo = "safer-sk64"
	case EncryptSAFER128:
		algo = "safer-sk128"
	case EncryptSAFERPLUS:
		algo = "saferplus"
	default:
		algo = "Unknown"
	}
	return algo
}

// Decrypt decrypts a buffer
func (p *DataPacket) Decrypt(buffer []byte) error {
	var (
		algo string
		err  error
	)

	switch p.Encryption {
	case EncryptNone: // Just don't do anything
	case EncryptXOR:
		p.xor(buffer)
	case EncryptIDEA: // Unsupported in standard NSCA
		err = fmt.Errorf("Unimplemented encryption algorithm")
	case EncryptRC6: // Unsupported in standard NSCA
		err = fmt.Errorf("Unimplemented encryption algorithm")
	case EncryptMARS: // Unsupported in standard NSCA
		err = fmt.Errorf("Unimplemented encryption algorithm")
	case EncryptPANAMA: // Unsupported in standard NSCA
		err = fmt.Errorf("Unimplemented encryption algorithm")
	default:
		algo = p.setAlgo()
		if algo == "Unknown" {
			err = fmt.Errorf("%d is an unrecognized encryption integer", p.Encryption)
		}
	}
	if err != nil {
		return err
	}
	if algo != "" {
		err = MCrypt(algo, buffer, p.Password, p.Ipkt.Iv, true)
	}
	return err
}

// Encrypt encrypts a buffer
func (p *DataPacket) Encrypt(buffer []byte) error {
	var (
		algo string
		err  error
	)

	switch p.Encryption {
	case EncryptNone:
	case EncryptXOR:
		p.xor(buffer)
	case EncryptIDEA: // Unsupported in standard NSCA
		err = fmt.Errorf("Unimplemented encryption algorithm")
	case EncryptRC6: // Unsupported in standard NSCA
		err = fmt.Errorf("Unimplemented encryption algorithm")
	case EncryptMARS: // Unsupported in standard NSCA
		err = fmt.Errorf("Unimplemented encryption algorithm")
	case EncryptPANAMA: // Unsupported in standard NSCA
		err = fmt.Errorf("Unimplemented encryption algorithm")
	default:
		algo = p.setAlgo()
		if algo == "Unknown" {
			err = fmt.Errorf("%d is an unrecognized encryption integer", p.Encryption)
		}
	}
	if err != nil {
		return err
	}
	if algo != "" {
		err = MCrypt(algo, buffer, p.Password, p.Ipkt.Iv, false)
	}
	return err
}

// Encrypt* are the encryptions supported by the standard NSCA configuration
const (
	EncryptNone        = iota // no encryption
	EncryptXOR                // Simple XOR  (No security, just obfuscation, but very fast)
	EncryptDES                // DES
	Encrypt3DES               // 3DES or Triple DES
	EncryptCAST128            // CAST-128
	EncryptCAST256            // CAST-256
	EncryptXTEA               // xTEA
	Encrypt3WAY               // 3-WAY
	EncryptBLOWFISH           // SKIPJACK
	EncryptTWOFISH            // TWOFISH
	EncryptLOKI97             // LOKI97
	EncryptRC2                // RC2
	EncryptARCFOUR            // RC4
	EncryptRC6                // RC6 - Unsupported in standard NSCA
	EncryptRIJNDAEL128        // AES-128
	EncryptRIJNDAEL192        // AES-192
	EncryptRIJNDAEL256        // AES-256
	EncryptMARS               // MARS - Unsupported in standard NSCA
	EncryptPANAMA             // PANAMA - Unsupported in standard NSCA
	EncryptWAKE               // WAKE
	EncryptSERPENT            // SERPENT
	EncryptIDEA               // IDEA - Unsupported in standard NSCA
	EncryptENIGMA             // ENIGMA (Unix crypt)
	EncryptGOST               // GOST
	EncryptSAFER64            // SAFER-sk64
	EncryptSAFER128           // SAFER-sk128
	EncryptSAFERPLUS          // SAFER+
)

// State* are the states understood by NSCA
const (
	StateOK = iota
	StateWarning
	StateCritical
	StateUnknown
)

// MaxPacketAge is the number of seconds difference allowed between the
// initialization packet epoch and the epoch of the data packet received
const MaxPacketAge = 30
<commit_msg>Add short packets support<commit_after>package nscatools

import (
	"bytes"
	"encoding/binary"
	"fmt"
	"hash/crc32"
	"io"
	"io/ioutil"
	"time"
)

// DataPacket stores the data received for the client-server communication
type DataPacket struct {
	Version      int16
	Crc          uint32
	Timestamp    uint32
	State        int16
	HostName     string
	Service      string
	PluginOutput string
	Ipkt         *InitPacket
	Password     []byte
	Encryption   int
}

// NewDataPacket initializes a new blank data packet
func NewDataPacket(encryption int, password []byte, ipkt *InitPacket) *DataPacket {
	packet := DataPacket{
		Version:      3,
		Crc:          0,
		Timestamp:    uint32(time.Now().Unix()),
		State:        StateUnknown,
		HostName:     "",
		Service:      "",
		PluginOutput: "",
		Ipkt:         ipkt,
		Password:     password,
		Encryption:   encryption,
	}
	return &packet
}

// CalculateCrc returns the Crc of a packet ready to be sent over the network,
// ignoring the Crc data part of it as it's done in the original nsca code
func (p *DataPacket) CalculateCrc(buffer []byte) uint32 {
	crcdPacket := make([]byte, len(buffer))
	copy(crcdPacket, buffer[0:4])
	copy(crcdPacket[8:], buffer[8:])
	return crc32.ChecksumIEEE(crcdPacket)
}

// Read gets the data packet and populates the attributes of the DataPacket
// according. When encountering an error, it returns the error and don't process
// further.
func (p *DataPacket) Read(conn io.Reader) error {
	// We need to read the full packet 1st to check the crc and decrypt it too
	fullPacket, err := ioutil.ReadAll(conn)
	if err != nil {
		return err
	}

	if len(fullPacket) != ShortPacketLength && len(fullPacket) != LongPacketLength {
		return fmt.Errorf("Dropping packet with invalid size: %d", len(fullPacket))
	}

	if err := p.Decrypt(fullPacket); err != nil {
		return err
	}

	p.Crc = binary.BigEndian.Uint32(fullPacket[4:8])
	if crc32 := p.CalculateCrc(fullPacket); p.Crc != crc32 {
		return fmt.Errorf("Dropping packet with invalid CRC32 - possibly due to client using wrong password or crypto algorithm?")
	}

	p.Timestamp = binary.BigEndian.Uint32(fullPacket[8:12])
	// MaxPacketAge <= 0 means that we don't check it
	if MaxPacketAge > 0 {
		if p.Timestamp > (p.Ipkt.Timestamp+MaxPacketAge) || p.Timestamp < (p.Ipkt.Timestamp-MaxPacketAge) {
			return fmt.Errorf("Dropping packet with stale timestamp - Max age difference is %d seconds", MaxPacketAge)
		}
	}

	sep := []byte("\x00") // sep is used to extract only the useful string
	p.Version = int16(binary.BigEndian.Uint16(fullPacket[0:2]))
	p.State = int16(binary.BigEndian.Uint16(fullPacket[12:14]))
	p.HostName = string(bytes.Split(fullPacket[14:78], sep)[0])
	p.Service = string(bytes.Split(fullPacket[78:206], sep)[0])
	p.PluginOutput = string(bytes.Split(fullPacket[206:], sep)[0])

	return nil
}

// Write generates the buffer to write to the writer based on the populated
// fields of the DataPacket instance encrypts it if needed and send it to the
// writer.
// When encountering an error, it returns the error and don't process
// further.
func (p *DataPacket) Write(w io.Writer) error {

	// Build network packet
	packet := new(bytes.Buffer)
	binary.Write(packet, binary.BigEndian, p.Version)
	binary.Write(packet, binary.BigEndian, make([]byte, 6))
	binary.Write(packet, binary.BigEndian, p.Timestamp)
	binary.Write(packet, binary.BigEndian, p.State)
	h := make([]byte, 64)
	copy(h, p.HostName)
	binary.Write(packet, binary.BigEndian, h)
	s := make([]byte, 128)
	copy(s, p.Service)
	binary.Write(packet, binary.BigEndian, s)
	o := make([]byte, 64)
	copy(o, p.PluginOutput)
	binary.Write(packet, binary.BigEndian, o)
	binary.Write(packet, binary.BigEndian, make([]byte, 2))

	buf := make([]byte, LongPacketLength)
	copy(buf, packet.Bytes())

	// Calculate the Crc
	binary.BigEndian.PutUint32(buf[4:8], p.CalculateCrc(buf))

	// Encrypt
	if err := p.Encrypt(buf); err != nil {
		return err
	}

	// Write + consistency check
	if n, err := w.Write(buf); err != nil || n != len(buf) {
		return fmt.Errorf("%d bytes written, expecting %d. Error: %s", n, len(buf), err)
	}

	return nil
}

// Performs a XOR operation on a buffer using the initialization vector and the
// password.
func (p *DataPacket) xor(buffer []byte) {
	bufferSize := len(buffer)
	ivSize := len(p.Ipkt.Iv)
	pwdSize := len(p.Password)
	// Rotating over the initialization vector of the connection and the password
	for y := 0; y < bufferSize; y++ {
		buffer[y] ^= p.Ipkt.Iv[y%ivSize] ^ p.Password[y%pwdSize]
	}
}

// setAlgo translates the encryption to an mcrypt-understandable from algorithm name
func (p *DataPacket) setAlgo() string {
	var algo string
	switch p.Encryption {
	case EncryptDES:
		algo = "des"
	case Encrypt3DES:
		algo = "tripledes"
	case EncryptCAST128:
		algo = "cast-128"
	case EncryptCAST256:
		algo = "cast-256"
	case EncryptXTEA:
		algo = "xtea"
	case Encrypt3WAY:
		algo = "threeway"
	case EncryptBLOWFISH:
		algo = "blowfish"
	case EncryptTWOFISH:
		algo = "twofish"
	case EncryptLOKI97:
		algo = "loki97"
	case EncryptRC2:
		algo = "rc2"
	case EncryptARCFOUR:
		algo = "arcfour"
	case EncryptRIJNDAEL128:
		algo = "rijndael-128"
	case EncryptRIJNDAEL192:
		algo = "rijndael-192"
	case EncryptRIJNDAEL256:
		algo = "rijndael-256"
	case EncryptWAKE:
		algo = "wake"
	case EncryptSERPENT:
		algo = "serpent"
	case EncryptENIGMA:
		algo = "enigma"
	case EncryptGOST:
		algo = "gost"
	case EncryptSAFER64:
		algo = "safer-sk64"
	case EncryptSAFER128:
		algo = "safer-sk128"
	case EncryptSAFERPLUS:
		algo = "saferplus"
	default:
		algo = "Unknown"
	}
	return algo
}

// Decrypt decrypts a buffer
func (p *DataPacket) Decrypt(buffer []byte) error {
	var (
		algo string
		err  error
	)

	switch p.Encryption {
	case EncryptNone: // Just don't do anything
	case EncryptXOR:
		p.xor(buffer)
	case EncryptIDEA: // Unsupported in standard NSCA
		err = fmt.Errorf("Unimplemented encryption algorithm")
	case EncryptRC6: // Unsupported in standard NSCA
		err = fmt.Errorf("Unimplemented encryption algorithm")
	case EncryptMARS: // Unsupported in standard NSCA
		err = fmt.Errorf("Unimplemented encryption algorithm")
	case EncryptPANAMA: // Unsupported in standard NSCA
		err = fmt.Errorf("Unimplemented encryption algorithm")
	default:
		algo = p.setAlgo()
		if algo == "Unknown" {
			err = fmt.Errorf("%d is an unrecognized encryption integer", p.Encryption)
		}
	}
	if err != nil {
		return err
	}
	if algo != "" {
		err = MCrypt(algo, buffer, p.Password, p.Ipkt.Iv, true)
	}
	return err
}

// Encrypt encrypts a buffer
func (p *DataPacket) Encrypt(buffer []byte) error {
	var (
		algo string
		err  error
	)

	switch p.Encryption {
	case EncryptNone:
	case EncryptXOR:
		p.xor(buffer)
	case EncryptIDEA: // Unsupported in standard NSCA
		err = fmt.Errorf("Unimplemented encryption algorithm")
	case EncryptRC6: // Unsupported in standard NSCA
		err = fmt.Errorf("Unimplemented encryption algorithm")
	case EncryptMARS: // Unsupported in standard NSCA
		err = fmt.Errorf("Unimplemented encryption algorithm")
	case EncryptPANAMA: // Unsupported in standard NSCA
		err = fmt.Errorf("Unimplemented encryption algorithm")
	default:
		algo = p.setAlgo()
		if algo == "Unknown" {
			err = fmt.Errorf("%d is an unrecognized encryption integer", p.Encryption)
		}
	}
	if err != nil {
		return err
	}
	if algo != "" {
		err = MCrypt(algo, buffer, p.Password, p.Ipkt.Iv, false)
	}
	return err
}

// Encrypt* are the encryptions supported by the standard NSCA configuration
const (
	EncryptNone        = iota // no encryption
	EncryptXOR                // Simple XOR  (No security, just obfuscation, but very fast)
	EncryptDES                // DES
	Encrypt3DES               // 3DES or Triple DES
	EncryptCAST128            // CAST-128
	EncryptCAST256            // CAST-256
	EncryptXTEA               // xTEA
	Encrypt3WAY               // 3-WAY
	EncryptBLOWFISH           // SKIPJACK
	EncryptTWOFISH            // TWOFISH
	EncryptLOKI97             // LOKI97
	EncryptRC2                // RC2
	EncryptARCFOUR            // RC4
	EncryptRC6                // RC6 - Unsupported in standard NSCA
	EncryptRIJNDAEL128        // AES-128
	EncryptRIJNDAEL192        // AES-192
	EncryptRIJNDAEL256        // AES-256
	EncryptMARS               // MARS - Unsupported in standard NSCA
	EncryptPANAMA             // PANAMA - Unsupported in standard NSCA
	EncryptWAKE               // WAKE
	EncryptSERPENT            // SERPENT
	EncryptIDEA               // IDEA - Unsupported in standard NSCA
	EncryptENIGMA             // ENIGMA (Unix crypt)
	EncryptGOST               // GOST
	EncryptSAFER64            // SAFER-sk64
	EncryptSAFER128           // SAFER-sk128
	EncryptSAFERPLUS          // SAFER+
)

// State* are the states understood by NSCA
const (
	StateOK = iota
	StateWarning
	StateCritical
	StateUnknown
)

// MaxPacketAge is the number of seconds difference allowed between the
// initialization packet epoch and the epoch of the data packet received
const MaxPacketAge = 30

const ShortPacketLength = 720
const LongPacketLength = 4304
<|endoftext|>