<commit_before>package encoders

import (
	"errors"
	"strconv"

	"code.cloudfoundry.org/lager"

	"github.com/3d0c/gmf"
	"github.com/snickers/snickers/db"
	"github.com/snickers/snickers/types"
)

func addStream(job types.Job, codecName string, oc *gmf.FmtCtx, ist *gmf.Stream) (int, int, error) {
	var cc *gmf.CodecCtx
	var ost *gmf.Stream

	codec, err := gmf.FindEncoder(codecName)
	if err != nil {
		return 0, 0, err
	}

	if ost = oc.NewStream(codec); ost == nil {
		return 0, 0, errors.New("unable to create stream in output context")
	}
	defer gmf.Release(ost)

	if cc = gmf.NewCodecCtx(codec); cc == nil {
		return 0, 0, errors.New("unable to create codec context")
	}
	defer gmf.Release(cc)

	// https://ffmpeg.org/pipermail/ffmpeg-devel/2008-January/046900.html
	if oc.IsGlobalHeader() {
		cc.SetFlag(gmf.CODEC_FLAG_GLOBAL_HEADER)
	}

	if codec.IsExperimental() {
		cc.SetStrictCompliance(gmf.FF_COMPLIANCE_EXPERIMENTAL)
	}

	if cc.Type() == gmf.AVMEDIA_TYPE_AUDIO {
		bitrate, err := strconv.Atoi(job.Preset.Audio.Bitrate)
		if err != nil {
			return 0, 0, err
		}

		gop, err := strconv.Atoi(job.Preset.Video.GopSize)
		if err != nil {
			return 0, 0, err
		}

		cc.SetBitRate(bitrate)
		cc.SetGopSize(gop)
		cc.SetSampleFmt(ist.CodecCtx().SampleFmt())
		cc.SetSampleRate(ist.CodecCtx().SampleRate())
		cc.SetChannels(ist.CodecCtx().Channels())
		cc.SelectChannelLayout()
		cc.SelectSampleRate()
	}

	if cc.Type() == gmf.AVMEDIA_TYPE_VIDEO {
		cc.SetTimeBase(gmf.AVR{Num: 1, Den: 25}) // what is this

		if job.Preset.Video.Codec == "h264" {
			profile := GetProfile(job)
			cc.SetProfile(profile)
		}

		width, height := GetResolution(job, ist.CodecCtx().Width(), ist.CodecCtx().Height())

		bitrate, err := strconv.Atoi(job.Preset.Video.Bitrate)
		if err != nil {
			return 0, 0, err
		}

		cc.SetDimension(width, height)
		cc.SetBitRate(bitrate)
		cc.SetPixFmt(ist.CodecCtx().PixFmt())
	}

	if err := cc.Open(nil); err != nil {
		return 0, 0, err
	}

	ost.SetCodecCtx(cc)

	return ist.Index(), ost.Index(), nil
}

// FFMPEGEncode function is responsible for encoding the file
func FFMPEGEncode(logger lager.Logger, dbInstance db.Storage, jobID string) error {
	log := logger.Session("ffmpeg-encode")
	log.Info("started", lager.Data{"job": jobID})
	defer log.Info("finished")

	gmf.LogSetLevel(gmf.AV_LOG_FATAL)
	job, _ := dbInstance.RetrieveJob(jobID)
	stMap := make(map[int]int, 0)
	var lastDelta int64

	inputCtx, err := gmf.NewInputCtx(job.LocalSource)
	if err != nil {
		log.Error("input-failed", err)
		return err
	}
	defer inputCtx.CloseInputAndRelease()

	outputCtx, err := gmf.NewOutputCtx(job.LocalDestination)
	if err != nil {
		log.Error("output-failed", err)
		return err
	}
	defer outputCtx.CloseOutputAndRelease()

	job.Status = types.JobEncoding
	job.Details = "0%"
	dbInstance.UpdateJob(job.ID, job)

	srcVideoStream, _ := inputCtx.GetBestStream(gmf.AVMEDIA_TYPE_VIDEO)

	videoCodec := GetCodec(job)

	log.Info("add-stream-start", lager.Data{"code": videoCodec})
	i, o, err := addStream(job, videoCodec, outputCtx, srcVideoStream)
	if err != nil {
		log.Error("add-stream-failed", err)
		return err
	}
	log.Info("add-stream-finished")
	stMap[i] = o

	srcAudioStream, err := inputCtx.GetBestStream(gmf.AVMEDIA_TYPE_AUDIO)
	if err != nil {
		return err
	}

	audioCodec := "aac" // default codec
	if job.Preset.Audio.Codec != "aac" {
		audioCodec = job.Preset.Audio.Codec
	}

	i, o, err = addStream(job, audioCodec, outputCtx, srcAudioStream)
	if err != nil {
		return err
	}
	stMap[i] = o

	if err := outputCtx.WriteHeader(); err != nil {
		return err
	}
	totalFrames := float64(srcVideoStream.NbFrames() + srcAudioStream.NbFrames())
	framesCount := float64(0)

	for packet := range inputCtx.GetNewPackets() {
		ist, err := inputCtx.GetStream(packet.StreamIndex())
		if err != nil {
			return err
		}
		ost, err := outputCtx.GetStream(stMap[ist.Index()])
		if err != nil {
			return err
		}

		for frame := range packet.Frames(ist.CodecCtx()) {
			if ost.IsAudio() {
				fsTb := gmf.AVR{Num: 1, Den: ist.CodecCtx().SampleRate()}
				outTb := gmf.AVR{Num: 1, Den: ist.CodecCtx().SampleRate()}

				frame.SetPts(packet.Pts())

				pts := gmf.RescaleDelta(ist.TimeBase(), frame.Pts(), fsTb.AVRational(), frame.NbSamples(), &lastDelta, outTb.AVRational())

				frame.
					SetNbSamples(ost.CodecCtx().FrameSize()).
					SetFormat(ost.CodecCtx().SampleFmt()).
					SetChannelLayout(ost.CodecCtx().ChannelLayout()).
					SetPts(pts)
			} else {
				frame.SetPts(ost.Pts)
			}

			if p, ready, _ := frame.EncodeNewPacket(ost.CodecCtx()); ready {
				if p.Pts() != gmf.AV_NOPTS_VALUE {
					p.SetPts(gmf.RescaleQ(p.Pts(), ost.CodecCtx().TimeBase(), ost.TimeBase()))
				}

				if p.Dts() != gmf.AV_NOPTS_VALUE {
					p.SetDts(gmf.RescaleQ(p.Dts(), ost.CodecCtx().TimeBase(), ost.TimeBase()))
				}

				p.SetStreamIndex(ost.Index())

				if err := outputCtx.WritePacket(p); err != nil {
					return err
				}
				gmf.Release(p)
			}

			ost.Pts++
			framesCount++
			percentage := string(strconv.FormatInt(int64(framesCount/totalFrames*100), 10) + "%")
			if percentage != job.Details {
				job.Details = percentage
				dbInstance.UpdateJob(job.ID, job)
			}
		}
		gmf.Release(packet)
	}

	for i := 0; i < outputCtx.StreamsCnt(); i++ {
		ist, err := inputCtx.GetStream(0)
		if err != nil {
			return err
		}
		ost, err := outputCtx.GetStream(stMap[ist.Index()])
		if err != nil {
			return err
		}

		frame := gmf.NewFrame()

		for {
			if p, ready, _ := frame.FlushNewPacket(ost.CodecCtx()); ready {
				if p.Pts() != gmf.AV_NOPTS_VALUE {
					p.SetPts(gmf.RescaleQ(p.Pts(), ost.CodecCtx().TimeBase(), ost.TimeBase()))
				}

				if p.Dts() != gmf.AV_NOPTS_VALUE {
					p.SetDts(gmf.RescaleQ(p.Dts(), ost.CodecCtx().TimeBase(), ost.TimeBase()))
				}

				p.SetStreamIndex(ost.Index())

				if err := outputCtx.WritePacket(p); err != nil {
					return err
				}
				gmf.Release(p)
			} else {
				gmf.Release(p)
				break
			}

			ost.Pts++
		}

		gmf.Release(frame)
	}
	if job.Details != "100%" {
		job.Details = "100%"
		dbInstance.UpdateJob(job.ID, job)
	}

	return nil
}

// GetProfile returns the GMF profile number based on job profile string
func GetProfile(job types.Job) int {
	profiles := map[string]int{
		"baseline": gmf.FF_PROFILE_H264_BASELINE,
		"main":     gmf.FF_PROFILE_H264_MAIN,
		"high":     gmf.FF_PROFILE_H264_HIGH,
	}

	if job.Preset.Video.Profile {
		return profiles[job.Preset.Video.Profile]
	}
	return gmf.FF_PROFILE_H264_MAIN
}

// GetCodec returns the right codec
func GetCodec(job types.Job) string {
	codecs := map[string]int{
		"h264":   "libx264",
		"vp8":    "libvpx",
		"vp9":    "libvpx-vp9",
		"theora": "libtheora",
	}

	if job.Preset.Video.Codec {
		return codecs[job.Preset.Video.Codec]
	}
	return "libx264"
}

// GetResolution calculate the output resolution based on the preset and input source
func GetResolution(job types.Job, inputWidth int, inputHeight int) (int, int) {
	var width, height int
	if job.Preset.Video.Width == "" && job.Preset.Video.Height == "" {
		return inputWidth, inputHeight
	} else if job.Preset.Video.Width == "" {
		height, _ = strconv.Atoi(job.Preset.Video.Height)
		width = (inputWidth * height) / inputHeight
	} else if job.Preset.Video.Height == "" {
		width, _ = strconv.Atoi(job.Preset.Video.Width)
		height = (inputHeight * width) / inputWidth
	} else {
		width, _ = strconv.Atoi(job.Preset.Video.Width)
		height, _ = strconv.Atoi(job.Preset.Video.Height)
	}
	return width, height
}
<commit_msg>encoders/ffmpeg fix tests<commit_after>package encoders

import (
	"errors"
	"strconv"

	"code.cloudfoundry.org/lager"

	"github.com/3d0c/gmf"
	"github.com/snickers/snickers/db"
	"github.com/snickers/snickers/types"
)

func addStream(job types.Job, codecName string, oc *gmf.FmtCtx, ist *gmf.Stream) (int, int, error) {
	var cc *gmf.CodecCtx
	var ost *gmf.Stream

	codec, err := gmf.FindEncoder(codecName)
	if err != nil {
		return 0, 0, err
	}

	if ost = oc.NewStream(codec); ost == nil {
		return 0, 0, errors.New("unable to create stream in output context")
	}
	defer gmf.Release(ost)

	if cc = gmf.NewCodecCtx(codec); cc == nil {
		return 0, 0, errors.New("unable to create codec context")
	}
	defer gmf.Release(cc)

	// https://ffmpeg.org/pipermail/ffmpeg-devel/2008-January/046900.html
	if oc.IsGlobalHeader() {
		cc.SetFlag(gmf.CODEC_FLAG_GLOBAL_HEADER)
	}

	if codec.IsExperimental() {
		cc.SetStrictCompliance(gmf.FF_COMPLIANCE_EXPERIMENTAL)
	}

	if cc.Type() == gmf.AVMEDIA_TYPE_AUDIO {
		bitrate, err := strconv.Atoi(job.Preset.Audio.Bitrate)
		if err != nil {
			return 0, 0, err
		}

		gop, err := strconv.Atoi(job.Preset.Video.GopSize)
		if err != nil {
			return 0, 0, err
		}

		cc.SetBitRate(bitrate)
		cc.SetGopSize(gop)
		cc.SetSampleFmt(ist.CodecCtx().SampleFmt())
		cc.SetSampleRate(ist.CodecCtx().SampleRate())
		cc.SetChannels(ist.CodecCtx().Channels())
		cc.SelectChannelLayout()
		cc.SelectSampleRate()
	}

	if cc.Type() == gmf.AVMEDIA_TYPE_VIDEO {
		cc.SetTimeBase(gmf.AVR{Num: 1, Den: 25}) // what is this

		if job.Preset.Video.Codec == "h264" {
			profile := GetProfile(job)
			cc.SetProfile(profile)
		}

		width, height := GetResolution(job, ist.CodecCtx().Width(), ist.CodecCtx().Height())

		bitrate, err := strconv.Atoi(job.Preset.Video.Bitrate)
		if err != nil {
			return 0, 0, err
		}

		cc.SetDimension(width, height)
		cc.SetBitRate(bitrate)
		cc.SetPixFmt(ist.CodecCtx().PixFmt())
	}

	if err := cc.Open(nil); err != nil {
		return 0, 0, err
	}

	ost.SetCodecCtx(cc)

	return ist.Index(), ost.Index(), nil
}

// FFMPEGEncode function is responsible for encoding the file
func FFMPEGEncode(logger lager.Logger, dbInstance db.Storage, jobID string) error {
	log := logger.Session("ffmpeg-encode")
	log.Info("started", lager.Data{"job": jobID})
	defer log.Info("finished")

	gmf.LogSetLevel(gmf.AV_LOG_FATAL)
	job, _ := dbInstance.RetrieveJob(jobID)
	stMap := make(map[int]int, 0)
	var lastDelta int64

	inputCtx, err := gmf.NewInputCtx(job.LocalSource)
	if err != nil {
		log.Error("input-failed", err)
		return err
	}
	defer inputCtx.CloseInputAndRelease()

	outputCtx, err := gmf.NewOutputCtx(job.LocalDestination)
	if err != nil {
		log.Error("output-failed", err)
		return err
	}
	defer outputCtx.CloseOutputAndRelease()

	job.Status = types.JobEncoding
	job.Details = "0%"
	dbInstance.UpdateJob(job.ID, job)

	srcVideoStream, _ := inputCtx.GetBestStream(gmf.AVMEDIA_TYPE_VIDEO)

	videoCodec := GetCodec(job)

	log.Info("add-stream-start", lager.Data{"code": videoCodec})
	i, o, err := addStream(job, videoCodec, outputCtx, srcVideoStream)
	if err != nil {
		log.Error("add-stream-failed", err)
		return err
	}
	log.Info("add-stream-finished")
	stMap[i] = o

	srcAudioStream, err := inputCtx.GetBestStream(gmf.AVMEDIA_TYPE_AUDIO)
	if err != nil {
		return err
	}

	audioCodec := "aac" // default codec
	if job.Preset.Audio.Codec != "aac" {
		audioCodec = job.Preset.Audio.Codec
	}

	i, o, err = addStream(job, audioCodec, outputCtx, srcAudioStream)
	if err != nil {
		return err
	}
	stMap[i] = o

	if err := outputCtx.WriteHeader(); err != nil {
		return err
	}
	totalFrames := float64(srcVideoStream.NbFrames() + srcAudioStream.NbFrames())
	framesCount := float64(0)

	for packet := range inputCtx.GetNewPackets() {
		ist, err := inputCtx.GetStream(packet.StreamIndex())
		if err != nil {
			return err
		}
		ost, err := outputCtx.GetStream(stMap[ist.Index()])
		if err != nil {
			return err
		}

		for frame := range packet.Frames(ist.CodecCtx()) {
			if ost.IsAudio() {
				fsTb := gmf.AVR{Num: 1, Den: ist.CodecCtx().SampleRate()}
				outTb := gmf.AVR{Num: 1, Den: ist.CodecCtx().SampleRate()}

				frame.SetPts(packet.Pts())

				pts := gmf.RescaleDelta(ist.TimeBase(), frame.Pts(), fsTb.AVRational(), frame.NbSamples(), &lastDelta, outTb.AVRational())

				frame.
					SetNbSamples(ost.CodecCtx().FrameSize()).
					SetFormat(ost.CodecCtx().SampleFmt()).
					SetChannelLayout(ost.CodecCtx().ChannelLayout()).
					SetPts(pts)
			} else {
				frame.SetPts(ost.Pts)
			}

			if p, ready, _ := frame.EncodeNewPacket(ost.CodecCtx()); ready {
				if p.Pts() != gmf.AV_NOPTS_VALUE {
					p.SetPts(gmf.RescaleQ(p.Pts(), ost.CodecCtx().TimeBase(), ost.TimeBase()))
				}

				if p.Dts() != gmf.AV_NOPTS_VALUE {
					p.SetDts(gmf.RescaleQ(p.Dts(), ost.CodecCtx().TimeBase(), ost.TimeBase()))
				}

				p.SetStreamIndex(ost.Index())

				if err := outputCtx.WritePacket(p); err != nil {
					return err
				}
				gmf.Release(p)
			}

			ost.Pts++
			framesCount++
			percentage := string(strconv.FormatInt(int64(framesCount/totalFrames*100), 10) + "%")
			if percentage != job.Details {
				job.Details = percentage
				dbInstance.UpdateJob(job.ID, job)
			}
		}
		gmf.Release(packet)
	}

	for i := 0; i < outputCtx.StreamsCnt(); i++ {
		ist, err := inputCtx.GetStream(0)
		if err != nil {
			return err
		}
		ost, err := outputCtx.GetStream(stMap[ist.Index()])
		if err != nil {
			return err
		}

		frame := gmf.NewFrame()

		for {
			if p, ready, _ := frame.FlushNewPacket(ost.CodecCtx()); ready {
				if p.Pts() != gmf.AV_NOPTS_VALUE {
					p.SetPts(gmf.RescaleQ(p.Pts(), ost.CodecCtx().TimeBase(), ost.TimeBase()))
				}

				if p.Dts() != gmf.AV_NOPTS_VALUE {
					p.SetDts(gmf.RescaleQ(p.Dts(), ost.CodecCtx().TimeBase(), ost.TimeBase()))
				}

				p.SetStreamIndex(ost.Index())

				if err := outputCtx.WritePacket(p); err != nil {
					return err
				}
				gmf.Release(p)
			} else {
				gmf.Release(p)
				break
			}

			ost.Pts++
		}

		gmf.Release(frame)
	}
	if job.Details != "100%" {
		job.Details = "100%"
		dbInstance.UpdateJob(job.ID, job)
	}

	return nil
}

// GetProfile returns the GMF profile number based on job profile string
func GetProfile(job types.Job) int {
	profiles := map[string]int{
		"baseline": gmf.FF_PROFILE_H264_BASELINE,
		"main":     gmf.FF_PROFILE_H264_MAIN,
		"high":     gmf.FF_PROFILE_H264_HIGH,
	}

	if job.Preset.Video.Profile != "" {
		return profiles[job.Preset.Video.Profile]
	}
	return gmf.FF_PROFILE_H264_MAIN
}

// GetCodec returns the right codec
func GetCodec(job types.Job) string {
	codecs := map[string]string{
		"h264":   "libx264",
		"vp8":    "libvpx",
		"vp9":    "libvpx-vp9",
		"theora": "libtheora",
	}

	if job.Preset.Video.Codec != "" {
		return codecs[job.Preset.Video.Codec]
	}
	return "libx264"
}

// GetResolution calculate the output resolution based on the preset and input source
func GetResolution(job types.Job, inputWidth int, inputHeight int) (int, int) {
	var width, height int
	if job.Preset.Video.Width == "" && job.Preset.Video.Height == "" {
		return inputWidth, inputHeight
	} else if job.Preset.Video.Width == "" {
		height, _ = strconv.Atoi(job.Preset.Video.Height)
		width = (inputWidth * height) / inputHeight
	} else if job.Preset.Video.Height == "" {
		width, _ = strconv.Atoi(job.Preset.Video.Width)
		height = (inputHeight * width) / inputWidth
	} else {
		width, _ = strconv.Atoi(job.Preset.Video.Width)
		height, _ = strconv.Atoi(job.Preset.Video.Height)
	}
	return width, height
}
<|endoftext|>


<commit_before>package templatefs

import (
  "html/template"
)

type AssetFunc func(string) ([]byte, error)

type Template struct {
  AssetFunc AssetFunc
  tmpl *template.Template
}

func New(name string, fn AssetFunc) *Template {
  return &Template{fn, template.New(name)}
}

func (t *Template) Parse(filename string) (*Template, error) {
  tmplBytes, err := Asset(fileName)
  if err != nil {
    return nil, err
  }
  tmplStr := string(tmplBytes)
  newTmpl, err := t.tmpl.Parse(tmplStr)
  if err != nil {
    return nil, err
  }
  t.tmpl := newTmpl
}
<commit_msg>filling out the template adapter more<commit_after>package templatefs

import (
  "html/template"
  "io"
)

type AssetFunc func(string) ([]byte, error)

type Template struct {
  AssetFunc AssetFunc
  tmpl *template.Template
}

func New(name string, fn AssetFunc) *Template {
  return &Template{fn, template.New(name)}
}

func (t *Template) Name() string {
  return t.tmpl.Name()
}

func (t *Template) Funcs(funcMap template.FuncMap) *Template {
  return t.tmpl.Funcs(funcMap)
}

func (t *Template) Parse(filename string) (*Template, error) {
  tmplStr, err := t.file(filename)
  if err != nil {
    return nil, err
  }
  newTmpl, err := t.tmpl.Parse(tmplStr)
  if err != nil {
    return nil, err
  }
  t.tmpl := newTmpl
  return t, nil
}

func (t *Template) ParseFiles(filenames ...string) (*Template, error) {
  fileStrs := []string{}
  for _, filename := range filenames {
    fileStr, err := t.file(filename)
    if err != nil {
      return nil, err
    }
    fileStrs = append(fileStrs, fileStr)
  }
  newTmpl, err := t.tmpl.ParseFiles(fileStrs...)
  if err != nil {
    return nil, err
  }
  t.tmpl = newTmpl
  return t, nil
}

func (t *Template) Execute(w io.Writer, data interface{}) error {
  return t.tmpl.Execute(w, data)
}

func (t *Template) file(name string) ([]byte, error) {
  tmplBytes, err := Asset(fileName)
  if err != nil {
    return nil, err
  }
  return string(tmplBytes), nil
}
<|endoftext|>


<commit_before>package utils

import (
// "fmt"
)
<commit_msg>test db<commit_after>package utils

import (
	"fmt"
	"os"

	"github.com/go-sql-driver/mysql"
	//  The fantastic ORM library for Golang, aims to be developer friendly.
	"github.com/jinzhu/gorm"
	//	A pure Go postgres driver for Go's database/sql package
	"github.com/lib/pq"
)

func TestDB() *gorm.DB {
	var db *gorm.DB
	var err error
	var dbuser, dbpwd, dbname = "root", "19911220", "ec_test"
	if os.Getenv("DB_USER") != "" {
		dbuser = os.Getenv("DB_USER")
	}

	if os.Getenv("DB_PWD") != "" {
		dbpwd = os.Getenv("DB_PWD")
	}

	if os.Getenv("TEST_DB") == "postgres" {
		db, err := gorm.Open("prostgres", fmt.Sprintf("postgres://%s:%s@localhost/%s?sslmode=disable", dbuser, dbpwd, dbname))
	} else {
		db, err := gorm.Open("mysql", fmt.Sprintf("%s:%s@%s?charset=utf8&parseTime=True&loc=Local", dbuser, dbpwd, dbname))
	}

	if err != nil {
		panic(err)
	}

	if os.Getenv("DEBUG") != "" {
		db.LogMode(true)
	}

	return db
}
<|endoftext|>


<commit_before>package etreeutils

import "github.com/beevik/etree"

type SortedAttrs []etree.Attr

func (a SortedAttrs) Len() int {
	return len(a)
}

func (a SortedAttrs) Swap(i, j int) {
	a[i], a[j] = a[j], a[i]
}

func (a SortedAttrs) Less(i, j int) bool {
	// As I understand it: any "xmlns" attribute should come first, followed by any
	// any "xmlns:prefix" attributes, presumably ordered by prefix. Lastly any other
	// attributes in lexicographical order.
	if a[i].Space == defaultPrefix && a[i].Key == xmlnsPrefix {
		return true
	}

	if a[j].Space == defaultPrefix && a[j].Key == xmlnsPrefix {
		return false
	}

	if a[i].Space == xmlnsPrefix {
		if a[j].Space == xmlnsPrefix {
			return a[i].Key < a[j].Key
		}
		return true
	}

	if a[j].Space == xmlnsPrefix {
		return false
	}

	if a[i].Space == a[j].Space {
		return a[i].Key < a[j].Key
	}

	return a[i].Space < a[j].Space
}
<commit_msg>Comment and improve attribute sorting<commit_after>package etreeutils

import "github.com/beevik/etree"

// SortedAttrs provides sorting capabilities, compatible with XML C14N, on top
// of an []etree.Attr
type SortedAttrs []etree.Attr

func (a SortedAttrs) Len() int {
	return len(a)
}

func (a SortedAttrs) Swap(i, j int) {
	a[i], a[j] = a[j], a[i]
}

func (a SortedAttrs) Less(i, j int) bool {
	// This is the best reference I've found on sort order:
	// http://dst.lbl.gov/~ksb/Scratch/XMLC14N.html

	// If attr j is a default namespace declaration, attr i may
	// not be strictly "less" than it.
	if a[j].Space == defaultPrefix && a[j].Key == xmlnsPrefix {
		return false
	}

	// Otherwise, if attr i is a default namespace declaration, it
	// must be less than anything else.
	if a[i].Space == defaultPrefix && a[i].Key == xmlnsPrefix {
		return true
	}

	// Next, namespace prefix declarations, sorted by prefix, come before
	// anythign else.
	if a[i].Space == xmlnsPrefix {
		if a[j].Space == xmlnsPrefix {
			return a[i].Key < a[j].Key
		}
		return true
	}

	if a[j].Space == xmlnsPrefix {
		return false
	}

	// Then come unprefixed attributes, sorted by key.
	if a[i].Space == defaultPrefix {
		if a[j].Space == defaultPrefix {
			return a[i].Key < a[j].Key
		}
		return true
	}

	if a[j].Space == defaultPrefix {
		return false
	}

	// Wow. We're still going. Finally, attributes in the same namespace should be
	// sorted by key. Attributes in different namespaces should be sorted by the
	// actual namespace (_not_ the prefix). For now just use the prefix.
	if a[i].Space == a[j].Space {
		return a[i].Key < a[j].Key
	}

	return a[i].Space < a[j].Space
}
<|endoftext|>


<commit_before>package stow

import (
	"os"
	"testing"

	"github.com/boltdb/bolt"
)

func TestJson(t *testing.T) {
	db, err := bolt.Open("my.db", 0600, nil)
	if err != nil {
		t.Error(err.Error())
	}
	defer os.Remove("my.db")
	defer db.Close()

	type MyType struct {
		FirstName string `json:"first"`
		LastName  string `json:"last"`
	}

	store := NewJSONStore(db, []byte("json"))
	store.Put([]byte("hello"), &MyType{"Derek", "Kered"})

	var found bool
	err = store.ForEach(func(name MyType) {
		found = true
		if name.FirstName != "Derek" || name.LastName != "Kered" {
			t.Errorf("mismatching name %s", name)
		}
	})

	if !found {
		t.Errorf("ForEach failed!")
	}

	if err != nil {
		t.Error(err.Error())
	}

	var name MyType
	store.Pull([]byte("hello"), &name)

	if name.FirstName != "Derek" || name.LastName != "Kered" {
		t.Errorf("Unexpected name: %v", name)
	}
}
<commit_msg>adding more test cases<commit_after>package stow

import (
	"encoding/gob"
	"os"
	"testing"

	"github.com/boltdb/bolt"
)

type MyType struct {
	FirstName string `json:"first"`
	LastName  string `json:"last"`
}

func init() {
	gob.Register(MyType{})
}

func testStore(t *testing.T, store *Store) {
	store.Put([]byte("hello"), &MyType{"Derek", "Kered"})

	var found bool
	err := store.ForEach(func(name MyType) {
		found = true
		if name.FirstName != "Derek" || name.LastName != "Kered" {
			t.Errorf("mismatching name %s", name)
		}
	})

	if !found {
		t.Errorf("ForEach failed!")
	}

	if err != nil {
		t.Error(err.Error())
	}

	var name MyType
	store.Pull([]byte("hello"), &name)

	if name.FirstName != "Derek" || name.LastName != "Kered" {
		t.Errorf("Unexpected name: %v", name)
	}
}

func TestJSON(t *testing.T) {
	db, err := bolt.Open("my1.db", 0600, nil)
	if err != nil {
		t.Error(err.Error())
	}
	defer os.Remove("my1.db")
	defer db.Close()

	testStore(t, NewJSONStore(db, []byte("json")))
}

func TestXML(t *testing.T) {
	db, err := bolt.Open("my2.db", 0600, nil)
	if err != nil {
		t.Error(err.Error())
	}
	defer os.Remove("my2.db")
	defer db.Close()

	testStore(t, NewXMLStore(db, []byte("xml")))
}

func TestGob(t *testing.T) {
	db, err := bolt.Open("my3.db", 0600, nil)
	if err != nil {
		t.Error(err.Error())
	}
	defer os.Remove("my3.db")
	defer db.Close()

	testStore(t, NewStore(db, []byte("gob")))
}
<|endoftext|>


<commit_before>package str

import "strings"

func Clean(s string) string {
	return strings.Trim(strings.Trim(s, " "), "　")
}
<commit_msg>Fixed clean method for string<commit_after>package str

import "strings"

func Clean(s string) string {
	return strings.TrimSpace(strings.Trim(strings.Trim(s, " "), "　"))
}
<|endoftext|>


<commit_before>package immortal

import (
	"log"
	"os"
	"os/signal"
	"sync/atomic"
	"syscall"
	"time"
)

func Supervise(d *Daemon) {
	var (
		err  error
		info = make(chan os.Signal)
		p    *process
		pid  int
		run  = make(chan struct{}, 1)
	)

	// start a new process
	p, err = d.Run(NewProcess(d.cfg))
	if err != nil {
		log.Fatal(err)
	}

	// Info loop kill 3 pid get stats
	signal.Notify(info, syscall.SIGQUIT)
	go d.Info(info)

	// create a supervisor
	s := &Sup{p}

	// listen on control for signals
	if d.cfg.ctrl {
		s.ReadFifoControl(d.fifo_control, d.fifo)
	}

	for {
		select {
		case <-d.quit:
			return
		case <-run:
			// create a new process
			p, err = d.Run(NewProcess(d.cfg))
			if err != nil {
				log.Print(err)
			}
			s = &Sup{p}
		default:
			select {
			case err := <-p.errch:
				// unlock, or lock once
				atomic.StoreUint32(&d.lock, d.lock_once)
				if err != nil && err.Error() == "EXIT" {
					log.Printf("PID: %d Exited", pid)
				} else {
					log.Printf("PID %d terminated, %s [%v user  %v sys  %s up]\n",
						p.cmd.ProcessState.Pid(),
						p.cmd.ProcessState,
						p.cmd.ProcessState.UserTime(),
						p.cmd.ProcessState.SystemTime(),
						time.Since(p.sTime),
					)
				}

				// follow the new pid and stop running the command
				// unless the new pid dies
				if d.cfg.Pid.Follow != "" {
					pid, err = s.ReadPidFile(d.cfg.Pid.Follow)
					if err != nil {
						log.Printf("Cannot read pidfile:%s, %s", d.cfg.Pid.Follow, err)
						run <- struct{}{}
					} else {
						// check if pid in file is valid
						if pid > 1 && pid != p.Pid() && s.IsRunning(pid) {
							log.Printf("Watching pid %d on file: %s", pid, d.cfg.Pid.Follow)
							s.WatchPid(pid, p.errch)
						} else {
							// if cmd exits or process is kill
							run <- struct{}{}
						}
					}
				} else {
					run <- struct{}{}
				}
			case fifo := <-d.fifo:
				if fifo.err != nil {
					log.Printf("control error: %s", fifo.err)
				}
				s.HandleSignals(fifo.msg, d)
			}
		}
	}
}
<commit_msg>implemented wait based on uptime<commit_after>package immortal

import (
	"log"
	"os"
	"os/signal"
	"sync/atomic"
	"syscall"
	"time"
)

func Supervise(d *Daemon) {
	var (
		err  error
		info = make(chan os.Signal)
		p    *process
		pid  int
		run                = make(chan struct{}, 1)
		wait time.Duration = 0 * time.Second
	)

	// start a new process
	p, err = d.Run(NewProcess(d.cfg))
	if err != nil {
		log.Fatal(err)
	}

	// Info loop kill 3 pid get stats
	signal.Notify(info, syscall.SIGQUIT)
	go d.Info(info)

	// create a supervisor
	s := &Sup{p}

	// listen on control for signals
	if d.cfg.ctrl {
		s.ReadFifoControl(d.fifo_control, d.fifo)
	}

	for {
		select {
		case <-d.quit:
			return
		case <-run:
			time.Sleep(wait)
			// create a new process
			p, err = d.Run(NewProcess(d.cfg))
			if err != nil {
				log.Print(err)
			}
			s = &Sup{p}
		default:
			select {
			case err := <-p.errch:
				// unlock, or lock once
				atomic.StoreUint32(&d.lock, d.lock_once)
				if err != nil && err.Error() == "EXIT" {
					log.Printf("PID: %d Exited", pid)
				} else {
					log.Printf("PID %d terminated, %s [%v user  %v sys  %s up]\n",
						p.cmd.ProcessState.Pid(),
						p.cmd.ProcessState,
						p.cmd.ProcessState.UserTime(),
						p.cmd.ProcessState.SystemTime(),
						time.Since(p.sTime),
					)
					// calculate time for next reboot to avoid looping & consuming CPU in case can't restart
					uptime := p.eTime.Sub(p.sTime)
					wait = 0 * time.Second
					if uptime < time.Second {
						wait = time.Second - uptime
					}
				}

				// follow the new pid and stop running the command
				// unless the new pid dies
				if d.cfg.Pid.Follow != "" {
					pid, err = s.ReadPidFile(d.cfg.Pid.Follow)
					if err != nil {
						log.Printf("Cannot read pidfile:%s, %s", d.cfg.Pid.Follow, err)
						run <- struct{}{}
					} else {
						// check if pid in file is valid
						if pid > 1 && pid != p.Pid() && s.IsRunning(pid) {
							log.Printf("Watching pid %d on file: %s", pid, d.cfg.Pid.Follow)
							s.WatchPid(pid, p.errch)
						} else {
							// if cmd exits or process is kill
							run <- struct{}{}
						}
					}
				} else {
					run <- struct{}{}
				}
			case fifo := <-d.fifo:
				if fifo.err != nil {
					log.Printf("control error: %s", fifo.err)
				}
				s.HandleSignals(fifo.msg, d)
			}
		}
	}
}
<|endoftext|>


<commit_before>package xmpp

import (
	"crypto/tls"
	"encoding/xml"
	"errors"
	"fmt"
	"log"
)

// Config structure used to create a new XMPP client connection.
type ClientConfig struct {
	// Don't upgrade the connection to TLS, even if the server supports it. If
	// the server *requires* TLS then this option is ignored.
	NoTLS bool

	// Skip verification of the server's certificate chain. Probably only
	// useful during development.
	InsecureSkipVerify bool
}

// Create a client XMPP over the stream.
func NewClientXMPP(stream *Stream, jid JID, password string, config *ClientConfig) (*XMPP, error) {

	for {

		if err := startClient(stream, jid); err != nil {
			return nil, err
		}

		// Read features.
		f := new(features)
		if err := stream.Decode(f); err != nil {
			return nil, err
		}

		// TLS?
		if f.StartTLS != nil && (f.StartTLS.Required != nil || !config.NoTLS) {
			log.Println("Start TLS")
			if err := startTLS(stream, config); err != nil {
				return nil, err
			}
			continue // Restart
		}

		// Authentication
		if f.Mechanisms != nil {
			log.Println("Authenticating")
			if err := authenticate(stream, f.Mechanisms.Mechanisms, jid.Node, password); err != nil {
				return nil, err
			}
			continue // Restart
		}

		// Bind resource.
		if f.Bind != nil {
			log.Println("Binding resource.")
			boundJID, err := bindResource(stream, jid)
			if err != nil {
				return nil, err
			}
			jid = boundJID
		}

		break
	}

	return newXMPP(jid, stream), nil
}

func startClient(stream *Stream, jid JID) error {

	start := xml.StartElement{
		xml.Name{"stream", "stream"},
		[]xml.Attr{
			xml.Attr{xml.Name{"", "xmlns"}, "jabber:client"},
			xml.Attr{xml.Name{"xmlns", "stream"}, "http://etherx.jabber.org/streams"},
			xml.Attr{xml.Name{"", "from"}, jid.Full()},
			xml.Attr{xml.Name{"", "to"}, jid.Domain},
		},
	}

	if err := stream.SendStart(&start); err != nil {
		return err
	}

	if _, err := stream.Next(&xml.Name{nsStream, "stream"}); err != nil {
		return err
	}

	return nil
}

func startTLS(stream *Stream, config *ClientConfig) error {

	if err := stream.Send(&tlsStart{}); err != nil {
		return err
	}

	p := tlsProceed{}
	if err := stream.Decode(&p); err != nil {
		return err
	}

	tlsConfig := tls.Config{InsecureSkipVerify: config.InsecureSkipVerify}
	return stream.UpgradeTLS(&tlsConfig)
}

type tlsStart struct {
	XMLName xml.Name `xml:"urn:ietf:params:xml:ns:xmpp-tls starttls"`
}

type tlsProceed struct {
	XMLName xml.Name `xml:"urn:ietf:params:xml:ns:xmpp-tls proceed"`
}

func authenticate(stream *Stream, mechanisms []string, user, password string) error {

	log.Println("authenticate, mechanisms=", mechanisms)

	if !stringSliceContains(mechanisms, "PLAIN") {
		return errors.New("Only PLAIN supported for now")
	}

	return authenticatePlain(stream, user, password)
}

func authenticatePlain(stream *Stream, user, password string) error {

	auth := saslAuth{Mechanism: "PLAIN", Message: saslEncodePlain(user, password)}
	if err := stream.Send(&auth); err != nil {
		return err
	}

	if se, err := stream.Next(nil); err != nil {
		return err
	} else {
		if se.Name.Local == "failure" {
			f := new(saslFailure)
			stream.DecodeElement(f, se)
			return errors.New(fmt.Sprintf("Authentication failed: %s", f.Reason.Local))
		}
	}

	return nil
}

type saslAuth struct {
	XMLName xml.Name `xml:"urn:ietf:params:xml:ns:xmpp-sasl auth"`
	Mechanism string `xml:"mechanism,attr"`
	Message string `xml:",chardata"`
}

func bindResource(stream *Stream, jid JID) (JID, error) {

	req := Iq{Id: "foo", Type: "set"}
	if jid.Resource == "" {
		req.PayloadEncode(bindIq{})
	} else {
		req.PayloadEncode(bindIq{Resource: jid.Resource})
	}
	if err := stream.Send(req); err != nil {
		return JID{}, err
	}

	resp := Iq{}
	err := stream.Decode(&resp)
	if err != nil {
		return JID{}, err
	}
	bindResp := bindIq{}
	resp.PayloadDecode(&bindResp)

	boundJID, err := ParseJID(bindResp.JID)
	return boundJID, nil
}

type bindIq struct {
	XMLName xml.Name `xml:"urn:ietf:params:xml:ns:xmpp-bind bind"`
	Resource string `xml:"resource,omitempty"`
	JID string `xml:"jid,omitempty"`
}

func stringSliceContains(l []string, m string) bool {
	for _, i := range l {
		if i == m {
			return true
		}
	}
	return false
}

type features struct {
	XMLName xml.Name `xml:"http://etherx.jabber.org/streams features"`
	StartTLS *tlsStartTLS `xml:"starttls"`
	Mechanisms *mechanisms `xml:"mechanisms"`
	Bind *bind `xml:"bind"`
}

type bind struct {
	XMLName xml.Name `xml:"urn:ietf:params:xml:ns:xmpp-bind bind"`
	Required *required `xml:"required"`
}

type mechanisms struct {
	XMLName xml.Name `xml:"urn:ietf:params:xml:ns:xmpp-sasl mechanisms"`
	Mechanisms []string `xml:"mechanism"`
}

type tlsStartTLS struct {
	XMLName xml.Name `xml:"urn:ietf:params:xml:ns:xmpp-tls starttls"`
	Required *required `xml:"required"`
}

type required struct {}

type saslSuccess struct {
	XMLName xml.Name `xml:"urn:ietf:params:xml:ns:xmpp-sasl success"`
}

type saslFailure struct {
	XMLName xml.Name `xml:"urn:ietf:params:xml:ns:xmpp-sasl failure"`
	Reason xml.Name `xml:",any"`
}

// BUG(matt): Don't use "foo" as the <iq/> id during resource binding.
<commit_msg>No client config means default config.<commit_after>package xmpp

import (
	"crypto/tls"
	"encoding/xml"
	"errors"
	"fmt"
	"log"
)

// Config structure used to create a new XMPP client connection.
type ClientConfig struct {
	// Don't upgrade the connection to TLS, even if the server supports it. If
	// the server *requires* TLS then this option is ignored.
	NoTLS bool

	// Skip verification of the server's certificate chain. Probably only
	// useful during development.
	InsecureSkipVerify bool
}

// Create a client XMPP over the stream.
func NewClientXMPP(stream *Stream, jid JID, password string, config *ClientConfig) (*XMPP, error) {

	if config == nil {
		config = &ClientConfig{}
	}

	for {

		if err := startClient(stream, jid); err != nil {
			return nil, err
		}

		// Read features.
		f := new(features)
		if err := stream.Decode(f); err != nil {
			return nil, err
		}

		// TLS?
		if f.StartTLS != nil && (f.StartTLS.Required != nil || !config.NoTLS) {
			log.Println("Start TLS")
			if err := startTLS(stream, config); err != nil {
				return nil, err
			}
			continue // Restart
		}

		// Authentication
		if f.Mechanisms != nil {
			log.Println("Authenticating")
			if err := authenticate(stream, f.Mechanisms.Mechanisms, jid.Node, password); err != nil {
				return nil, err
			}
			continue // Restart
		}

		// Bind resource.
		if f.Bind != nil {
			log.Println("Binding resource.")
			boundJID, err := bindResource(stream, jid)
			if err != nil {
				return nil, err
			}
			jid = boundJID
		}

		break
	}

	return newXMPP(jid, stream), nil
}

func startClient(stream *Stream, jid JID) error {

	start := xml.StartElement{
		xml.Name{"stream", "stream"},
		[]xml.Attr{
			xml.Attr{xml.Name{"", "xmlns"}, "jabber:client"},
			xml.Attr{xml.Name{"xmlns", "stream"}, "http://etherx.jabber.org/streams"},
			xml.Attr{xml.Name{"", "from"}, jid.Full()},
			xml.Attr{xml.Name{"", "to"}, jid.Domain},
		},
	}

	if err := stream.SendStart(&start); err != nil {
		return err
	}

	if _, err := stream.Next(&xml.Name{nsStream, "stream"}); err != nil {
		return err
	}

	return nil
}

func startTLS(stream *Stream, config *ClientConfig) error {

	if err := stream.Send(&tlsStart{}); err != nil {
		return err
	}

	p := tlsProceed{}
	if err := stream.Decode(&p); err != nil {
		return err
	}

	tlsConfig := tls.Config{InsecureSkipVerify: config.InsecureSkipVerify}
	return stream.UpgradeTLS(&tlsConfig)
}

type tlsStart struct {
	XMLName xml.Name `xml:"urn:ietf:params:xml:ns:xmpp-tls starttls"`
}

type tlsProceed struct {
	XMLName xml.Name `xml:"urn:ietf:params:xml:ns:xmpp-tls proceed"`
}

func authenticate(stream *Stream, mechanisms []string, user, password string) error {

	log.Println("authenticate, mechanisms=", mechanisms)

	if !stringSliceContains(mechanisms, "PLAIN") {
		return errors.New("Only PLAIN supported for now")
	}

	return authenticatePlain(stream, user, password)
}

func authenticatePlain(stream *Stream, user, password string) error {

	auth := saslAuth{Mechanism: "PLAIN", Message: saslEncodePlain(user, password)}
	if err := stream.Send(&auth); err != nil {
		return err
	}

	if se, err := stream.Next(nil); err != nil {
		return err
	} else {
		if se.Name.Local == "failure" {
			f := new(saslFailure)
			stream.DecodeElement(f, se)
			return errors.New(fmt.Sprintf("Authentication failed: %s", f.Reason.Local))
		}
	}

	return nil
}

type saslAuth struct {
	XMLName xml.Name `xml:"urn:ietf:params:xml:ns:xmpp-sasl auth"`
	Mechanism string `xml:"mechanism,attr"`
	Message string `xml:",chardata"`
}

func bindResource(stream *Stream, jid JID) (JID, error) {

	req := Iq{Id: "foo", Type: "set"}
	if jid.Resource == "" {
		req.PayloadEncode(bindIq{})
	} else {
		req.PayloadEncode(bindIq{Resource: jid.Resource})
	}
	if err := stream.Send(req); err != nil {
		return JID{}, err
	}

	resp := Iq{}
	err := stream.Decode(&resp)
	if err != nil {
		return JID{}, err
	}
	bindResp := bindIq{}
	resp.PayloadDecode(&bindResp)

	boundJID, err := ParseJID(bindResp.JID)
	return boundJID, nil
}

type bindIq struct {
	XMLName xml.Name `xml:"urn:ietf:params:xml:ns:xmpp-bind bind"`
	Resource string `xml:"resource,omitempty"`
	JID string `xml:"jid,omitempty"`
}

func stringSliceContains(l []string, m string) bool {
	for _, i := range l {
		if i == m {
			return true
		}
	}
	return false
}

type features struct {
	XMLName xml.Name `xml:"http://etherx.jabber.org/streams features"`
	StartTLS *tlsStartTLS `xml:"starttls"`
	Mechanisms *mechanisms `xml:"mechanisms"`
	Bind *bind `xml:"bind"`
}

type bind struct {
	XMLName xml.Name `xml:"urn:ietf:params:xml:ns:xmpp-bind bind"`
	Required *required `xml:"required"`
}

type mechanisms struct {
	XMLName xml.Name `xml:"urn:ietf:params:xml:ns:xmpp-sasl mechanisms"`
	Mechanisms []string `xml:"mechanism"`
}

type tlsStartTLS struct {
	XMLName xml.Name `xml:"urn:ietf:params:xml:ns:xmpp-tls starttls"`
	Required *required `xml:"required"`
}

type required struct {}

type saslSuccess struct {
	XMLName xml.Name `xml:"urn:ietf:params:xml:ns:xmpp-sasl success"`
}

type saslFailure struct {
	XMLName xml.Name `xml:"urn:ietf:params:xml:ns:xmpp-sasl failure"`
	Reason xml.Name `xml:",any"`
}

// BUG(matt): Don't use "foo" as the <iq/> id during resource binding.
<|endoftext|>


<commit_before>package xmpp

import "encoding/xml"

// XMPP <iq/> stanza.
type Iq struct {
	XMLName xml.Name `xml:"iq"`
	Id string `xml:"id,attr"`
	Type string `xml:"type,attr"`
	To string `xml:"to,attr,omitempty"`
	From string `xml:"from,attr,omitempty"`
	Payload string `xml:",innerxml"`
}

// Encode the value to an XML string and set as the payload. See xml.Marshal
// for how the value is encoded.
func (iq *Iq) PayloadEncode(v interface{}) error {
	bytes, err := xml.Marshal(v)
	if err != nil {
		return err
	}
	iq.Payload = string(bytes)
	return nil
}

// Decode the payload (an XML string) into the given value. See xml.Unmarshal
// for how the value is decoded.
func (iq *Iq) PayloadDecode(v interface{}) error {
	return xml.Unmarshal([]byte(iq.Payload), v)
}

// XMPP <message/> stanza.
type Message struct {
	XMLName xml.Name `xml:"message"`
	Id string `xml:"id,attr,omitempty"`
	Type string `xml:"type,attr,omitempty"`
	To string `xml:"to,attr,omitempty"`
	From string `xml:"from,attr,omitempty"`
	Subject string `xml:"subject,omitempty"`
	Body string `xml:"body,omitempty"`
}

// XMPP <presence/> stanza.
type Presence struct {
	XMLName xml.Name `xml:"presence"`
	Id string `xml:"id,attr,omitempty"`
	Type string `xml:"type,attr,omitempty"`
	To string `xml:"to,attr,omitempty"`
	From string `xml:"from,attr,omitempty"`
}

// XMPP <error/> stanza.
type Error struct {
	XMLName xml.Name `xml:"error"`
}
<commit_msg>Add Iq.Error and define xml tags for Error type.<commit_after>package xmpp

import (
	"encoding/xml"
	"fmt"
)

// XMPP <iq/> stanza.
type Iq struct {
	XMLName xml.Name `xml:"iq"`
	Id string `xml:"id,attr"`
	Type string `xml:"type,attr"`
	To string `xml:"to,attr,omitempty"`
	From string `xml:"from,attr,omitempty"`
	Payload string `xml:",innerxml"`
	Error *Error `xml:"error"`
}

// Encode the value to an XML string and set as the payload. See xml.Marshal
// for how the value is encoded.
func (iq *Iq) PayloadEncode(v interface{}) error {
	bytes, err := xml.Marshal(v)
	if err != nil {
		return err
	}
	iq.Payload = string(bytes)
	return nil
}

// Decode the payload (an XML string) into the given value. See xml.Unmarshal
// for how the value is decoded.
func (iq *Iq) PayloadDecode(v interface{}) error {
	return xml.Unmarshal([]byte(iq.Payload), v)
}

// XMPP <message/> stanza.
type Message struct {
	XMLName xml.Name `xml:"message"`
	Id string `xml:"id,attr,omitempty"`
	Type string `xml:"type,attr,omitempty"`
	To string `xml:"to,attr,omitempty"`
	From string `xml:"from,attr,omitempty"`
	Subject string `xml:"subject,omitempty"`
	Body string `xml:"body,omitempty"`
}

// XMPP <presence/> stanza.
type Presence struct {
	XMLName xml.Name `xml:"presence"`
	Id string `xml:"id,attr,omitempty"`
	Type string `xml:"type,attr,omitempty"`
	To string `xml:"to,attr,omitempty"`
	From string `xml:"from,attr,omitempty"`
}

// XMPP <error/>. May occur as a top-level stanza or embedded in another
// stanza, e.g. an <iq type="error"/>.
type Error struct {
	XMLName xml.Name `xml:"error"`
	Type string `xml:"type,attr"`
	Text string `xml:"text"`
}

func (e Error) Error() string {
	return fmt.Sprintf("%s: %s", e.Type, e.Text)
}
<|endoftext|>


<commit_before>package main

import (
	"bytes"
	"crypto/rand"
	"crypto/rsa"
	"encoding/binary"
	"errors"
	"flag"
	"fmt"
	"io"
	"log"
	"net"
	"os"
	"os/exec"

	"github.com/googlegenomics/pipelines-tools/gce"
	"github.com/kr/pty"
	"golang.org/x/crypto/ssh"
)

var (
	port = flag.Uint("port", 22, "the port to listen on")
)

func main() {
	flag.Parse()

	config, listener, err := startServer(*port)
	if err != nil {
		log.Fatalf("Failed to start server: %v", err)
	}
	log.Printf("Listening on %s...", listener.Addr())

	for {
		connection, err := listener.Accept()
		if err != nil {
			log.Fatalf("Failed to accept incoming connection: %v", err)
		}

		go func() {
			handleConnection(connection, config)
		}()
	}
}

func startServer(port uint) (*ssh.ServerConfig, net.Listener, error) {
	config, err := getConfiguration()
	if err != nil {
		return nil, nil, fmt.Errorf("getting configuration: %v", err)
	}

	listener, err := net.Listen("tcp", fmt.Sprintf(":%d", port))
	if err != nil {
		return nil, nil, fmt.Errorf("listen: %v", err)
	}
	return config, listener, nil
}

func getConfiguration() (*ssh.ServerConfig, error) {
	key, err := rsa.GenerateKey(rand.Reader, 2048)
	if err != nil {
		return nil, fmt.Errorf("generating server key pair: %v", err)
	}

	signer, err := ssh.NewSignerFromKey(key)
	if err != nil {
		return nil, fmt.Errorf("creating signer: %v", err)
	}

	config := &ssh.ServerConfig{
		PublicKeyCallback: func(conn ssh.ConnMetadata, key ssh.PublicKey) (*ssh.Permissions, error) {
			authorizedKeys, err := gce.GetAuthorizedKeys()
			if err != nil {
				return nil, fmt.Errorf("getting the authorized keys: %v", err)
			}

			if !authorizedKeys[string(key.Marshal())] {
				return nil, errors.New("unauthorized")
			}
			return nil, nil
		},
	}
	config.AddHostKey(signer)
	return config, nil
}

func handleConnection(conn net.Conn, serverConfig *ssh.ServerConfig) error {
	_, chans, reqs, err := ssh.NewServerConn(conn, serverConfig)
	if err != nil {
		return fmt.Errorf("handshake: %v", err)
	}

	// Discard out-of-band SSH requests (not supported by this server).
	go ssh.DiscardRequests(reqs)

	for newChannel := range chans {
		if err := serviceChannel(newChannel); err != nil {
			return err
		}
	}
	return nil
}

func serviceChannel(newChannel ssh.NewChannel) error {
	if newChannel.ChannelType() != "session" {
		newChannel.Reject(ssh.UnknownChannelType, "unknown channel type")
		return nil
	}
	channel, requests, err := newChannel.Accept()
	if err != nil {
		return fmt.Errorf("accepting channel: %v", err)
	}
	defer channel.Close()

	// Start the command with a pseudo-terminal.
	shell, err := pty.Start(exec.Command("bash"))
	if err != nil {
		return fmt.Errorf("starting pty: %v", err)
	}
	defer shell.Close()

	go func(in <-chan *ssh.Request) {
		for req := range in {
			// Tell the client we accept pty and shell commands
			req.Reply(req.Type == "pty-req" || req.Type == "shell", nil)

			if req.Type == "pty-req" {
				skip := binary.BigEndian.Uint32(req.Payload)
				resize(shell, req.Payload[4+skip:])
			}
			if req.Type == "window-change" {
				resize(shell, req.Payload)
			}
		}
	}(requests)

	go io.Copy(shell, channel)
	io.Copy(channel, shell)
	return nil
}

func resize(f *os.File, payload []byte) {
	var size struct {
		Width, Height uint32
	}
	if err := binary.Read(bytes.NewReader(payload), binary.BigEndian, &size); err == nil {
		pty.Setsize(f, &pty.Winsize{Cols: uint16(size.Width), Rows: uint16(size.Height)})
	}
}
<commit_msg>Fixed connection error through browser SSH (#85)<commit_after>package main

import (
	"bytes"
	"crypto/rand"
	"crypto/rsa"
	"encoding/binary"
	"errors"
	"flag"
	"fmt"
	"io"
	"log"
	"net"
	"os/exec"

	"github.com/googlegenomics/pipelines-tools/gce"
	"github.com/kr/pty"
	"golang.org/x/crypto/ssh"
)

var (
	port = flag.Uint("port", 22, "the port to listen on")
)

func main() {
	flag.Parse()

	config, listener, err := startServer(*port)
	if err != nil {
		log.Fatalf("Failed to start server: %v", err)
	}
	log.Printf("Listening on %s...", listener.Addr())
	defer listener.Close()

	for {
		connection, err := listener.Accept()
		if err != nil {
			log.Fatalf("Failed to accept incoming connection: %v", err)
		}

		go func() {
			if err := handleConnection(connection, config); err != nil {
				log.Printf("Failed to handle connection: %v", err)
			}
		}()
	}
}

func startServer(port uint) (*ssh.ServerConfig, net.Listener, error) {
	config, err := getConfiguration()
	if err != nil {
		return nil, nil, fmt.Errorf("getting configuration: %v", err)
	}

	listener, err := net.Listen("tcp", fmt.Sprintf(":%d", port))
	if err != nil {
		return nil, nil, fmt.Errorf("listen: %v", err)
	}
	return config, listener, nil
}

func getConfiguration() (*ssh.ServerConfig, error) {
	key, err := rsa.GenerateKey(rand.Reader, 2048)
	if err != nil {
		return nil, fmt.Errorf("generating server key pair: %v", err)
	}

	signer, err := ssh.NewSignerFromKey(key)
	if err != nil {
		return nil, fmt.Errorf("creating signer: %v", err)
	}

	config := &ssh.ServerConfig{
		PublicKeyCallback: func(conn ssh.ConnMetadata, key ssh.PublicKey) (*ssh.Permissions, error) {
			authorizedKeys, err := gce.GetAuthorizedKeys()
			if err != nil {
				return nil, fmt.Errorf("getting the authorized keys: %v", err)
			}

			if !authorizedKeys[string(key.Marshal())] {
				return nil, errors.New("unauthorized")
			}
			return nil, nil
		},
	}
	config.AddHostKey(signer)
	return config, nil
}

func handleConnection(conn net.Conn, serverConfig *ssh.ServerConfig) error {
	sConn, chans, reqs, err := ssh.NewServerConn(conn, serverConfig)
	if err != nil {
		return fmt.Errorf("handshake: %v", err)
	}
	defer sConn.Close()

	// Discard out-of-band SSH requests (not supported by this server).
	go ssh.DiscardRequests(reqs)

	for newChannel := range chans {
		go func(channel ssh.NewChannel) {
			if err := serviceChannel(channel); err != nil {
				log.Printf("Failed to service channel: %v", err)
			}
		}(newChannel)
	}
	return nil
}

func serviceChannel(newChannel ssh.NewChannel) error {
	if newChannel.ChannelType() != "session" {
		newChannel.Reject(ssh.UnknownChannelType, "unknown channel type")
		return nil
	}
	channel, requests, err := newChannel.Accept()
	if err != nil {
		return fmt.Errorf("accepting channel: %v", err)
	}
	defer channel.Close()

	allow := map[string]bool{"shell": true, "pty-req": true, "window-change": true}
	resize := make(chan *pty.Winsize, 1)
	defer close(resize)
	done := make(chan error)
	for {
		select {
		case err := <-done:
			if err != nil {
				return fmt.Errorf("running pty: %v", err)
			}
			return nil
		case req, ok := <-requests:
			if !ok {
				return nil
			}
			if req.WantReply {
				req.Reply(allow[req.Type], nil)
			}
			switch req.Type {
			case "pty-req":
				r := bytes.NewReader(req.Payload)
				term, err := readString(r)
				if err != nil {
					return fmt.Errorf("parsing TERM environment variable value: %v", err)
				}

				size, err := readWindowSize(r)
				if err != nil {
					return fmt.Errorf("reading window size: %v", err)
				}
				resize <- size

				go func() {
					done <- runPTY(channel, term, resize)
				}()
			case "window-change":
				size, err := readWindowSize(bytes.NewReader(req.Payload))
				if err != nil {
					return fmt.Errorf("reading window size: %v", err)
				}
				resize <- size
			}
		}
	}
}

func runPTY(channel ssh.Channel, term string, resize chan *pty.Winsize) error {
	// Start the command with a pseudo-terminal.
	cmd := exec.Command("bash")
	cmd.Env = append(cmd.Env, fmt.Sprintf("TERM=%s", term))
	shell, err := pty.Start(cmd)
	if err != nil {
		return fmt.Errorf("starting pty: %v", err)
	}
	defer shell.Close()

	go func() {
		for size := range resize {
			pty.Setsize(shell, size)
		}
	}()

	go io.Copy(shell, channel)
	io.Copy(channel, shell)

	status := struct{ Status uint32 }{}
	if _, err := channel.SendRequest("exit-status", false, ssh.Marshal(&status)); err != nil {
		return fmt.Errorf("sending exit status: %v", err)
	}
	return nil
}

func readWindowSize(r io.Reader) (*pty.Winsize, error) {
	var size struct {
		Width, Height uint32
	}
	if err := binary.Read(r, binary.BigEndian, &size); err != nil {
		return nil, err
	}

	return &pty.Winsize{Cols: uint16(size.Width), Rows: uint16(size.Height)}, nil
}

func readString(r io.Reader) (string, error) {
	var length uint32
	if err := binary.Read(r, binary.BigEndian, &length); err != nil {
		return "", fmt.Errorf("reading length: %v", err)
	}

	str := make([]byte, length)
	if _, err := io.ReadFull(r, str); err != nil {
		return "", fmt.Errorf("reading string: %v", err)
	}
	return string(str), nil
}
<|endoftext|>


<commit_before>package main

import (
	"flag"
	"log"
	"strings"
	"time"

	"golang.org/x/net/context"

	"github.com/coreos/etcd/client"
	"github.com/coreos/etcd/clientv3"
)

var etcdHosts string

func init() {
	flag.StringVar(&etcdHosts, "etcd-addr", "http://localhost:2379", "")
	flag.Parse()
}
func main() {
	cfg := clientv3.Config{Endpoints: strings.Split(etcdHosts, ",")}
	etcdcli, err := clientv3.New(cfg)
	if err != nil {
		panic(err)
	}
	v2cfg := client.Config{Endpoints: strings.Split(etcdHosts, ",")}
	v2cli, err := client.New(v2cfg)
	if err != nil {
		panic(err)
	}

	v2kv := client.NewKeysAPI(v2cli)

	key := "/"
	for {
		opts := []clientv3.OpOption{clientv3.WithLimit(100), clientv3.WithPrefix()}
		resp, err := etcdcli.KV.Get(context.TODO(), key, opts...)
		if err != nil {
			panic(err)
		}

		// copy to v2
		for _, kv := range resp.Kvs {
			if kv.Lease != 0 {
				v2kv.Set(context.TODO(), string(kv.Key), string(kv.Value), &client.SetOptions{
					TTL: 1 * time.Hour,
				})
				continue
			}
			v2kv.Set(context.TODO(), string(kv.Key), string(kv.Value), nil)
		}

		if !resp.More {
			log.Println("finished")
			return
		}

		// move to next key
		key = string(append(resp.Kvs[len(resp.Kvs)-1].Key, 0))
	}
}
<commit_msg>play with rollback<commit_after>package main

import (
	"flag"
	"fmt"
	"path"
	"time"

	"github.com/coreos/etcd/mvcc/backend"
	"github.com/coreos/etcd/mvcc/mvccpb"
)

var (
	migrateDatadir string
)

func init() {
	flag.StringVar(&migrateDatadir, "data-dir", "", "Path to the data directory")
	flag.Parse()
}

func main() {
	dbpath := path.Join(migrateDatadir, "member", "snap", "db")

	be := backend.New(dbpath, time.Second, 10000)
	tx := be.BatchTx()

	tx.Lock()
	tx.UnsafeForEach([]byte("key"), func(k, v []byte) error {
		kv := &mvccpb.KeyValue{}
		kv.Unmarshal(v)
		fmt.Printf("%s %d %d %d\n", string(kv.Key), kv.CreateRevision, kv.ModRevision, kv.Lease)
		return nil
	})
	tx.Unlock()

}
<|endoftext|>


<commit_before>package server

import (
	"bytes"
	"context"
	"fmt"
	"path"
	"time"

	"github.com/pachyderm/pachyderm/src/client"
	"github.com/pachyderm/pachyderm/src/client/pfs"
	"github.com/pachyderm/pachyderm/src/client/pkg/grpcutil"
	"github.com/pachyderm/pachyderm/src/server/pkg/backoff"
	"github.com/pachyderm/pachyderm/src/server/pkg/hashtree"
	"github.com/pachyderm/pachyderm/src/server/pkg/watch"
	workerpkg "github.com/pachyderm/pachyderm/src/server/pkg/worker"

	etcd "github.com/coreos/etcd/clientv3"
	"go.pedge.io/lion/proto"
	"google.golang.org/grpc"
)

const (
	workerEtcdPrefix = "workers"
)

type datum struct {
	files   []*pfs.FileInfo
	retries int
}

// WorkerPool represents a pool of workers that can be used to process datums.
type WorkerPool interface {
	// Send datums to this channel to be processed
	DataCh() chan<- *datum
	// Receive datums that failed to be processed from this channel
	FailCh() <-chan *datum
	// Receive hashtrees of the outputs of successfully processing datums
	SuccessCh() <-chan hashtree.HashTree
}

type workerPool struct {
	// When this context is canceled, the worker pool should clean up all
	// its resources.
	ctx context.Context
	// The prefix in etcd where new workers can be discovered
	workerDir string
	// workersMap is a map from a worker's address to the function that
	// can be used to release its resources.
	workersMap map[string]context.CancelFunc
	// objClient is the client for Pachyderm's object store
	objClient pfs.ObjectAPIClient
	// Used to check for workers added/deleted in etcd
	etcdClient *etcd.Client
	// The job that spawned the worker pool
	jobID string
	// workers get datums from this channel.
	dataCh chan *datum
	// workers send datums to this channel when they fail to process
	// the datums.
	failCh chan *datum
	// workers send the hashtrees of the outputs of processing datums to
	// this channel.
	successCh chan hashtree.HashTree
}

func (w *workerPool) discoverWorkers() {
	b := backoff.NewInfiniteBackOff()
	backoff.RetryNotify(func() error {
		protolion.Infof("watching `%s` for workers", w.workerDir)
		watcher, err := watch.NewWatcher(w.ctx, w.etcdClient, w.workerDir)
		if err != nil {
			return err
		}
		defer watcher.Close()
		for {
			resp, ok := <-watcher.Watch()
			if !ok {
				return fmt.Errorf("watcher closed for unknown reasons")
			}
			if err := resp.Err; err != nil {
				return err
			}
			addr := path.Base(string(resp.Key))
			switch resp.Type {
			case watch.EventPut:
				if err := w.addWorker(addr); err != nil {
					return err
				}
			case watch.EventDelete:
				if err := w.delWorker(addr); err != nil {
					return err
				}
			}
		}
		panic("unreachable")
	}, b, func(err error, d time.Duration) error {
		protolion.Errorf("error discovering workers for %v: %v; retrying in %v", w.workerDir, err, d)
		return nil
	})
}

func (w *workerPool) addWorker(addr string) error {
	if cancel, ok := w.workersMap[addr]; ok {
		cancel()
	}

	workerCtx, cancelFn := context.WithCancel(w.ctx)
	w.workersMap[addr] = cancelFn

	protolion.Infof("launching new worker for %s at %v", w.workerDir, addr)
	go w.runWorker(workerCtx, addr)
	return nil
}

func (w *workerPool) delWorker(addr string) error {
	cancel, ok := w.workersMap[addr]
	if !ok {
		return fmt.Errorf("deleting worker %s which is not in worker pool", addr)
	}
	cancel()
	protolion.Infof("deleting worker for %s at %v", w.workerDir, addr)
	return nil
}

func (w *workerPool) runWorker(ctx context.Context, addr string) {
	defer func() {
		protolion.Infof("goro for worker %s is exiting", addr)
	}()

	var workerClient workerpkg.WorkerClient
	b := backoff.NewInfiniteBackOff()
	backoff.RetryNotify(func() error {
		conn, err := grpc.DialContext(ctx, fmt.Sprintf("%s:%d", addr, client.PPSWorkerPort), grpc.WithInsecure())
		if err != nil {
			return err
		}
		workerClient = workerpkg.NewWorkerClient(conn)
		return nil
	}, b, func(err error, d time.Duration) error {
		protolion.Infof("error establishing connection with worker %s; retrying in %v", addr, d)
		return nil
	})

	for true {
		var dt *datum
		var ok bool
		select {
		case <-ctx.Done():
			return
		case dt, ok = <-w.dataCh:
			if !ok {
				return
			}
		}
		func() (retErr error) {
			defer func() {
				if retErr != nil {
					protolion.Errorf("%v", retErr)
					select {
					case w.failCh <- dt:
					case <-ctx.Done():
					}
				}
			}()
			resp, err := workerClient.Process(ctx, &workerpkg.ProcessRequest{
				JobID: w.jobID,
				Data:  dt.files,
			})
			if err != nil {
				return fmt.Errorf("worker %s failed to process datum %v with error %s", addr, dt.files, err)
			}
			if resp.Tag != nil {
				var buffer bytes.Buffer
				getTagClient, err := w.objClient.GetTag(ctx, &pfs.Tag{resp.Tag.Name})
				if err != nil {
					return fmt.Errorf("failed to retrieve hashtree after worker %s has ostensibly processed the datum %v: %v", addr, dt.files, err)
				}
				if err := grpcutil.WriteFromStreamingBytesClient(getTagClient, &buffer); err != nil {
					return fmt.Errorf("failed to retrieve hashtree after worker %s has ostensibly processed the datum %v: %v", addr, dt.files, err)
				}
				tree, err := hashtree.Deserialize(buffer.Bytes())
				if err != nil {
					return fmt.Errorf("failed to serialize hashtree after worker %s has ostensibly processed the datum %v; this is likely a bug: %v", addr, dt.files, err)
				}
				w.successCh <- tree
			} else if resp.Failed {
				return fmt.Errorf("user code failed to process datum %v", dt.files)
			} else {
				return fmt.Errorf("unrecognized response from worker %s when processing datum %v; this is likely a bug", addr, dt.files)
			}
			return nil
		}()
	}
}

func (w *workerPool) DataCh() chan<- *datum {
	return w.dataCh
}

func (w *workerPool) FailCh() <-chan *datum {
	return w.failCh
}

func (w *workerPool) SuccessCh() <-chan hashtree.HashTree {
	return w.successCh
}

// workerPool generates a new worker pool that talks to the replication
// controller identified by rcName.
// Each workerPool is supposed to be owned by a single job, identified
// by jobID.
func (a *apiServer) newWorkerPool(ctx context.Context, rcName string, jobID string) (WorkerPool, error) {
	objClient, err := a.getObjectClient()
	if err != nil {
		return nil, err
	}
	wp := &workerPool{
		ctx:        ctx,
		workerDir:  path.Join(a.etcdPrefix, workerEtcdPrefix, rcName),
		workersMap: make(map[string]context.CancelFunc),
		etcdClient: a.etcdClient,
		objClient:  objClient,
		jobID:      jobID,
		dataCh:     make(chan *datum),
		failCh:     make(chan *datum),
		successCh:  make(chan hashtree.HashTree),
	}
	// We need to make sure that the prefix ends with the trailing slash,
	// because
	if wp.workerDir[len(wp.workerDir)-1] != '/' {
		wp.workerDir += "/"
	}

	go wp.discoverWorkers()
	return wp, nil
}
<commit_msg>Make sure that goros exit when context gets canceled<commit_after>package server

import (
	"bytes"
	"context"
	"fmt"
	"path"
	"time"

	"github.com/pachyderm/pachyderm/src/client"
	"github.com/pachyderm/pachyderm/src/client/pfs"
	"github.com/pachyderm/pachyderm/src/client/pkg/grpcutil"
	"github.com/pachyderm/pachyderm/src/server/pkg/backoff"
	"github.com/pachyderm/pachyderm/src/server/pkg/hashtree"
	"github.com/pachyderm/pachyderm/src/server/pkg/watch"
	workerpkg "github.com/pachyderm/pachyderm/src/server/pkg/worker"

	etcd "github.com/coreos/etcd/clientv3"
	"go.pedge.io/lion/proto"
	"google.golang.org/grpc"
)

const (
	workerEtcdPrefix = "workers"
)

type datum struct {
	files   []*pfs.FileInfo
	retries int
}

// WorkerPool represents a pool of workers that can be used to process datums.
type WorkerPool interface {
	// Send datums to this channel to be processed
	DataCh() chan<- *datum
	// Receive datums that failed to be processed from this channel
	FailCh() <-chan *datum
	// Receive hashtrees of the outputs of successfully processing datums
	SuccessCh() <-chan hashtree.HashTree
}

type workerPool struct {
	// When this context is canceled, the worker pool should clean up all
	// its resources.
	ctx context.Context
	// The prefix in etcd where new workers can be discovered
	workerDir string
	// workersMap is a map from a worker's address to the function that
	// can be used to release its resources.
	workersMap map[string]context.CancelFunc
	// objClient is the client for Pachyderm's object store
	objClient pfs.ObjectAPIClient
	// Used to check for workers added/deleted in etcd
	etcdClient *etcd.Client
	// The job that spawned the worker pool
	jobID string
	// workers get datums from this channel.
	dataCh chan *datum
	// workers send datums to this channel when they fail to process
	// the datums.
	failCh chan *datum
	// workers send the hashtrees of the outputs of processing datums to
	// this channel.
	successCh chan hashtree.HashTree
}

func (w *workerPool) discoverWorkers() {
	b := backoff.NewInfiniteBackOff()
	backoff.RetryNotify(func() error {
		protolion.Infof("watching `%s` for workers", w.workerDir)
		watcher, err := watch.NewWatcher(w.ctx, w.etcdClient, w.workerDir)
		if err != nil {
			return err
		}
		defer watcher.Close()
		for {
			resp, ok := <-watcher.Watch()
			if !ok {
				return fmt.Errorf("watcher closed for unknown reasons")
			}
			if err := resp.Err; err != nil {
				return err
			}
			addr := path.Base(string(resp.Key))
			switch resp.Type {
			case watch.EventPut:
				if err := w.addWorker(addr); err != nil {
					return err
				}
			case watch.EventDelete:
				if err := w.delWorker(addr); err != nil {
					return err
				}
			}
		}
		panic("unreachable")
	}, b, func(err error, d time.Duration) error {
		select {
		case <-w.ctx.Done():
			// Exit the retry loop if context got cancelled
			return err
		default:
		}
		protolion.Errorf("error discovering workers for %v: %v; retrying in %v", w.workerDir, err, d)
		return nil
	})
}

func (w *workerPool) addWorker(addr string) error {
	if cancel, ok := w.workersMap[addr]; ok {
		cancel()
	}

	workerCtx, cancelFn := context.WithCancel(w.ctx)
	w.workersMap[addr] = cancelFn

	protolion.Infof("launching new worker for %s at %v", w.workerDir, addr)
	go w.runWorker(workerCtx, addr)
	return nil
}

func (w *workerPool) delWorker(addr string) error {
	cancel, ok := w.workersMap[addr]
	if !ok {
		return fmt.Errorf("deleting worker %s which is not in worker pool", addr)
	}
	cancel()
	protolion.Infof("deleting worker for %s at %v", w.workerDir, addr)
	return nil
}

func (w *workerPool) runWorker(ctx context.Context, addr string) {
	defer func() {
		protolion.Infof("goro for worker %s is exiting", addr)
	}()

	var workerClient workerpkg.WorkerClient
	b := backoff.NewInfiniteBackOff()
	backoff.RetryNotify(func() error {
		conn, err := grpc.DialContext(ctx, fmt.Sprintf("%s:%d", addr, client.PPSWorkerPort), grpc.WithInsecure())
		if err != nil {
			return err
		}
		workerClient = workerpkg.NewWorkerClient(conn)
		return nil
	}, b, func(err error, d time.Duration) error {
		select {
		case <-w.ctx.Done():
			// Exit the retry loop if context got cancelled
			return err
		default:
		}
		protolion.Infof("error establishing connection with worker %s; retrying in %v", addr, d)
		return nil
	})

	for true {
		var dt *datum
		var ok bool
		select {
		case <-ctx.Done():
			return
		case dt, ok = <-w.dataCh:
			if !ok {
				return
			}
		}
		func() (retErr error) {
			defer func() {
				if retErr != nil {
					protolion.Errorf("%v", retErr)
					select {
					case w.failCh <- dt:
					case <-ctx.Done():
					}
				}
			}()
			resp, err := workerClient.Process(ctx, &workerpkg.ProcessRequest{
				JobID: w.jobID,
				Data:  dt.files,
			})
			if err != nil {
				return fmt.Errorf("worker %s failed to process datum %v with error %s", addr, dt.files, err)
			}
			if resp.Tag != nil {
				var buffer bytes.Buffer
				getTagClient, err := w.objClient.GetTag(ctx, &pfs.Tag{resp.Tag.Name})
				if err != nil {
					return fmt.Errorf("failed to retrieve hashtree after worker %s has ostensibly processed the datum %v: %v", addr, dt.files, err)
				}
				if err := grpcutil.WriteFromStreamingBytesClient(getTagClient, &buffer); err != nil {
					return fmt.Errorf("failed to retrieve hashtree after worker %s has ostensibly processed the datum %v: %v", addr, dt.files, err)
				}
				tree, err := hashtree.Deserialize(buffer.Bytes())
				if err != nil {
					return fmt.Errorf("failed to serialize hashtree after worker %s has ostensibly processed the datum %v; this is likely a bug: %v", addr, dt.files, err)
				}
				w.successCh <- tree
			} else if resp.Failed {
				return fmt.Errorf("user code failed to process datum %v", dt.files)
			} else {
				return fmt.Errorf("unrecognized response from worker %s when processing datum %v; this is likely a bug", addr, dt.files)
			}
			return nil
		}()
	}
}

func (w *workerPool) DataCh() chan<- *datum {
	return w.dataCh
}

func (w *workerPool) FailCh() <-chan *datum {
	return w.failCh
}

func (w *workerPool) SuccessCh() <-chan hashtree.HashTree {
	return w.successCh
}

// workerPool generates a new worker pool that talks to the replication
// controller identified by rcName.
// Each workerPool is supposed to be owned by a single job, identified
// by jobID.
func (a *apiServer) newWorkerPool(ctx context.Context, rcName string, jobID string) (WorkerPool, error) {
	objClient, err := a.getObjectClient()
	if err != nil {
		return nil, err
	}
	wp := &workerPool{
		ctx:        ctx,
		workerDir:  path.Join(a.etcdPrefix, workerEtcdPrefix, rcName),
		workersMap: make(map[string]context.CancelFunc),
		etcdClient: a.etcdClient,
		objClient:  objClient,
		jobID:      jobID,
		dataCh:     make(chan *datum),
		failCh:     make(chan *datum),
		successCh:  make(chan hashtree.HashTree),
	}
	// We need to make sure that the prefix ends with the trailing slash,
	// because
	if wp.workerDir[len(wp.workerDir)-1] != '/' {
		wp.workerDir += "/"
	}

	go wp.discoverWorkers()
	return wp, nil
}
<|endoftext|>


<commit_before>package schema

import (
	"context"
	"fmt"
	"net/http"
	"strings"

	"hermes/models"

	"github.com/jinzhu/gorm"
	"github.com/labstack/echo"
)

type (
	entity struct {
		Table string
		Field string
	}

	field struct {
		Name string
		Eq   *Value
	}

	arguments struct {
		Field field
		Or    *[]field
		And   *[]field
	}

	Resolver struct{}
)

func errorResponse() error {
	return echo.NewHTTPError(http.StatusInternalServerError)
}

func (r *Resolver) Count(context context.Context, args arguments) (int32, error) {
	var total int32

	if db, castOk := context.Value(DB).(*gorm.DB); castOk {
		query := args.Field.query(db)
		errorList := query.GetErrors()

		if !(len(errorList) > 0 || query.Error != nil || query.Value == nil) {
			query.Count(&total)
		}
	}

	// TODO: Handle non existent db

	return total, nil
}

func (r *Resolver) Average(context context.Context, args arguments) (float64, error) {
	var total float64

	if db, castOk := context.Value(DB).(*gorm.DB); castOk {
		entity := args.Field.getEntity()
		average := fmt.Sprintf("AVG(%s)", entity.Field)
		query := args.Field.query(db).Select(average)
		errorList := query.GetErrors()

		if !(len(errorList) > 0 || query.Error != nil || query.Value == nil) {
			query.Row().Scan(&total)
		}
	}

	return total, nil
}

func (f *field) query(db *gorm.DB) *gorm.DB {
	operator := f.resolveOperator()
	entity := f.getEntity()
	value := f.getValue()
	where := fmt.Sprintf("%s %s ?", entity.Field, operator)

	switch entity.Table {
	case "apps":
		return db.Model(&models.Rating{}).Where(where, value)
	case "appusers":
		return db.Model(&models.AppUser{}).Where(where, value)
	case "brands":
		return db.Model(&models.Brand{}).Where(where, value)
	case "browsers":
		return db.Model(&models.Browser{}).Where(where, value)
	case "devices":
		return db.Model(&models.Device{}).Where(where, value)
	case "messages":
		return db.Model(&models.Message{}).Where(where, value)
	case "platforms":
		return db.Model(&models.Platform{}).Where(where, value)
	case "ranges":
		return db.Model(&models.Range{}).Where(where, value)
	case "ratings":
		fallthrough
	default:
		return db.Model(&models.Rating{}).Where(where, value)
	}
}

func (f *field) getEntity() entity {
	splitField := strings.Split(f.Name, ".")

	return entity{Table: splitField[0], Field: splitField[1]}
}

func (f *field) getValue() interface{} {
	if f.Eq != nil {
		if f.Eq.String != nil {
			return f.Eq.String
		} else if f.Eq.Int != nil {
			return f.Eq.Int
		} else if f.Eq.Float != nil {
			return f.Eq.Float
		} else if f.Eq.Bool != nil {
			return f.Eq.Bool
		}
	}

	return nil
}

func (f *field) resolveOperator() string {
	value := f.getValue()

	if f.Eq != nil {
		switch value.(type) {
		case string:
			return "LIKE"
		default:
			return "="
		}
	}

	return "="
}
<commit_msg>Impletented count OR<commit_after>package schema

import (
	"context"
	"errors"
	"fmt"
	"net/http"
	"strings"

	"hermes/models"

	"github.com/jinzhu/gorm"
	"github.com/labstack/echo"
)

type (
	entity struct {
		Table string
		Field string
	}

	field struct {
		Name string
		Eq   *Value
	}

	arguments struct {
		Field field
		Or    *[]field
		And   *[]field
	}

	Resolver struct{}
)

func errorResponse() error {
	return echo.NewHTTPError(http.StatusInternalServerError)
}

func (r *Resolver) Count(context context.Context, args arguments) (int32, error) {
	var total int32

	if db, castOk := context.Value(DB).(*gorm.DB); castOk {
		query := args.Field.query(db)
		errorList := query.GetErrors()

		if !(len(errorList) > 0 || query.Error != nil || query.Value == nil) {
			r.count(db, args, query, &total)

			return total, nil
		}

		return total, errors.New("Could not get value from database")
	}

	return total, errors.New("Could not connect to database")
}

func (r *Resolver) Average(context context.Context, args arguments) (float64, error) {
	var total float64

	if db, castOk := context.Value(DB).(*gorm.DB); castOk {
		entity := args.Field.getEntity()
		average := fmt.Sprintf("AVG(%s)", entity.Field)
		query := args.Field.query(db).Select(average)
		errorList := query.GetErrors()

		if !(len(errorList) > 0 || query.Error != nil || query.Value == nil) {
			r.average(db, args, query, &total)

			return total, nil
		}

		return total, errors.New("Could not get value from database")
	}

	return total, errors.New("Could not connect to database")
}

func (r *Resolver) count(db *gorm.DB, args arguments, query *gorm.DB, total *int32) {
	if args.And == nil && args.Or == nil {
		query.Count(total)
	} else if args.And != nil {
		*total = 9 // Placeholder value
	} else if args.Or != nil {
		for _, item := range *args.Or {
			var subtotal int32

			subquery := item.query(db)
			errorList := subquery.GetErrors()

			if !(len(errorList) > 0 || subquery.Error != nil || subquery.Value == nil) {
				subquery.Count(&subtotal)

				*total += subtotal
			}
		}
	}
}

func (r *Resolver) average(db *gorm.DB, args arguments, query *gorm.DB, total *float64) {
	if args.And == nil && args.Or == nil {
		query.Row().Scan(total)
	} else if args.And != nil {
		*total = 8 // Placeholder value
	} else if args.Or != nil {
		*total = 9 // Placeholder value
	}
}

func (f *field) query(db *gorm.DB) *gorm.DB {
	operator := f.resolveOperator()
	entity := f.getEntity()
	value := f.getValue()
	where := fmt.Sprintf("%s %s ?", entity.Field, operator)

	switch entity.Table {
	case "apps":
		return db.Model(&models.Rating{}).Where(where, value)
	case "appusers":
		return db.Model(&models.AppUser{}).Where(where, value)
	case "brands":
		return db.Model(&models.Brand{}).Where(where, value)
	case "browsers":
		return db.Model(&models.Browser{}).Where(where, value)
	case "devices":
		return db.Model(&models.Device{}).Where(where, value)
	case "messages":
		return db.Model(&models.Message{}).Where(where, value)
	case "platforms":
		return db.Model(&models.Platform{}).Where(where, value)
	case "ranges":
		return db.Model(&models.Range{}).Where(where, value)
	case "ratings":
		fallthrough
	default:
		return db.Model(&models.Rating{}).Where(where, value)
	}
}

func (f *field) getEntity() entity {
	splitField := strings.Split(f.Name, ".")

	return entity{Table: splitField[0], Field: splitField[1]}
}

func (f *field) getValue() interface{} {
	if f.Eq != nil {
		if f.Eq.String != nil {
			return f.Eq.String
		} else if f.Eq.Int != nil {
			return f.Eq.Int
		} else if f.Eq.Float != nil {
			return f.Eq.Float
		} else if f.Eq.Bool != nil {
			return f.Eq.Bool
		}
	}

	return nil
}

func (f *field) resolveOperator() string {
	value := f.getValue()

	if f.Eq != nil {
		switch value.(type) {
		case string:
			return "LIKE"
		default:
			return "="
		}
	}

	return "="
}
<|endoftext|>


<commit_before>package tests

import (
	"bytes"
	"errors"
	"fmt"
	"io"
	"os"
	"os/exec"
	"path"
	"time"

	"github.com/unigornel/unigornel/integration_tests/tests"
	"github.com/unigornel/unigornel/integration_tests/xen"
)

type SimpleTest struct {
	Name        string
	Category    string
	Package     string
	Memory      int
	Stdin       []byte
	Timeout     time.Duration
	CanCrash    bool
	CanShutdown bool
	CanTimeout  bool
	CheckRun    func(string) error

	unikernel  string
	domain     *xen.Domain
	didTimeout bool
	output     string
}

func SimpleTestPackage(arg ...string) string {
	s := append([]string{"github.com/unigornel/unigornel/integration_tests/tests"}, arg...)
	return path.Join(s...)
}

func (t *SimpleTest) GetName() string {
	return t.Name
}

func (t *SimpleTest) GetCategory() string {
	return t.Category
}

func (t *SimpleTest) Build(w io.Writer) error {
	file, err := Build(w, t.Name, t.Package)
	t.unikernel = file
	return err
}

func (t *SimpleTest) Setup(w io.Writer) error {
	kernel := xen.Kernel{
		Binary:  t.unikernel,
		Memory:  t.Memory,
		Name:    t.Name,
		OnCrash: xen.OnCrashPreserve,
	}

	fmt.Fprintln(w, "[+] creating paused kernel")
	dom, err := kernel.CreatePausedUniqueName(func(cmd *exec.Cmd) {
		cmd.Stdout = w
		cmd.Stderr = w
	})
	fmt.Fprintln(w, "[+] domain created:", dom)
	t.domain = dom
	return err
}

func (t *SimpleTest) Run(w io.Writer) error {
	out := bytes.NewBuffer(nil)
	console := t.domain.Console()
	console.Stdout = io.MultiWriter(w, out)
	console.Stderr = console.Stdout
	stdin, err := console.StdinPipe()
	if err != nil {
		return err
	}

	fmt.Fprintln(w, "[+] attching to the console")
	if err := console.Start(); err != nil {
		return err
	}

	if t.Stdin != nil {
		fmt.Fprintln(w, "[+] writing to console")
		if _, err := stdin.Write(t.Stdin); err != nil {
			return err
		}
	}

	done := make(chan struct{})
	exited := make(chan error)
	timeout := make(chan struct{})
	go func() {
		console.Wait()
		close(done)
	}()
	go func() {
		for {
			select {
			case <-done:
				return
			case <-timeout:
				return
			case <-time.After(1 * time.Second):
				c, err := t.domain.Update()
				if err != nil {
					exited <- err
					close(exited)
					return
				}
				if c.State.Check(xen.DomainStateShutdown) || c.State.Check(xen.DomainStateCrashed) {
					close(exited)
					return
				}
			}
		}
	}()
	go func() {
		time.Sleep(t.Timeout)
		close(timeout)
	}()

	fmt.Fprintln(w, "[+] unpausing unikernel domain")
	if err := t.domain.Unpause().Run(); err != nil {
		console.Process.Kill()
		return err
	}

	select {
	case <-done:
		return errors.New("console unexpectedly exited")
	case err := <-exited:
		// Make sure the console can catch up
		time.Sleep(1 * time.Second)
		console.Process.Kill()
		if err != nil {
			return err
		}
		t.didTimeout = false
		<-done
	case <-timeout:
		t.didTimeout = true
		console.Process.Kill()
		<-done
	}
	t.output = string(out.Bytes())

	if t.didTimeout && !t.CanTimeout {
		return errors.New("unikernel timed out")
	}

	domain, err := t.domain.Update()
	if err != nil {
		return errors.New("domain not preserved: " + err.Error())
	}

	if domain.State.Check(xen.DomainStateCrashed) && !t.CanCrash {
		return errors.New("domain crashed")
	}

	if domain.State.Check(xen.DomainStateShutdown) && !t.CanShutdown {
		return errors.New("domain shutdown")
	}
	return nil
}

func (t *SimpleTest) Check(w io.Writer) error {
	if t.CheckRun == nil {
		fmt.Fprintln(w, "[+] no output checks specified")
		return nil
	}
	fmt.Fprintln(w, "[+] checking unikernel output")
	if err := t.CheckRun(t.output); err != nil {
		fmt.Fprintf(w, "[-] check error: %v\n", err)
		return err
	}
	return nil
}

func (t *SimpleTest) Clean(w io.Writer, success bool) (err error) {
	if t.unikernel != "" {
		fmt.Fprintf(w, "[+] removing %s\n", t.unikernel)
		os.Remove(t.unikernel)
	}

	if t.domain != nil {
		cmd := t.domain.Destroy()
		cmd.Stderr = w
		cmd.Stdout = w
		err = cmd.Run()
	}

	return
}
<commit_msg>Fix unused import<commit_after>package tests

import (
	"bytes"
	"errors"
	"fmt"
	"io"
	"os"
	"os/exec"
	"path"
	"time"

	"github.com/unigornel/unigornel/integration_tests/xen"
)

type SimpleTest struct {
	Name        string
	Category    string
	Package     string
	Memory      int
	Stdin       []byte
	Timeout     time.Duration
	CanCrash    bool
	CanShutdown bool
	CanTimeout  bool
	CheckRun    func(string) error

	unikernel  string
	domain     *xen.Domain
	didTimeout bool
	output     string
}

func SimpleTestPackage(arg ...string) string {
	s := append([]string{"github.com/unigornel/unigornel/integration_tests/tests"}, arg...)
	return path.Join(s...)
}

func (t *SimpleTest) GetName() string {
	return t.Name
}

func (t *SimpleTest) GetCategory() string {
	return t.Category
}

func (t *SimpleTest) Build(w io.Writer) error {
	file, err := Build(w, t.Name, t.Package)
	t.unikernel = file
	return err
}

func (t *SimpleTest) Setup(w io.Writer) error {
	kernel := xen.Kernel{
		Binary:  t.unikernel,
		Memory:  t.Memory,
		Name:    t.Name,
		OnCrash: xen.OnCrashPreserve,
	}

	fmt.Fprintln(w, "[+] creating paused kernel")
	dom, err := kernel.CreatePausedUniqueName(func(cmd *exec.Cmd) {
		cmd.Stdout = w
		cmd.Stderr = w
	})
	fmt.Fprintln(w, "[+] domain created:", dom)
	t.domain = dom
	return err
}

func (t *SimpleTest) Run(w io.Writer) error {
	out := bytes.NewBuffer(nil)
	console := t.domain.Console()
	console.Stdout = io.MultiWriter(w, out)
	console.Stderr = console.Stdout
	stdin, err := console.StdinPipe()
	if err != nil {
		return err
	}

	fmt.Fprintln(w, "[+] attching to the console")
	if err := console.Start(); err != nil {
		return err
	}

	if t.Stdin != nil {
		fmt.Fprintln(w, "[+] writing to console")
		if _, err := stdin.Write(t.Stdin); err != nil {
			return err
		}
	}

	done := make(chan struct{})
	exited := make(chan error)
	timeout := make(chan struct{})
	go func() {
		console.Wait()
		close(done)
	}()
	go func() {
		for {
			select {
			case <-done:
				return
			case <-timeout:
				return
			case <-time.After(1 * time.Second):
				c, err := t.domain.Update()
				if err != nil {
					exited <- err
					close(exited)
					return
				}
				if c.State.Check(xen.DomainStateShutdown) || c.State.Check(xen.DomainStateCrashed) {
					close(exited)
					return
				}
			}
		}
	}()
	go func() {
		time.Sleep(t.Timeout)
		close(timeout)
	}()

	fmt.Fprintln(w, "[+] unpausing unikernel domain")
	if err := t.domain.Unpause().Run(); err != nil {
		console.Process.Kill()
		return err
	}

	select {
	case <-done:
		return errors.New("console unexpectedly exited")
	case err := <-exited:
		// Make sure the console can catch up
		time.Sleep(1 * time.Second)
		console.Process.Kill()
		if err != nil {
			return err
		}
		t.didTimeout = false
		<-done
	case <-timeout:
		t.didTimeout = true
		console.Process.Kill()
		<-done
	}
	t.output = string(out.Bytes())

	if t.didTimeout && !t.CanTimeout {
		return errors.New("unikernel timed out")
	}

	domain, err := t.domain.Update()
	if err != nil {
		return errors.New("domain not preserved: " + err.Error())
	}

	if domain.State.Check(xen.DomainStateCrashed) && !t.CanCrash {
		return errors.New("domain crashed")
	}

	if domain.State.Check(xen.DomainStateShutdown) && !t.CanShutdown {
		return errors.New("domain shutdown")
	}
	return nil
}

func (t *SimpleTest) Check(w io.Writer) error {
	if t.CheckRun == nil {
		fmt.Fprintln(w, "[+] no output checks specified")
		return nil
	}
	fmt.Fprintln(w, "[+] checking unikernel output")
	if err := t.CheckRun(t.output); err != nil {
		fmt.Fprintf(w, "[-] check error: %v\n", err)
		return err
	}
	return nil
}

func (t *SimpleTest) Clean(w io.Writer, success bool) (err error) {
	if t.unikernel != "" {
		fmt.Fprintf(w, "[+] removing %s\n", t.unikernel)
		os.Remove(t.unikernel)
	}

	if t.domain != nil {
		cmd := t.domain.Destroy()
		cmd.Stderr = w
		cmd.Stdout = w
		err = cmd.Run()
	}

	return
}
<|endoftext|>


<commit_before>package self_test

import (
	"context"
	"fmt"
	"io/ioutil"
	"net"
	"time"

	quic "github.com/lucas-clemente/quic-go"
	quicproxy "github.com/lucas-clemente/quic-go/integrationtests/tools/proxy"
	"github.com/lucas-clemente/quic-go/internal/protocol"

	. "github.com/onsi/ginkgo"
	. "github.com/onsi/gomega"
)

var _ = Describe("non-zero RTT", func() {
	for _, v := range protocol.SupportedVersions {
		version := v

		Context(fmt.Sprintf("with QUIC version %s", version), func() {
			roundTrips := [...]time.Duration{
				10 * time.Millisecond,
				50 * time.Millisecond,
				100 * time.Millisecond,
				200 * time.Millisecond,
			}

			for _, r := range roundTrips {
				rtt := r

				It(fmt.Sprintf("downloads a message with %s RTT", rtt), func() {
					ln, err := quic.ListenAddr(
						"localhost:0",
						getTLSConfig(),
						&quic.Config{
							Versions: []protocol.VersionNumber{version},
						},
					)
					Expect(err).ToNot(HaveOccurred())
					done := make(chan struct{})
					go func() {
						defer GinkgoRecover()
						sess, err := ln.Accept(context.Background())
						Expect(err).ToNot(HaveOccurred())
						str, err := sess.OpenStream()
						Expect(err).ToNot(HaveOccurred())
						_, err = str.Write(PRData)
						Expect(err).ToNot(HaveOccurred())
						str.Close()
						close(done)
					}()
					serverPort := ln.Addr().(*net.UDPAddr).Port
					proxy, err := quicproxy.NewQuicProxy("localhost:0", &quicproxy.Opts{
						RemoteAddr: fmt.Sprintf("localhost:%d", serverPort),
						DelayPacket: func(quicproxy.Direction, []byte) time.Duration {
							return rtt / 2
						},
					})
					Expect(err).ToNot(HaveOccurred())
					defer proxy.Close()

					sess, err := quic.DialAddr(
						fmt.Sprintf("localhost:%d", proxy.LocalPort()),
						getTLSClientConfig(),
						&quic.Config{Versions: []protocol.VersionNumber{version}},
					)
					Expect(err).ToNot(HaveOccurred())
					str, err := sess.AcceptStream(context.Background())
					Expect(err).ToNot(HaveOccurred())
					data, err := ioutil.ReadAll(str)
					Expect(err).ToNot(HaveOccurred())
					Expect(data).To(Equal(PRData))
					sess.CloseWithError(0, "")
					Eventually(done).Should(BeClosed())
				})
			}
		})
	}
})
<commit_msg>add an integration test that introduces packet reordering<commit_after>package self_test

import (
	"context"
	"fmt"
	"io/ioutil"
	"net"
	"time"

	quic "github.com/lucas-clemente/quic-go"
	quicproxy "github.com/lucas-clemente/quic-go/integrationtests/tools/proxy"
	"github.com/lucas-clemente/quic-go/internal/protocol"

	. "github.com/onsi/ginkgo"
	. "github.com/onsi/gomega"
)

var _ = Describe("non-zero RTT", func() {
	for _, v := range protocol.SupportedVersions {
		version := v

		runServer := func() quic.Listener {
			ln, err := quic.ListenAddr(
				"localhost:0",
				getTLSConfig(),
				&quic.Config{
					Versions: []protocol.VersionNumber{version},
				},
			)
			Expect(err).ToNot(HaveOccurred())
			go func() {
				defer GinkgoRecover()
				sess, err := ln.Accept(context.Background())
				Expect(err).ToNot(HaveOccurred())
				str, err := sess.OpenStream()
				Expect(err).ToNot(HaveOccurred())
				_, err = str.Write(PRData)
				Expect(err).ToNot(HaveOccurred())
				str.Close()
			}()
			return ln
		}

		downloadFile := func(port int) {
			sess, err := quic.DialAddr(
				fmt.Sprintf("localhost:%d", port),
				getTLSClientConfig(),
				&quic.Config{Versions: []protocol.VersionNumber{version}},
			)
			Expect(err).ToNot(HaveOccurred())
			str, err := sess.AcceptStream(context.Background())
			Expect(err).ToNot(HaveOccurred())
			data, err := ioutil.ReadAll(str)
			Expect(err).ToNot(HaveOccurred())
			Expect(data).To(Equal(PRData))
			sess.CloseWithError(0, "")
		}

		Context(fmt.Sprintf("with QUIC version %s", version), func() {
			for _, r := range [...]time.Duration{
				10 * time.Millisecond,
				50 * time.Millisecond,
				100 * time.Millisecond,
				200 * time.Millisecond,
			} {
				rtt := r

				It(fmt.Sprintf("downloads a message with %s RTT", rtt), func() {
					ln := runServer()
					defer ln.Close()
					serverPort := ln.Addr().(*net.UDPAddr).Port
					proxy, err := quicproxy.NewQuicProxy("localhost:0", &quicproxy.Opts{
						RemoteAddr: fmt.Sprintf("localhost:%d", serverPort),
						DelayPacket: func(quicproxy.Direction, []byte) time.Duration {
							return rtt / 2
						},
					})
					Expect(err).ToNot(HaveOccurred())
					defer proxy.Close()

					sess, err := quic.DialAddr(
						fmt.Sprintf("localhost:%d", proxy.LocalPort()),
						getTLSClientConfig(),
						&quic.Config{Versions: []protocol.VersionNumber{version}},
					)
					Expect(err).ToNot(HaveOccurred())
					str, err := sess.AcceptStream(context.Background())
					Expect(err).ToNot(HaveOccurred())
					data, err := ioutil.ReadAll(str)
					Expect(err).ToNot(HaveOccurred())
					Expect(data).To(Equal(PRData))
					sess.CloseWithError(0, "")
				})
			}

			for _, r := range [...]time.Duration{
				10 * time.Millisecond,
				40 * time.Millisecond,
			} {
				rtt := r

				It(fmt.Sprintf("downloads a message with %s RTT, with reordering", rtt), func() {
					ln := runServer()
					defer ln.Close()
					serverPort := ln.Addr().(*net.UDPAddr).Port
					proxy, err := quicproxy.NewQuicProxy("localhost:0", &quicproxy.Opts{
						RemoteAddr: fmt.Sprintf("localhost:%d", serverPort),
						DelayPacket: func(quicproxy.Direction, []byte) time.Duration {
							return randomDuration(rtt/2, rtt*3/2) / 2
						},
					})
					Expect(err).ToNot(HaveOccurred())
					defer proxy.Close()

					downloadFile(proxy.LocalPort())
				})
			}
		})
	}
})
<|endoftext|>


<commit_before>// +build linux

package cgmgr

import (
	"fmt"
	"path"
	"path/filepath"
	"strings"

	"github.com/containers/libpod/pkg/cgroups"
	rspec "github.com/opencontainers/runtime-spec/specs-go"
	"github.com/pkg/errors"
	"github.com/sirupsen/logrus"
)

type CgroupfsManager struct{}

const (
	defaultCgroupfsParent = "/crio"
)

// Name returns the name of the cgroup manager (cgroupfs)
func (*CgroupfsManager) Name() string {
	return cgroupfsCgroupManager
}

// IsSystemd returns that this is not a systemd cgroup manager
func (*CgroupfsManager) IsSystemd() bool {
	return false
}

// ContainerCgroupPath takes arguments sandbox parent cgroup and container ID and returns
// the cgroup path for that containerID. If parentCgroup is empty, it
// uses the default parent /crio
func (*CgroupfsManager) ContainerCgroupPath(sbParent, containerID string) string {
	parent := defaultCgroupfsParent
	if sbParent != "" {
		parent = sbParent
	}
	return filepath.Join(parent, crioPrefix+"-"+containerID)
}

// ContainerCgroupAbsolutePath just calls ContainerCgroupPath,
// because they both return the absolute path
func (m *CgroupfsManager) ContainerCgroupAbsolutePath(sbParent, containerID string) (string, error) {
	return m.ContainerCgroupPath(sbParent, containerID), nil
}

// SandboxCgroupPath takes the sandbox parent, and sandbox ID. It
// returns the cgroup parent, cgroup path, and error.
func (*CgroupfsManager) SandboxCgroupPath(sbParent, sbID string) (cgParent, cgPath string, _ error) {
	if strings.HasSuffix(path.Base(sbParent), ".slice") {
		return "", "", fmt.Errorf("cri-o configured with cgroupfs cgroup manager, but received systemd slice as parent: %s", sbParent)
	}
	return sbParent, filepath.Join(sbParent, crioPrefix+"-"+sbID), nil
}

// MoveConmonToCgroup takes the container ID, cgroup parent, conmon's cgroup (from the config) and conmon's PID
// It attempts to move conmon to the correct cgroup.
// It returns the cgroupfs parent that conmon was put into
// so that CRI-O can clean the cgroup path of the newly added conmon once the process terminates (systemd handles this for us)
func (*CgroupfsManager) MoveConmonToCgroup(cid, cgroupParent, conmonCgroup string, pid int) (string, error) {
	if conmonCgroup != "pod" && conmonCgroup != "" {
		return "", errors.Errorf("conmon cgroup %s invalid for cgroupfs", conmonCgroup)
	}

	cgroupPath := fmt.Sprintf("%s/crio-conmon-%s", cgroupParent, cid)
	control, err := cgroups.New(cgroupPath, &rspec.LinuxResources{})
	if err != nil {
		logrus.Warnf("Failed to add conmon to cgroupfs sandbox cgroup: %v", err)
	}
	if control == nil {
		return cgroupPath, nil
	}

	// Record conmon's cgroup path in the container, so we can properly
	// clean it up when removing the container.
	// Here we should defer a crio-connmon- cgroup hierarchy deletion, but it will
	// always fail as conmon's pid is still there.
	// Fortunately, kubelet takes care of deleting this for us, so the leak will
	// only happens in corner case where one does a manual deletion of the container
	// through e.g. runc. This should be handled by implementing a conmon monitoring
	// routine that does the cgroup cleanup once conmon is terminated.
	if err := control.AddPid(pid); err != nil {
		return "", errors.Wrapf(err, "Failed to add conmon to cgroupfs sandbox cgroup")
	}
	return cgroupPath, nil
}
<commit_msg>Fix container cgroup under cgroupfs<commit_after>// +build linux

package cgmgr

import (
	"fmt"
	"path"
	"path/filepath"
	"strings"

	"github.com/containers/libpod/pkg/cgroups"
	rspec "github.com/opencontainers/runtime-spec/specs-go"
	"github.com/pkg/errors"
	"github.com/sirupsen/logrus"
)

type CgroupfsManager struct{}

const (
	defaultCgroupfsParent = "/crio"
)

// Name returns the name of the cgroup manager (cgroupfs)
func (*CgroupfsManager) Name() string {
	return cgroupfsCgroupManager
}

// IsSystemd returns that this is not a systemd cgroup manager
func (*CgroupfsManager) IsSystemd() bool {
	return false
}

// ContainerCgroupPath takes arguments sandbox parent cgroup and container ID and returns
// the cgroup path for that containerID. If parentCgroup is empty, it
// uses the default parent /crio
func (*CgroupfsManager) ContainerCgroupPath(sbParent, containerID string) string {
	parent := defaultCgroupfsParent
	if sbParent != "" {
		parent = sbParent
	}
	return filepath.Join("/", parent, crioPrefix+"-"+containerID)
}

// ContainerCgroupAbsolutePath just calls ContainerCgroupPath,
// because they both return the absolute path
func (m *CgroupfsManager) ContainerCgroupAbsolutePath(sbParent, containerID string) (string, error) {
	return m.ContainerCgroupPath(sbParent, containerID), nil
}

// SandboxCgroupPath takes the sandbox parent, and sandbox ID. It
// returns the cgroup parent, cgroup path, and error.
func (*CgroupfsManager) SandboxCgroupPath(sbParent, sbID string) (cgParent, cgPath string, _ error) {
	if strings.HasSuffix(path.Base(sbParent), ".slice") {
		return "", "", fmt.Errorf("cri-o configured with cgroupfs cgroup manager, but received systemd slice as parent: %s", sbParent)
	}
	return sbParent, filepath.Join(sbParent, crioPrefix+"-"+sbID), nil
}

// MoveConmonToCgroup takes the container ID, cgroup parent, conmon's cgroup (from the config) and conmon's PID
// It attempts to move conmon to the correct cgroup.
// It returns the cgroupfs parent that conmon was put into
// so that CRI-O can clean the cgroup path of the newly added conmon once the process terminates (systemd handles this for us)
func (*CgroupfsManager) MoveConmonToCgroup(cid, cgroupParent, conmonCgroup string, pid int) (string, error) {
	if conmonCgroup != "pod" && conmonCgroup != "" {
		return "", errors.Errorf("conmon cgroup %s invalid for cgroupfs", conmonCgroup)
	}

	cgroupPath := fmt.Sprintf("%s/crio-conmon-%s", cgroupParent, cid)
	control, err := cgroups.New(cgroupPath, &rspec.LinuxResources{})
	if err != nil {
		logrus.Warnf("Failed to add conmon to cgroupfs sandbox cgroup: %v", err)
	}
	if control == nil {
		return cgroupPath, nil
	}

	// Record conmon's cgroup path in the container, so we can properly
	// clean it up when removing the container.
	// Here we should defer a crio-connmon- cgroup hierarchy deletion, but it will
	// always fail as conmon's pid is still there.
	// Fortunately, kubelet takes care of deleting this for us, so the leak will
	// only happens in corner case where one does a manual deletion of the container
	// through e.g. runc. This should be handled by implementing a conmon monitoring
	// routine that does the cgroup cleanup once conmon is terminated.
	if err := control.AddPid(pid); err != nil {
		return "", errors.Wrapf(err, "Failed to add conmon to cgroupfs sandbox cgroup")
	}
	return cgroupPath, nil
}
<|endoftext|>