<commit_before>package carto

import (
	"fmt"
	"log"
	"runtime"

	"sync"

	"github.com/mathuin/gdal"
	"github.com/mathuin/terroir/world"
)

// JMT: convert to XZ, biome value, and list 0->n of points
type Column struct {
	xz     world.XZ
	biome  string
	blocks []string
}

func (r Region) genFeatures(in chan Feature) {
	ds, err := gdal.Open(r.mapfile, gdal.ReadOnly)
	if err != nil {
		panic(err)
	}
	if Debug {
		datasetInfo(ds, "genFeatures Input")
	}
	inx := ds.RasterXSize()
	iny := ds.RasterYSize()
	srs := gdal.CreateSpatialReference(ds.Projection())
	bufferLen := inx * iny

	lcarr := make([]int16, bufferLen)
	lcBand := ds.RasterBand(Landcover)
	lcrerr := lcBand.IO(gdal.Read, 0, 0, inx, iny, lcarr, inx, iny, 0, 0)
	if notnil(lcrerr) {
		panic(lcrerr)
	}

	// shapefile driver
	outdrv := gdal.OGRDriverByName("Memory")
	outDS, ok := outdrv.Create("out", nil)
	if !ok {
		panic(fmt.Errorf("OGR Driver Create Fail"))
	}
	outLayer := outDS.CreateLayer("polygons", srs, gdal.GT_Polygon, nil)

	// field definition
	outField := gdal.CreateFieldDefinition("lc", gdal.FT_Integer)
	outLayer.CreateField(outField, false)
	field := 0

	// options!
	options := []string{}

	// do it!
	err = lcBand.Polygonize(lcBand, outLayer, field, options, gdal.DummyProgress, nil)
	if notnil(err) {
		panic(err)
	}

	// iterate over features
	fc, ok := outLayer.FeatureCount(true)
	if !ok {
		panic(fmt.Errorf("outLayer.FeatureCount NOT OK"))
	}
	if Debug {
		log.Print("outLayer.FeatureCount(true): ", fc)
	}
	outLayer.ResetReading()
	for i := 0; i < fc; i++ {
		in <- Feature{outLayer.NextFeature()}
	}
	close(in)
}

func (r *Region) BuildWorld() (*world.World, error) {
	w := world.MakeWorld(r.name)
	w.SetRandomSeed(0)
	// JMT: need a sane storage location for files
	w.SetSaveDir(".")
	spawnpt := world.MakePoint(0, 0, 0)

	in := make(chan Feature)
	out := make(chan Column)

	var wg sync.WaitGroup

	numWorkers := runtime.NumCPU()
	// if Debug {
	// 	log.Print("debug mode - only starting one worker")
	// 	numWorkers = 1
	// }

	for i := 0; i < numWorkers; i++ {
		wg.Add(1)
		go func(i int) {
			defer wg.Done()
			r.processFeatures(in, out, i)
		}(i)
	}
	go func() { wg.Wait(); close(out) }()
	go r.genFeatures(in)

	columncount := 0
	for column := range out {
		columncount++

		b, ok := world.Biome[column.biome]
		if !ok {
			err := fmt.Errorf("biome %s not in world.Biome", column.biome)
			panic(err)
		}
		w.SetBiome(column.xz, byte(b))

		for k, v := range column.blocks {
			pt := world.Point{X: column.xz.X, Y: int32(k), Z: column.xz.Z}
			b, err := world.BlockNamed(v)
			if err != nil {
				panic(err)
			}
			w.SetBlock(pt, b)
		}

		topBlock := world.Point{X: column.xz.X, Y: int32(len(column.blocks)), Z: column.xz.Z}

		if topBlock.Y > spawnpt.Y {
			if Debug {
				log.Printf("new spawn: %s", topBlock)
			}
			spawnpt = topBlock
		}

		// JMT: naive lighting here
		w.SetSkyLight(topBlock, 15)
	}

	w.SetSpawn(spawnpt)

	return &w, nil
}

var Arrind2Debug = false

func arrind2(x int32, y int32, inx int32, iny int32, gti [6]int32) (int32, error) {
	if Arrind2Debug {
		log.Printf("x: %d, y: %d, inx: %d", x, y, inx)
		log.Printf("0: %d, 1: %d, 2: %d, 3: %d, 4: %d, 5: %d",
			gti[0], gti[1], gti[2], gti[3], gti[4], gti[5])
		log.Printf("x-0: %d, y-3: %d", x-gti[0], y-gti[3])
		log.Printf("x-0/1: %d, y-3/5: %d", (x-gti[0])/gti[1], (y-gti[3])/gti[5])
	}
	realx := (x - gti[0]) / gti[1]
	if realx < 0 {
		return 0, fmt.Errorf("realx %d < 0", realx)
	}
	if realx > inx {
		return 0, fmt.Errorf("realx %d >= inx %d", realx, inx)
	}
	realy := (y-gti[3])/gti[5] - 1
	if realy < 0 {
		return 0, fmt.Errorf("realx %d < 0", realx)
	}
	if realy > iny {
		return 0, fmt.Errorf("realy %d > iny %d", realy, iny)
	}
	return realx + realy*inx, nil
}

func (r *Region) processFeatures(in chan Feature, out chan Column, i int) {
	ds, err := gdal.Open(r.mapfile, gdal.ReadOnly)
	if err != nil {
		panic(err)
	}
	if Debug && i == 0 {
		datasetInfo(ds, "processFeatures")
	}
	inx := ds.RasterXSize()
	iny := ds.RasterYSize()
	bufferLen := inx * iny

	lcarr := make([]int16, bufferLen)
	lcBand := ds.RasterBand(Landcover)
	lcrerr := lcBand.IO(gdal.Read, 0, 0, inx, iny, lcarr, inx, iny, 0, 0)
	if notnil(lcrerr) {
		panic(lcrerr)
	}

	elevarr := make([]int16, bufferLen)
	elevBand := ds.RasterBand(Elevation)
	elevrerr := elevBand.IO(gdal.Read, 0, 0, inx, iny, elevarr, inx, iny, 0, 0)
	if notnil(elevrerr) {
		panic(elevrerr)
	}

	bathyarr := make([]int16, bufferLen)
	bathyBand := ds.RasterBand(Bathy)
	bathyrerr := bathyBand.IO(gdal.Read, 0, 0, inx, iny, bathyarr, inx, iny, 0, 0)
	if notnil(bathyrerr) {
		panic(bathyrerr)
	}

	crustarr := make([]int16, bufferLen)
	crustBand := ds.RasterBand(Crust)
	crustrerr := crustBand.IO(gdal.Read, 0, 0, inx, iny, crustarr, inx, iny, 0, 0)
	if notnil(crustrerr) {
		panic(crustrerr)
	}

	processed := 0

	// first pass at memoize
	// memo := make(map[string]Column)

	for f := range in {
		processed++

		head := fmt.Sprintf("%d: feature #%d", i, processed)

		// if Debug {
		// 	log.Printf("%s begins", head)
		// }
		pts := f.Points(ds, head)
		if len(pts) == 0 {
			log.Printf("%s: No points in geometry!", head)
			log.Print("SCRATCH ONE FEATURE")
			continue
		}

		lc := f.LCValue()
		switch lc {
		case 11:
			// "open water"
			for _, pt := range pts {
				elev := elevarr[pt.index]
				bathy := bathyarr[pt.index]
				crust := crustarr[pt.index]

				// key := fmt.Sprintf("%d|%d|%d|%d", lc, elev, bathy, crust)
				// if col, ok := memo[key]; ok {
				// 	out <- col
				// 	continue
				// }

				col := Column{}
				col.xz = pt.xz

				if int(bathy) <= r.maxdepth-1 {
					col.biome = "Deep Ocean"
				} else {
					col.biome = "Ocean"
				}

				blocks := make([]string, elev)

				for y := int16(0); y < elev; y++ {
					if y == 0 {
						blocks[y] = "Bedrock"
					} else if y < (elev - bathy - crust) {
						blocks[y] = "Stone"
					} else if y < (elev - bathy) {
						blocks[y] = "Gravel"
					} else {
						blocks[y] = "Water"
					}
				}
				col.blocks = blocks
				// memo[key] = col
				out <- col
			}
		default:
			// anything else
			for _, pt := range pts {
				elev := elevarr[pt.index]
				// bathy := bathyarr[pt.index]
				crust := crustarr[pt.index]

				// key := fmt.Sprintf("%d|%d|%d|%d", lc, elev, bathy, crust)
				// if col, ok := memo[key]; ok {
				// 	out <- col
				// 	continue
				// }

				col := Column{}
				col.xz = pt.xz
				col.biome = "Plains"

				blocks := make([]string, elev)

				for y := int16(0); y < elev; y++ {
					if y == 0 {
						blocks[y] = "Bedrock"
					} else if y < (elev - crust - 1) {
						blocks[y] = "Stone"
					} else if y < elev-1 {
						blocks[y] = "Dirt"
					} else {
						blocks[y] = "Grass Block"
					}
				}
				col.blocks = blocks
				// memo[key] = col
				out <- col
			}
		}
	}
}

type Feature struct {
	gdal.Feature
}

// returns true if the feature is valid
func (f Feature) isValid() bool {
	return f.Geometry().IsValid()
}

// returns the landcover value
func (f Feature) LCValue() int {
	// field 0 is the only field here
	return f.FieldAsInteger(0)
}

// generates list of points and sends them to a channel
func (f Feature) genPoints(in chan world.XZ, gti [6]int32) {
	g := f.Geometry()
	e := g.Envelope()
	eminx := int32(e.MinX())
	eminy := int32(e.MinY())
	emaxx := int32(e.MaxX())
	emaxy := int32(e.MaxY())

	for y := eminy; y < emaxy; y -= gti[5] {
		for x := eminx; x < emaxx; x += gti[1] {
			in <- world.XZ{X: x, Z: y}
		}
	}
	close(in)
}

type XZIndex struct {
	xz    world.XZ
	index int32
}

func (f Feature) Points(ds gdal.Dataset, head string) []XZIndex {
	inx := ds.RasterXSize()
	iny := ds.RasterYSize()
	gt := ds.GeoTransform()
	srs := gdal.CreateSpatialReference(ds.Projection())

	lcarr := make([]int16, inx*iny)
	lcBand := ds.RasterBand(Landcover)
	lcrerr := lcBand.IO(gdal.Read, 0, 0, inx, iny, lcarr, inx, iny, 0, 0)
	if notnil(lcrerr) {
		panic(lcrerr)
	}

	var gti [6]int32
	for i, v := range gt {
		gti[i] = int32(v)
	}

	in := make(chan world.XZ)
	out := make(chan XZIndex)

	var wg sync.WaitGroup

	numWorkers := runtime.NumCPU()

	for i := 0; i < numWorkers; i++ {
		wg.Add(1)
		go func(i int) {
			defer wg.Done()
			ppHead := fmt.Sprintf("%s %d", head, i)
			f.processPoints(in, out, inx, iny, gti, srs, lcarr, ppHead)
		}(i)
	}
	go func() { wg.Wait(); close(out) }()
	go f.genPoints(in, gti)

	pts := []XZIndex{}
	ptcount := 0
	for pt := range out {
		if Debug {
			if ptcount > 1 && ptcount%10000 == 0 {
				log.Printf("%s: %d columns", head, ptcount)
			}
		}
		ptcount++
		pts = append(pts, pt)
	}
	return pts
}

func (f Feature) processPoints(in chan world.XZ, out chan XZIndex, inx int, iny int, gti [6]int32, srs gdal.SpatialReference, lcarr []int16, ppHead string) {
	g := f.Geometry()

	lc := f.LCValue()

	for xz := range in {
		index, aerr := arrind2(xz.X, xz.Z, int32(inx), int32(iny), gti)
		// JMT: arrind returns nil if coordinates are invalid
		if aerr != nil {
			log.Printf("%s: aerr was not nil: %s", ppHead, aerr.Error())
			continue
		}
		if lcarr[index] != int16(lc) {
			continue
		}
		wkt := fmt.Sprintf("POINT (%f %f)", float64(xz.X)+0.5*float64(gti[1]), float64(xz.Z)-0.5*float64(gti[5]))
		pt, err := gdal.CreateFromWKT(wkt, srs)
		if notnil(err) {
			log.Printf("%s: gdal.CreateFromWKT() error: %s", ppHead, err.Error())
			continue
		}
		if g.Contains(pt) {
			out <- XZIndex{xz: world.XZ{X: xz.X / gti[1], Z: xz.Z / gti[5]}, index: index}
		}
	}
}
<commit_msg>Moved transformation to a better place.<commit_after>package carto

import (
	"fmt"
	"log"
	"runtime"

	"sync"

	"github.com/mathuin/gdal"
	"github.com/mathuin/terroir/world"
)

// JMT: convert to XZ, biome value, and list 0->n of points
type Column struct {
	xz     world.XZ
	biome  string
	blocks []string
}

func (r Region) genFeatures(in chan Feature) {
	ds, err := gdal.Open(r.mapfile, gdal.ReadOnly)
	if err != nil {
		panic(err)
	}
	if Debug {
		datasetInfo(ds, "genFeatures Input")
	}
	inx := ds.RasterXSize()
	iny := ds.RasterYSize()
	srs := gdal.CreateSpatialReference(ds.Projection())
	bufferLen := inx * iny

	lcarr := make([]int16, bufferLen)
	lcBand := ds.RasterBand(Landcover)
	lcrerr := lcBand.IO(gdal.Read, 0, 0, inx, iny, lcarr, inx, iny, 0, 0)
	if notnil(lcrerr) {
		panic(lcrerr)
	}

	// shapefile driver
	outdrv := gdal.OGRDriverByName("Memory")
	outDS, ok := outdrv.Create("out", nil)
	if !ok {
		panic(fmt.Errorf("OGR Driver Create Fail"))
	}
	outLayer := outDS.CreateLayer("polygons", srs, gdal.GT_Polygon, nil)

	// field definition
	outField := gdal.CreateFieldDefinition("lc", gdal.FT_Integer)
	outLayer.CreateField(outField, false)
	field := 0

	// options!
	options := []string{}

	// do it!
	err = lcBand.Polygonize(lcBand, outLayer, field, options, gdal.DummyProgress, nil)
	if notnil(err) {
		panic(err)
	}

	// iterate over features
	fc, ok := outLayer.FeatureCount(true)
	if !ok {
		panic(fmt.Errorf("outLayer.FeatureCount NOT OK"))
	}
	if Debug {
		log.Print("outLayer.FeatureCount(true): ", fc)
	}
	outLayer.ResetReading()
	for i := 0; i < fc; i++ {
		in <- Feature{outLayer.NextFeature()}
	}
	close(in)
}

func (r *Region) BuildWorld() (*world.World, error) {
	w := world.MakeWorld(r.name)
	w.SetRandomSeed(0)
	// JMT: need a sane storage location for files
	w.SetSaveDir(".")
	spawnpt := world.MakePoint(0, 0, 0)

	in := make(chan Feature)
	out := make(chan Column)

	var wg sync.WaitGroup

	numWorkers := runtime.NumCPU()
	// if Debug {
	// 	log.Print("debug mode - only starting one worker")
	// 	numWorkers = 1
	// }

	for i := 0; i < numWorkers; i++ {
		wg.Add(1)
		go func(i int) {
			defer wg.Done()
			r.processFeatures(in, out, i)
		}(i)
	}
	go func() { wg.Wait(); close(out) }()
	go r.genFeatures(in)

	columncount := 0
	for column := range out {
		columncount++

		b, ok := world.Biome[column.biome]
		if !ok {
			err := fmt.Errorf("biome %s not in world.Biome", column.biome)
			panic(err)
		}
		w.SetBiome(column.xz, byte(b))

		for k, v := range column.blocks {
			pt := world.Point{X: column.xz.X, Y: int32(k), Z: column.xz.Z}
			b, err := world.BlockNamed(v)
			if err != nil {
				panic(err)
			}
			w.SetBlock(pt, b)
		}

		topBlock := world.Point{X: column.xz.X, Y: int32(len(column.blocks)), Z: column.xz.Z}

		if topBlock.Y > spawnpt.Y {
			if Debug {
				log.Printf("new spawn: %s", topBlock)
			}
			spawnpt = topBlock
		}

		// JMT: naive lighting here
		w.SetSkyLight(topBlock, 15)
	}

	w.SetSpawn(spawnpt)

	return &w, nil
}

var Arrind2Debug = false

func arrind2(x int32, y int32, inx int32, iny int32, gti [6]int32) (int32, error) {
	if Arrind2Debug {
		log.Printf("x: %d, y: %d, inx: %d", x, y, inx)
		log.Printf("0: %d, 1: %d, 2: %d, 3: %d, 4: %d, 5: %d",
			gti[0], gti[1], gti[2], gti[3], gti[4], gti[5])
		log.Printf("x-0: %d, y-3: %d", x-gti[0], y-gti[3])
		log.Printf("x-0/1: %d, y-3/5: %d", (x-gti[0])/gti[1], (y-gti[3])/gti[5])
	}
	realx := (x - gti[0]) / gti[1]
	if realx < 0 {
		return 0, fmt.Errorf("realx %d < 0", realx)
	}
	if realx > inx {
		return 0, fmt.Errorf("realx %d >= inx %d", realx, inx)
	}
	realy := (y-gti[3])/gti[5] - 1
	if realy < 0 {
		return 0, fmt.Errorf("realx %d < 0", realx)
	}
	if realy > iny {
		return 0, fmt.Errorf("realy %d > iny %d", realy, iny)
	}
	return realx + realy*inx, nil
}

func (r *Region) processFeatures(in chan Feature, out chan Column, i int) {
	ds, err := gdal.Open(r.mapfile, gdal.ReadOnly)
	if err != nil {
		panic(err)
	}
	if Debug && i == 0 {
		datasetInfo(ds, "processFeatures")
	}
	inx := ds.RasterXSize()
	iny := ds.RasterYSize()
	bufferLen := inx * iny

	lcarr := make([]int16, bufferLen)
	lcBand := ds.RasterBand(Landcover)
	lcrerr := lcBand.IO(gdal.Read, 0, 0, inx, iny, lcarr, inx, iny, 0, 0)
	if notnil(lcrerr) {
		panic(lcrerr)
	}

	elevarr := make([]int16, bufferLen)
	elevBand := ds.RasterBand(Elevation)
	elevrerr := elevBand.IO(gdal.Read, 0, 0, inx, iny, elevarr, inx, iny, 0, 0)
	if notnil(elevrerr) {
		panic(elevrerr)
	}

	bathyarr := make([]int16, bufferLen)
	bathyBand := ds.RasterBand(Bathy)
	bathyrerr := bathyBand.IO(gdal.Read, 0, 0, inx, iny, bathyarr, inx, iny, 0, 0)
	if notnil(bathyrerr) {
		panic(bathyrerr)
	}

	crustarr := make([]int16, bufferLen)
	crustBand := ds.RasterBand(Crust)
	crustrerr := crustBand.IO(gdal.Read, 0, 0, inx, iny, crustarr, inx, iny, 0, 0)
	if notnil(crustrerr) {
		panic(crustrerr)
	}

	processed := 0

	// first pass at memoize
	// memo := make(map[string]Column)

	for f := range in {
		processed++

		head := fmt.Sprintf("%d: feature #%d", i, processed)

		// if Debug {
		// 	log.Printf("%s begins", head)
		// }
		pts := f.Points(ds, head)
		if len(pts) == 0 {
			log.Printf("%s: No points in geometry!", head)
			log.Print("SCRATCH ONE FEATURE")
			continue
		}

		lc := f.LCValue()
		switch lc {
		case 11:
			// "open water"
			for _, pt := range pts {
				elev := elevarr[pt.index]
				bathy := bathyarr[pt.index]
				crust := crustarr[pt.index]

				// key := fmt.Sprintf("%d|%d|%d|%d", lc, elev, bathy, crust)
				// if col, ok := memo[key]; ok {
				// 	out <- col
				// 	continue
				// }

				col := Column{}
				col.xz = pt.xz

				if int(bathy) <= r.maxdepth-1 {
					col.biome = "Deep Ocean"
				} else {
					col.biome = "Ocean"
				}

				blocks := make([]string, elev)

				for y := int16(0); y < elev; y++ {
					if y == 0 {
						blocks[y] = "Bedrock"
					} else if y < (elev - bathy - crust) {
						blocks[y] = "Stone"
					} else if y < (elev - bathy) {
						blocks[y] = "Gravel"
					} else {
						blocks[y] = "Water"
					}
				}
				col.blocks = blocks
				// memo[key] = col
				out <- col
			}
		default:
			// anything else
			for _, pt := range pts {
				elev := elevarr[pt.index]
				// bathy := bathyarr[pt.index]
				crust := crustarr[pt.index]

				// key := fmt.Sprintf("%d|%d|%d|%d", lc, elev, bathy, crust)
				// if col, ok := memo[key]; ok {
				// 	out <- col
				// 	continue
				// }

				col := Column{}
				col.xz = pt.xz
				col.biome = "Plains"

				blocks := make([]string, elev)

				for y := int16(0); y < elev; y++ {
					if y == 0 {
						blocks[y] = "Bedrock"
					} else if y < (elev - crust - 1) {
						blocks[y] = "Stone"
					} else if y < elev-1 {
						blocks[y] = "Dirt"
					} else {
						blocks[y] = "Grass Block"
					}
				}
				col.blocks = blocks
				// memo[key] = col
				out <- col
			}
		}
	}
}

type Feature struct {
	gdal.Feature
}

// returns true if the feature is valid
func (f Feature) isValid() bool {
	return f.Geometry().IsValid()
}

// returns the landcover value
func (f Feature) LCValue() int {
	// field 0 is the only field here
	return f.FieldAsInteger(0)
}

// generates list of points and sends them to a channel
func (f Feature) genPoints(in chan world.XZ, gti [6]int32) {
	g := f.Geometry()
	e := g.Envelope()
	eminx := int32(e.MinX())
	eminy := int32(e.MinY())
	emaxx := int32(e.MaxX())
	emaxy := int32(e.MaxY())

	for y := eminy; y < emaxy; y -= gti[5] {
		for x := eminx; x < emaxx; x += gti[1] {
			in <- world.XZ{X: x, Z: y}
		}
	}
	close(in)
}

type XZIndex struct {
	xz    world.XZ
	index int32
}

func makeXZIndex(xz world.XZ, index int32, gti [6]int32) XZIndex {
	return XZIndex{xz: world.XZ{X: xz.X / gti[1], Z: xz.Z / gti[5]}, index: index}
}

func (f Feature) Points(ds gdal.Dataset, head string) []XZIndex {
	inx := ds.RasterXSize()
	iny := ds.RasterYSize()
	gt := ds.GeoTransform()
	srs := gdal.CreateSpatialReference(ds.Projection())

	lcarr := make([]int16, inx*iny)
	lcBand := ds.RasterBand(Landcover)
	lcrerr := lcBand.IO(gdal.Read, 0, 0, inx, iny, lcarr, inx, iny, 0, 0)
	if notnil(lcrerr) {
		panic(lcrerr)
	}

	var gti [6]int32
	for i, v := range gt {
		gti[i] = int32(v)
	}

	in := make(chan world.XZ)
	out := make(chan XZIndex)

	var wg sync.WaitGroup

	numWorkers := runtime.NumCPU()

	for i := 0; i < numWorkers; i++ {
		wg.Add(1)
		go func(i int) {
			defer wg.Done()
			ppHead := fmt.Sprintf("%s %d", head, i)
			f.processPoints(in, out, inx, iny, gti, srs, lcarr, ppHead)
		}(i)
	}
	go func() { wg.Wait(); close(out) }()
	go f.genPoints(in, gti)

	pts := []XZIndex{}
	ptcount := 0
	for pt := range out {
		if Debug {
			if ptcount > 1 && ptcount%10000 == 0 {
				log.Printf("%s: %d columns", head, ptcount)
			}
		}
		ptcount++
		pts = append(pts, pt)
	}
	return pts
}

func (f Feature) processPoints(in chan world.XZ, out chan XZIndex, inx int, iny int, gti [6]int32, srs gdal.SpatialReference, lcarr []int16, ppHead string) {
	g := f.Geometry()

	lc := f.LCValue()

	for xz := range in {
		index, aerr := arrind2(xz.X, xz.Z, int32(inx), int32(iny), gti)
		// JMT: arrind returns nil if coordinates are invalid
		if aerr != nil {
			log.Printf("%s: aerr was not nil: %s", ppHead, aerr.Error())
			continue
		}
		if lcarr[index] != int16(lc) {
			continue
		}
		wkt := fmt.Sprintf("POINT (%f %f)", float64(xz.X)+0.5*float64(gti[1]), float64(xz.Z)-0.5*float64(gti[5]))
		pt, err := gdal.CreateFromWKT(wkt, srs)
		if notnil(err) {
			log.Printf("%s: gdal.CreateFromWKT() error: %s", ppHead, err.Error())
			continue
		}
		if g.Contains(pt) {
			out <- makeXZIndex(xz, index, gti)
		}
	}
}
<|endoftext|>


<commit_before>// Copyright 2015 Liu Dong <ddliuhb@gmail.com>.
// Licensed under the MIT license.

package spider

import (
    "time"
    // "fmt"
    "sync"
    "log"
)

const defaultConcurrency = 3

const (
    ON_START = iota
    ON_STOP = iota
)

type Listener func(*Spider, *Task)

// Create a spider.
func NewSpider() *Spider {
    spider := &Spider {
    }

    spider.prepare()
    return spider
}

type Spider struct {
    Concurrency int
    pipes []Pipe
    tasks []*Task
    events map[int][]Listener
    Stats map[Status]uint64
    IsPaused bool
    IsStopped bool
    IsDebug bool
    m sync.Mutex
    statsLock sync.Mutex
}

// Chain a pipe.
func (this *Spider) Pipe(pipe Pipe) *Spider {
    this.pipes = append(this.pipes, pipe)

    return this
}

// Initialize the spider objects.
func (this *Spider) prepare() {
    this.events = make(map[int][]Listener)
    this.Stats = make(map[Status]uint64)
    this.Stats[PENDING] = 0
    this.Stats[WORKING] = 0
    this.Stats[FAILED] = 0
    this.Stats[IGNORED] = 0
    this.Stats[DONE] = 0
}

// Run spider forever, and accept a quit channel to close it.
//
// Loop through the task list and run each of them with the help of a buffered channel.
func (this *Spider) RunForever(quit chan bool) {
    this.IsStopped = false
    this.IsPaused = false
    log.Println("[INFO] Spider started")

    if this.Concurrency <= 0 {
        this.Concurrency = defaultConcurrency
    }

    chs := make(chan bool, this.Concurrency)

    for {
        select {
        case <- quit:
            break
        default:
            // do nothing
        }

        if this.IsStopped {
            break
        }

        if this.IsPaused {
            time.Sleep(10 * time.Millisecond)
            continue
        }

        this.m.Lock()
        if len(this.tasks) > 0 {
            task := this.tasks[len(this.tasks) - 1]
            this.tasks = this.tasks[:len(this.tasks) - 1]

            this.m.Unlock()
            chs <- true
            go func() {
                this.do(task)
                <-chs
            }()
        } else {
            this.m.Unlock()

            // there is nothing to do, sleep for 10 ms
            time.Sleep(10 * time.Millisecond)
        }
    }
}

// Run spider and stop when complete.
func (this *Spider) Run() {
    quit := make(chan bool)
    go this.RunForever(quit)

    // check finish
    for {
        // if all tasks are finished, we can go out of the loop
        if this.IsFinished() {
            quit <- true
            log.Println("[INFO] Spider finished")
            break
        } else {
            time.Sleep(10 * time.Millisecond)
        }
    }
}

// Run spider and start a RPC server
func (this *Spider) RunAndServe(listen string) error {
    quit := make(chan bool)
    go this.RunForever(quit)

    return StartRPCServer(this, listen)
}

// Run a single task (should never panic)
func (this *Spider) do(task *Task) {
    defer func() {
        // error occured
        if r := recover(); r != nil {
            this.FailTask(task, r)
            return
        }

        if !task.IsEnded() {
            this.DoneTask(task)
        }
    }()

    this.StartTask(task)

    Series(this.pipes...)(this, task)
}

func (this *Spider) Pause() {
    this.IsPaused = true
}

func (this *Spider) Resume() {
    this.IsPaused = false
}

func (this *Spider) Stop() {
    this.IsStopped = true
}

// Check if all tasks have been processed.
func (this *Spider) IsFinished() bool {
    return this.Stats[PENDING] == 0 && this.Stats[WORKING] == 0
}

// Add tasks from uri.
func (this *Spider) AddUri(uris ...string) *Spider {
    for _, uri := range uris {
        this.AddTask(NewTask(uri))
    }

    return this
}

// Add a task to queue
func (this *Spider) AddTask(task *Task) *Spider {
    task.Spider = this
    this.m.Lock()
    this.tasks = append(this.tasks, task)
    this.m.Unlock()

    this.statsLock.Lock()
    this.Stats[PENDING]++
    this.statsLock.Unlock()

    return this
}

// Mark a task as failed.
func (this *Spider) FailTask(task *Task, reason interface{}) {
    task.Status = FAILED

    this.statsLock.Lock()
    this.Stats[FAILED]++
    this.Stats[WORKING]--
    this.statsLock.Unlock()

    log.Println("[WARN] Task failed: ", task.Uri, "\t", reason)
}

// Mark a task as done.
func (this *Spider) DoneTask(task *Task) {
    task.Status = DONE
    this.statsLock.Lock()
    this.Stats[DONE]++
    this.Stats[WORKING]--
    this.statsLock.Unlock()

    if this.IsDebug {
        log.Println("[DEBUG] Task done: ", task.Uri)
    }
}

// Mark a task as ignored.
func (this *Spider) IgnoreTask(task *Task, reason interface{}) {
    task.Status = IGNORED

    this.statsLock.Lock()
    this.Stats[IGNORED]++
    this.Stats[WORKING]--
    this.statsLock.Unlock()

    if this.IsDebug {
        log.Println("[DEBUG] Task ignored: ", task.Uri, "\t", reason)
    }
}

// Mark a task as started.
func (this *Spider) StartTask(task *Task) {
    task.Status = WORKING

    this.statsLock.Lock()
    this.Stats[WORKING]++
    this.Stats[PENDING]--
    this.statsLock.Unlock()

    if this.IsDebug {
        log.Println("[DEBUG] Task started: ", task.Uri)
    }
}

// Register events
func (this *Spider) On(e int, f Listener) *Spider {
    this.events[e] = append(this.events[e], f)

    return this
}

// Trigger an event
func (this *Spider) Trigger(e int, t *Task) {
    if events, ok := this.events[e]; ok {
        for _, e := range events {
            e(this, t)
        }
    }
}<commit_msg>fix memory leak<commit_after>// Copyright 2015 Liu Dong <ddliuhb@gmail.com>.
// Licensed under the MIT license.

package spider

import (
    "time"
    // "fmt"
    "sync"
    "log"
)

const defaultConcurrency = 3

const (
    ON_START = iota
    ON_STOP = iota
)

type Listener func(*Spider, *Task)

// Create a spider.
func NewSpider() *Spider {
    spider := &Spider {
    }

    spider.prepare()
    return spider
}

type Spider struct {
    Concurrency int
    pipes []Pipe
    tasks []*Task
    events map[int][]Listener
    Stats map[Status]uint64
    IsPaused bool
    IsStopped bool
    IsDebug bool
    m sync.Mutex
    statsLock sync.Mutex
}

// Chain a pipe.
func (this *Spider) Pipe(pipe Pipe) *Spider {
    this.pipes = append(this.pipes, pipe)

    return this
}

// Initialize the spider objects.
func (this *Spider) prepare() {
    this.events = make(map[int][]Listener)
    this.Stats = make(map[Status]uint64)
    this.Stats[PENDING] = 0
    this.Stats[WORKING] = 0
    this.Stats[FAILED] = 0
    this.Stats[IGNORED] = 0
    this.Stats[DONE] = 0
}

// Run spider forever, and accept a quit channel to close it.
//
// Loop through the task list and run each of them with the help of a buffered channel.
func (this *Spider) RunForever(quit chan bool) {
    this.IsStopped = false
    this.IsPaused = false
    log.Println("[INFO] Spider started")

    if this.Concurrency <= 0 {
        this.Concurrency = defaultConcurrency
    }

    chs := make(chan bool, this.Concurrency)

    for {
        select {
        case <- quit:
            break
        default:
            // do nothing
        }

        if this.IsStopped {
            break
        }

        if this.IsPaused {
            time.Sleep(10 * time.Millisecond)
            continue
        }

        this.m.Lock()
        if len(this.tasks) > 0 {
            task := this.tasks[len(this.tasks) - 1]
            this.tasks = this.tasks[:len(this.tasks) - 1]

            this.m.Unlock()
            chs <- true
            go func() {
                this.do(task)
                <-chs
            }()
        } else {
            this.m.Unlock()

            // there is nothing to do, sleep for 10 ms
            time.Sleep(10 * time.Millisecond)
        }
    }
}

// Run spider and stop when complete.
func (this *Spider) Run() {
    quit := make(chan bool)
    go this.RunForever(quit)

    // check finish
    for {
        // if all tasks are finished, we can go out of the loop
        if this.IsFinished() {
            quit <- true
            log.Println("[INFO] Spider finished")
            break
        } else {
            time.Sleep(10 * time.Millisecond)
        }
    }
}

// Run spider and start a RPC server
func (this *Spider) RunAndServe(listen string) error {
    quit := make(chan bool)
    go this.RunForever(quit)

    return StartRPCServer(this, listen)
}

// Run a single task (should never panic)
func (this *Spider) do(task *Task) {
    defer func() {
        // error occured
        if r := recover(); r != nil {
            this.FailTask(task, r)
        } else if !task.IsEnded() {
            this.DoneTask(task)
        }

        // cleanup
        task.Parent = nil
        task.Data = nil
    }()

    this.StartTask(task)

    Series(this.pipes...)(this, task)
}

func (this *Spider) Pause() {
    this.IsPaused = true
}

func (this *Spider) Resume() {
    this.IsPaused = false
}

func (this *Spider) Stop() {
    this.IsStopped = true
}

// Check if all tasks have been processed.
func (this *Spider) IsFinished() bool {
    return this.Stats[PENDING] == 0 && this.Stats[WORKING] == 0
}

// Add tasks from uri.
func (this *Spider) AddUri(uris ...string) *Spider {
    for _, uri := range uris {
        this.AddTask(NewTask(uri))
    }

    return this
}

// Add a task to queue
func (this *Spider) AddTask(task *Task) *Spider {
    task.Spider = this
    this.m.Lock()
    this.tasks = append(this.tasks, task)
    this.m.Unlock()

    this.statsLock.Lock()
    this.Stats[PENDING]++
    this.statsLock.Unlock()

    return this
}

// Mark a task as failed.
func (this *Spider) FailTask(task *Task, reason interface{}) {
    task.Status = FAILED

    this.statsLock.Lock()
    this.Stats[FAILED]++
    this.Stats[WORKING]--
    this.statsLock.Unlock()

    log.Println("[WARN] Task failed: ", task.Uri, "\t", reason)
}

// Mark a task as done.
func (this *Spider) DoneTask(task *Task) {
    task.Status = DONE
    this.statsLock.Lock()
    this.Stats[DONE]++
    this.Stats[WORKING]--
    this.statsLock.Unlock()

    if this.IsDebug {
        log.Println("[DEBUG] Task done: ", task.Uri)
    }
}

// Mark a task as ignored.
func (this *Spider) IgnoreTask(task *Task, reason interface{}) {
    task.Status = IGNORED

    this.statsLock.Lock()
    this.Stats[IGNORED]++
    this.Stats[WORKING]--
    this.statsLock.Unlock()

    if this.IsDebug {
        log.Println("[DEBUG] Task ignored: ", task.Uri, "\t", reason)
    }
}

// Mark a task as started.
func (this *Spider) StartTask(task *Task) {
    task.Status = WORKING

    this.statsLock.Lock()
    this.Stats[WORKING]++
    this.Stats[PENDING]--
    this.statsLock.Unlock()

    if this.IsDebug {
        log.Println("[DEBUG] Task started: ", task.Uri)
    }
}

// Register events
func (this *Spider) On(e int, f Listener) *Spider {
    this.events[e] = append(this.events[e], f)

    return this
}

// Trigger an event
func (this *Spider) Trigger(e int, t *Task) {
    if events, ok := this.events[e]; ok {
        for _, e := range events {
            e(this, t)
        }
    }
}<|endoftext|>


<commit_before>package cli

import (
	"fmt"
	"io"
	"os"
	"os/exec"
	"os/signal"
	"strings"
	"syscall"

	"github.com/spf13/cobra"

	"github.com/datawire/telepresence2/pkg/client"
	manager "github.com/datawire/telepresence2/pkg/rpc"
	"github.com/datawire/telepresence2/pkg/rpc/connector"
)

// runner contains all parameters needed in order to run an intercepted command.
type runner struct {
	connector.ConnectRequest
	manager.CreateInterceptRequest
	DNS             string
	Fallback        string
	RemoveIntercept string
	List            bool
	NoWait          bool
	Quit            bool
	Status          bool
	Version         bool
}

// run will ensure that an intercept is in place and then execute the command given by args[0]
// and the arguments starting at args[1:].
func (p *runner) run(cmd *cobra.Command, args []string) error {
	switch {
	case p.List:
		return listIntercepts(cmd, []string{})
	case p.Quit:
		return quit(cmd, []string{})
	case p.Status:
		return status(cmd, []string{})
	case p.RemoveIntercept != "":
		return removeIntercept(cmd, []string{p.RemoveIntercept})
	case p.Version:
		return printVersion(cmd, []string{})
	}

	doWithResource := func(context string) error {
		switch {
		case p.NoWait:
			return nil
		case len(args) == 0:
			return p.startSubshell(cmd, context)
		default:
			return start(args[0], args[1:], true, cmd.InOrStdin(), cmd.OutOrStdout(), cmd.ErrOrStderr())
		}
	}

	if p.CreateInterceptRequest.InterceptSpec.Name != "" {
		if err := prepareIntercept(&p.CreateInterceptRequest); err != nil {
			return err
		}
		return p.runWithIntercept(cmd, func(is *interceptState) error { return doWithResource(is.cs.info.ClusterContext) })
	}
	return p.runWithConnector(cmd, func(cs *connectorState) error { return doWithResource(cs.info.ClusterContext) })
}

func (p *runner) startSubshell(cmd *cobra.Command, ctx string) error {
	exe := os.Getenv("SHELL")
	var envArg string
	var args []string
	if strings.HasSuffix(exe, "/zsh") {
		// TODO: Find a way to alter zsh prompt. Not sure it's possible since the option
		//  prompt_subst must be set for prompt substitution to take place. The only way
		//  to set it is from files being sourced when the command starts. If that is
		//  enabled (and it really should be, it's very common), then the PS1 passed from
		//  here is overwritten.
		exe = "/bin/bash"
	}
	envArg = fmt.Sprintf(`PROMPT_COMMAND=export PS1="@%s $PS1";unset PROMPT_COMMAND`, ctx)
	args = []string{"-i"}
	out := cmd.OutOrStdout()
	fmt.Fprintf(out, "Starting a %s subshell\n", exe)
	return start(exe, args, true, cmd.InOrStdin(), out, cmd.ErrOrStderr(), envArg)
}

func (p *runner) runWithDaemon(cmd *cobra.Command, f func(ds *daemonState) error) error {
	ds, err := newDaemonState(cmd, p.DNS, p.Fallback)
	if err != nil && err != errDaemonIsNotRunning {
		return err
	}
	return client.WithEnsuredState(ds, p.NoWait, func() error { return f(ds) })
}

func (p *runner) runWithConnector(cmd *cobra.Command, f func(cs *connectorState) error) error {
	return p.runWithDaemon(cmd, func(ds *daemonState) error {
		p.InterceptEnabled = true
		cs, err := newConnectorState(ds.grpc, &p.ConnectRequest, cmd)
		if err != nil && err != errConnectorIsNotRunning {
			return err
		}
		return client.WithEnsuredState(cs, p.NoWait, func() error { return f(cs) })
	})
}

func (p *runner) runWithIntercept(cmd *cobra.Command, f func(is *interceptState) error) error {
	return p.runWithConnector(cmd, func(cs *connectorState) error {
		is := newInterceptState(cs, &p.CreateInterceptRequest, cmd)
		return client.WithEnsuredState(is, p.NoWait, func() error { return f(is) })
	})
}

func runAsRoot(exe string, args []string) error {
	if os.Geteuid() != 0 {
		if err := exec.Command("sudo", "-n", "true").Run(); err != nil {
			fmt.Printf("Need root privileges to run %q\n", client.ShellString(exe, args))
			if err = exec.Command("sudo", "true").Run(); err != nil {
				return err
			}
		}
		args = append([]string{"-n", "-E", exe}, args...)
		exe = "sudo"
	}
	return start(exe, args, false, nil, nil, nil)
}

func start(exe string, args []string, wait bool, stdin io.Reader, stdout, stderr io.Writer, env ...string) error {
	cmd := exec.Command(exe, args...)
	cmd.Stdout = stdout
	cmd.Stderr = stderr
	cmd.Stdin = stdin
	if len(env) > 0 {
		cmd.Env = append(os.Environ(), env...)
	}
	if !wait {
		// Process must live in a process group of its own to prevent
		// getting affected by <ctrl-c> in the terminal
		cmd.SysProcAttr = &syscall.SysProcAttr{Setpgid: true}
	}

	var err error
	if err = cmd.Start(); err != nil {
		return fmt.Errorf("%s: %v", client.ShellString(exe, args), err)
	}
	if !wait {
		_ = cmd.Process.Release()
		return nil
	}

	// Ensure that SIGINT and SIGTERM are propagated to the child process
	sigCh := make(chan os.Signal, 1)
	signal.Notify(sigCh, syscall.SIGINT, syscall.SIGTERM)
	go func() {
		sig := <-sigCh
		if sig == nil {
			return
		}
		_ = cmd.Process.Signal(sig)
	}()
	s, err := cmd.Process.Wait()
	if err != nil {
		return fmt.Errorf("%s: %v", client.ShellString(exe, args), err)
	}

	sigCh <- nil
	exitCode := s.ExitCode()
	if exitCode != 0 {
		return fmt.Errorf("%s %s: exited with %d", exe, strings.Join(args, " "), exitCode)
	}
	return nil
}
<commit_msg>Remove code that attempts to change prompt for subshell<commit_after>package cli

import (
	"fmt"
	"io"
	"os"
	"os/exec"
	"os/signal"
	"strings"
	"syscall"

	"github.com/spf13/cobra"

	"github.com/datawire/telepresence2/pkg/client"
	manager "github.com/datawire/telepresence2/pkg/rpc"
	"github.com/datawire/telepresence2/pkg/rpc/connector"
)

// runner contains all parameters needed in order to run an intercepted command.
type runner struct {
	connector.ConnectRequest
	manager.CreateInterceptRequest
	DNS             string
	Fallback        string
	RemoveIntercept string
	List            bool
	NoWait          bool
	Quit            bool
	Status          bool
	Version         bool
}

// run will ensure that an intercept is in place and then execute the command given by args[0]
// and the arguments starting at args[1:].
func (p *runner) run(cmd *cobra.Command, args []string) error {
	switch {
	case p.List:
		return listIntercepts(cmd, []string{})
	case p.Quit:
		return quit(cmd, []string{})
	case p.Status:
		return status(cmd, []string{})
	case p.RemoveIntercept != "":
		return removeIntercept(cmd, []string{p.RemoveIntercept})
	case p.Version:
		return printVersion(cmd, []string{})
	}

	doWithResource := func(context string) error {
		switch {
		case p.NoWait:
			return nil
		case len(args) == 0:
			return p.startSubshell(cmd, context)
		default:
			return start(args[0], args[1:], true, cmd.InOrStdin(), cmd.OutOrStdout(), cmd.ErrOrStderr())
		}
	}

	if p.CreateInterceptRequest.InterceptSpec.Name != "" {
		if err := prepareIntercept(&p.CreateInterceptRequest); err != nil {
			return err
		}
		return p.runWithIntercept(cmd, func(is *interceptState) error { return doWithResource(is.cs.info.ClusterContext) })
	}
	return p.runWithConnector(cmd, func(cs *connectorState) error { return doWithResource(cs.info.ClusterContext) })
}

func (p *runner) startSubshell(cmd *cobra.Command, ctx string) error {
	exe := os.Getenv("SHELL")
	out := cmd.OutOrStdout()
	fmt.Fprintf(out, "Starting a %s subshell\n", exe)
	return start(exe, []string{"i"}, true, cmd.InOrStdin(), out, cmd.ErrOrStderr())
}

func (p *runner) runWithDaemon(cmd *cobra.Command, f func(ds *daemonState) error) error {
	ds, err := newDaemonState(cmd, p.DNS, p.Fallback)
	if err != nil && err != errDaemonIsNotRunning {
		return err
	}
	return client.WithEnsuredState(ds, p.NoWait, func() error { return f(ds) })
}

func (p *runner) runWithConnector(cmd *cobra.Command, f func(cs *connectorState) error) error {
	return p.runWithDaemon(cmd, func(ds *daemonState) error {
		p.InterceptEnabled = true
		cs, err := newConnectorState(ds.grpc, &p.ConnectRequest, cmd)
		if err != nil && err != errConnectorIsNotRunning {
			return err
		}
		return client.WithEnsuredState(cs, p.NoWait, func() error { return f(cs) })
	})
}

func (p *runner) runWithIntercept(cmd *cobra.Command, f func(is *interceptState) error) error {
	return p.runWithConnector(cmd, func(cs *connectorState) error {
		is := newInterceptState(cs, &p.CreateInterceptRequest, cmd)
		return client.WithEnsuredState(is, p.NoWait, func() error { return f(is) })
	})
}

func runAsRoot(exe string, args []string) error {
	if os.Geteuid() != 0 {
		if err := exec.Command("sudo", "-n", "true").Run(); err != nil {
			fmt.Printf("Need root privileges to run %q\n", client.ShellString(exe, args))
			if err = exec.Command("sudo", "true").Run(); err != nil {
				return err
			}
		}
		args = append([]string{"-n", "-E", exe}, args...)
		exe = "sudo"
	}
	return start(exe, args, false, nil, nil, nil)
}

func start(exe string, args []string, wait bool, stdin io.Reader, stdout, stderr io.Writer, env ...string) error {
	cmd := exec.Command(exe, args...)
	cmd.Stdout = stdout
	cmd.Stderr = stderr
	cmd.Stdin = stdin
	if len(env) > 0 {
		cmd.Env = append(os.Environ(), env...)
	}
	if !wait {
		// Process must live in a process group of its own to prevent
		// getting affected by <ctrl-c> in the terminal
		cmd.SysProcAttr = &syscall.SysProcAttr{Setpgid: true}
	}

	var err error
	if err = cmd.Start(); err != nil {
		return fmt.Errorf("%s: %v", client.ShellString(exe, args), err)
	}
	if !wait {
		_ = cmd.Process.Release()
		return nil
	}

	// Ensure that SIGINT and SIGTERM are propagated to the child process
	sigCh := make(chan os.Signal, 1)
	signal.Notify(sigCh, syscall.SIGINT, syscall.SIGTERM)
	go func() {
		sig := <-sigCh
		if sig == nil {
			return
		}
		_ = cmd.Process.Signal(sig)
	}()
	s, err := cmd.Process.Wait()
	if err != nil {
		return fmt.Errorf("%s: %v", client.ShellString(exe, args), err)
	}

	sigCh <- nil
	exitCode := s.ExitCode()
	if exitCode != 0 {
		return fmt.Errorf("%s %s: exited with %d", exe, strings.Join(args, " "), exitCode)
	}
	return nil
}
<|endoftext|>


<commit_before>package network

import (
	"github.com/name5566/leaf/log"
	"net"
	"sync"
)

type TCPServer struct {
	Addr            string
	MaxConnNum      int
	PendingWriteNum int
	NewAgent        func(*TCPConn) Agent
	ln              net.Listener
	conns           ConnSet
	mutexConns      sync.Mutex
	wg              sync.WaitGroup
	closeFlag       bool

	// msg parser
	LenMsgLen    int
	MinMsgLen    uint32
	MaxMsgLen    uint32
	LittleEndian bool
	msgParser    *MsgParser
}

func (server *TCPServer) Start() {
	server.init()
	go server.run()
}

func (server *TCPServer) init() {
	ln, err := net.Listen("tcp", server.Addr)
	if err != nil {
		log.Fatal("%v", err)
	}

	if server.MaxConnNum <= 0 {
		server.MaxConnNum = 100
		log.Release("invalid MaxConnNum, reset to %v", server.MaxConnNum)
	}
	if server.PendingWriteNum <= 0 {
		server.PendingWriteNum = 100
		log.Release("invalid PendingWriteNum, reset to %v", server.PendingWriteNum)
	}
	if server.NewAgent == nil {
		log.Fatal("NewAgent must not be nil")
	}

	server.ln = ln
	server.conns = make(ConnSet)
	server.closeFlag = false

	// msg parser
	msgParser := NewMsgParser()
	msgParser.SetMsgLen(server.LenMsgLen, server.MinMsgLen, server.MaxMsgLen)
	msgParser.SetByteOrder(server.LittleEndian)
	server.msgParser = msgParser
}

func (server *TCPServer) run() {
	for {
		conn, err := server.ln.Accept()
		if err != nil {
			if server.closeFlag {
				return
			} else {
				log.Error("accept error: %v", err)
				continue
			}
		}

		server.mutexConns.Lock()
		if len(server.conns) >= server.MaxConnNum {
			server.mutexConns.Unlock()
			conn.Close()
			log.Debug("too many connections")
			continue
		}
		server.conns[conn] = struct{}{}
		server.mutexConns.Unlock()

		server.wg.Add(1)

		tcpConn := newTCPConn(conn, server.PendingWriteNum, server.msgParser)
		agent := server.NewAgent(tcpConn)
		go func() {
			agent.Run()

			// cleanup
			tcpConn.Close()
			server.mutexConns.Lock()
			delete(server.conns, conn)
			server.mutexConns.Unlock()
			agent.OnClose()

			server.wg.Done()
		}()
	}
}

func (server *TCPServer) Close() {
	server.closeFlag = true
	server.ln.Close()

	server.mutexConns.Lock()
	for conn := range server.conns {
		conn.Close()
	}
	server.conns = make(ConnSet)
	server.mutexConns.Unlock()

	server.wg.Wait()
}
<commit_msg>bug fix on Accept<commit_after>package network

import (
	"github.com/name5566/leaf/log"
	"net"
	"sync"
)

type TCPServer struct {
	Addr            string
	MaxConnNum      int
	PendingWriteNum int
	NewAgent        func(*TCPConn) Agent
	ln              net.Listener
	conns           ConnSet
	mutexConns      sync.Mutex
	wg              sync.WaitGroup

	// msg parser
	LenMsgLen    int
	MinMsgLen    uint32
	MaxMsgLen    uint32
	LittleEndian bool
	msgParser    *MsgParser
}

func (server *TCPServer) Start() {
	server.init()
	go server.run()
}

func (server *TCPServer) init() {
	ln, err := net.Listen("tcp", server.Addr)
	if err != nil {
		log.Fatal("%v", err)
	}

	if server.MaxConnNum <= 0 {
		server.MaxConnNum = 100
		log.Release("invalid MaxConnNum, reset to %v", server.MaxConnNum)
	}
	if server.PendingWriteNum <= 0 {
		server.PendingWriteNum = 100
		log.Release("invalid PendingWriteNum, reset to %v", server.PendingWriteNum)
	}
	if server.NewAgent == nil {
		log.Fatal("NewAgent must not be nil")
	}

	server.ln = ln
	server.conns = make(ConnSet)

	// msg parser
	msgParser := NewMsgParser()
	msgParser.SetMsgLen(server.LenMsgLen, server.MinMsgLen, server.MaxMsgLen)
	msgParser.SetByteOrder(server.LittleEndian)
	server.msgParser = msgParser
}

func (server *TCPServer) run() {
	for {
		conn, err := server.ln.Accept()
		if err != nil {
			return
		}

		server.mutexConns.Lock()
		if len(server.conns) >= server.MaxConnNum {
			server.mutexConns.Unlock()
			conn.Close()
			log.Debug("too many connections")
			continue
		}
		server.conns[conn] = struct{}{}
		server.mutexConns.Unlock()

		server.wg.Add(1)

		tcpConn := newTCPConn(conn, server.PendingWriteNum, server.msgParser)
		agent := server.NewAgent(tcpConn)
		go func() {
			agent.Run()

			// cleanup
			tcpConn.Close()
			server.mutexConns.Lock()
			delete(server.conns, conn)
			server.mutexConns.Unlock()
			agent.OnClose()

			server.wg.Done()
		}()
	}
}

func (server *TCPServer) Close() {
	server.ln.Close()

	server.mutexConns.Lock()
	for conn := range server.conns {
		conn.Close()
	}
	server.conns = make(ConnSet)
	server.mutexConns.Unlock()

	server.wg.Wait()
}
<|endoftext|>


<commit_before>// File: taorpc.go
// Author: Kevin Walsh <kwalsh@holycross.edu>
// Description: Support for RPC from hosted program to host Tao.
//
// Copyright (c) 2013, Google Inc.  All rights reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package tao

import (
	"errors"
	"io"
	"math"
	"net/rpc"
	"strings"

	"code.google.com/p/goprotobuf/proto"

	"cloudproxy/util"
	"cloudproxy/util/protorpc"
)

var opRPCName = map[string]string{
	"Tao.GetRandomBytes":  "TAO_RPC_GET_RANDOM_BYTES",
	"Tao.Seal":            "TAO_RPC_SEAL",
	"Tao.Unseal":          "TAO_RPC_UNSEAL",
	"Tao.Attest":          "TAO_RPC_ATTEST",
	"Tao.GetTaoName":      "TAO_RPC_GET_TAO_NAME",
	"Tao.ExtendTaoName":   "TAO_RPC_EXTEND_TAO_NAME",
	"Tao.GetSharedSecret": "TAO_RPC_GET_SHARED_SECRET",
}

var opGoName = make(map[string]string)

func init() {
	for goName, rpcName := range opRPCName {
		opGoName[rpcName] = goName
	}
}

// Convert string "Tao.FooBar" into integer TaoRPCOperation_TAO_RPC_FOO_BAR.
func goToRPC(m string) (TaoRPCOperation, error) {
	op := TaoRPCOperation(TaoRPCOperation_value[opRPCName[m]])
	if op == TaoRPCOperation(0) {
		return op, protorpc.ErrBadRequestType
	}
	return op, nil
}

// Convert integer TaoRPCOperation_TAO_RPC_FOO_BAR into string "Tao.FooBar".
func rpcToGo(op TaoRPCOperation) (string, error) {
	s := opGoName[TaoRPCOperation_name[int32(op)]]
	if s == "" {
		return "", protorpc.ErrBadRequestType
	}
	return s, nil
}

type taoMux struct{}

func (taoMux) SetRequestHeader(req proto.Message, servicemethod string, seq uint64) error {
	m, ok := req.(*TaoRPCRequest)
	if !ok || m == nil {
		return protorpc.ErrBadRequestType
	}
	rpc, err := goToRPC(servicemethod)
	if err != nil {
		return err
	}
	m.Rpc = &rpc
	m.Seq = &seq
	return nil
}

func (taoMux) SetResponseHeader(req proto.Message, servicemethod string, seq uint64) error {
	m, ok := req.(*TaoRPCResponse)
	if !ok || m == nil {
		return protorpc.ErrBadResponseType
	}
	rpc, err := goToRPC(servicemethod)
	if err != nil {
		return err
	}
	m.Rpc = &rpc
	m.Seq = &seq
	return nil
}

func (taoMux) GetServiceMethod(number uint64) (string, error) {
	return rpcToGo(TaoRPCOperation(int32(number)))
}

// TaoRPC sends requests between this hosted program and the host Tao.
type TaoRPC struct {
	rpc *rpc.Client
}

func DeserializeTaoRPC(s string) (*TaoRPC, error) {
	if s == "" {
		return nil, errors.New("taorpc: missing host Tao spec" +
			" (ensure $" + HostTaoEnvVar + " is set)")
	}
	r := strings.TrimPrefix(s, "tao::TaoRPC+")
	if r == s {
		return nil, errors.New("taorpc: unrecognized $" + HostTaoEnvVar + " string " + s)
	}
	ms, err := util.DeserializeFDMessageStream(r)
	if err != nil {
		return nil, errors.New("taorpc: unrecognized $" + HostTaoEnvVar + " string " + s +
			" (" + err.Error() + ")")
	}
	return &TaoRPC{protorpc.NewClient(ms, taoMux{})}, nil
}

type expectedResponse int

const (
	wantNothing                  = 0
	wantData    expectedResponse = 1 << iota
	wantPolicy
)

var ErrMalformedResponse = errors.New("taorpc: malformed response")

func (t *TaoRPC) call(method string, r *TaoRPCRequest, e expectedResponse) (data []byte, policy string, err error) {
	s := new(TaoRPCResponse)
	err = t.rpc.Call(method, r, s)
	if err != nil {
		return
	}
	if s.Error != nil {
		err = errors.New(*s.Error)
		return
	}
	if (s.Data != nil) != (e&wantData != 0) ||
		(s.Policy != nil) != (e&wantPolicy != 0) {
		err = ErrMalformedResponse
		return
	}
	if s.Data != nil {
		data = s.Data
	}
	if s.Policy != nil {
		policy = *s.Policy
	}
	return
}

func (t *TaoRPC) GetTaoName() (string, error) {
	r := &TaoRPCRequest{}
	data, _, err := t.call("Tao.GetTaoName", r, wantData)
	return string(data), err
}

func (t *TaoRPC) ExtendTaoName(subprin string) error {
	r := &TaoRPCRequest{Data: []byte(subprin)}
	_, _, err := t.call("Tao.ExtendTaoName", r, wantNothing)
	return err
}

type taoRandReader TaoRPC

func (t *taoRandReader) Read(p []byte) (n int, err error) {
	bytes, err := (*TaoRPC)(t).GetRandomBytes(len(p))
	if err != nil {
		return 0, err
	}
	copy(p, bytes)
	return len(p), nil
}

// TODO(kwalsh) Can Rand be made generic, or does it need to be defined for the
// concrete type TaoRPC?
func (t *TaoRPC) Rand() io.Reader {
	return (*taoRandReader)(t)
}

func (t *TaoRPC) GetRandomBytes(n int) ([]byte, error) {
	if n > math.MaxUint32 {
		return nil, errors.New("taorpc: request for too many random bytes")
	}
	r := &TaoRPCRequest{Size: proto.Int32(int32(n))}
	bytes, _, err := t.call("Tao.GetRandomBytes", r, wantData)
	return bytes, err
}

func (t *TaoRPC) GetSharedSecret(n int, policy string) ([]byte, error) {
	if n > math.MaxUint32 {
		return nil, errors.New("taorpc: request for too many secret bytes")
	}
	r := &TaoRPCRequest{Size: proto.Int32(int32(n)), Policy: proto.String(policy)}
	bytes, _, err := t.call("Tao.GetSharedSecret", r, wantData)
	return bytes, err
}

func (t *TaoRPC) Attest(stmt *Statement) (*Attestation, error) {
	data, err := proto.Marshal(stmt)
	if _, ok := err.(*proto.RequiredNotSetError); err != nil && !ok {
		return nil, err
	}
	r := &TaoRPCRequest{Data: data}
	bytes, _, err := t.call("Tao.Attest", r, wantData)
	if err != nil {
		return nil, err
	}
	var a Attestation
	err = proto.Unmarshal(bytes, &a)
	if err != nil {
		return nil, err
	}
	return &a, nil
}

func (t *TaoRPC) Seal(data []byte, policy string) (sealed []byte, err error) {
	r := &TaoRPCRequest{Data: data, Policy: proto.String(policy)}
	sealed, _, err = t.call("Tao.Seal", r, wantData)
	return
}

func (t *TaoRPC) Unseal(sealed []byte) (data []byte, policy string, err error) {
	r := &TaoRPCRequest{Data: sealed}
	data, policy, err = t.call("Tao.Unseal", r, wantData|wantPolicy)
	return
}
<commit_msg>taorpc: missing type on const<commit_after>// File: taorpc.go
// Author: Kevin Walsh <kwalsh@holycross.edu>
// Description: Support for RPC from hosted program to host Tao.
//
// Copyright (c) 2013, Google Inc.  All rights reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package tao

import (
	"errors"
	"io"
	"math"
	"net/rpc"
	"strings"

	"code.google.com/p/goprotobuf/proto"

	"cloudproxy/util"
	"cloudproxy/util/protorpc"
)

var opRPCName = map[string]string{
	"Tao.GetRandomBytes":  "TAO_RPC_GET_RANDOM_BYTES",
	"Tao.Seal":            "TAO_RPC_SEAL",
	"Tao.Unseal":          "TAO_RPC_UNSEAL",
	"Tao.Attest":          "TAO_RPC_ATTEST",
	"Tao.GetTaoName":      "TAO_RPC_GET_TAO_NAME",
	"Tao.ExtendTaoName":   "TAO_RPC_EXTEND_TAO_NAME",
	"Tao.GetSharedSecret": "TAO_RPC_GET_SHARED_SECRET",
}

var opGoName = make(map[string]string)

func init() {
	for goName, rpcName := range opRPCName {
		opGoName[rpcName] = goName
	}
}

// Convert string "Tao.FooBar" into integer TaoRPCOperation_TAO_RPC_FOO_BAR.
func goToRPC(m string) (TaoRPCOperation, error) {
	op := TaoRPCOperation(TaoRPCOperation_value[opRPCName[m]])
	if op == TaoRPCOperation(0) {
		return op, protorpc.ErrBadRequestType
	}
	return op, nil
}

// Convert integer TaoRPCOperation_TAO_RPC_FOO_BAR into string "Tao.FooBar".
func rpcToGo(op TaoRPCOperation) (string, error) {
	s := opGoName[TaoRPCOperation_name[int32(op)]]
	if s == "" {
		return "", protorpc.ErrBadRequestType
	}
	return s, nil
}

type taoMux struct{}

func (taoMux) SetRequestHeader(req proto.Message, servicemethod string, seq uint64) error {
	m, ok := req.(*TaoRPCRequest)
	if !ok || m == nil {
		return protorpc.ErrBadRequestType
	}
	rpc, err := goToRPC(servicemethod)
	if err != nil {
		return err
	}
	m.Rpc = &rpc
	m.Seq = &seq
	return nil
}

func (taoMux) SetResponseHeader(req proto.Message, servicemethod string, seq uint64) error {
	m, ok := req.(*TaoRPCResponse)
	if !ok || m == nil {
		return protorpc.ErrBadResponseType
	}
	rpc, err := goToRPC(servicemethod)
	if err != nil {
		return err
	}
	m.Rpc = &rpc
	m.Seq = &seq
	return nil
}

func (taoMux) GetServiceMethod(number uint64) (string, error) {
	return rpcToGo(TaoRPCOperation(int32(number)))
}

// TaoRPC sends requests between this hosted program and the host Tao.
type TaoRPC struct {
	rpc *rpc.Client
}

func DeserializeTaoRPC(s string) (*TaoRPC, error) {
	if s == "" {
		return nil, errors.New("taorpc: missing host Tao spec" +
			" (ensure $" + HostTaoEnvVar + " is set)")
	}
	r := strings.TrimPrefix(s, "tao::TaoRPC+")
	if r == s {
		return nil, errors.New("taorpc: unrecognized $" + HostTaoEnvVar + " string " + s)
	}
	ms, err := util.DeserializeFDMessageStream(r)
	if err != nil {
		return nil, errors.New("taorpc: unrecognized $" + HostTaoEnvVar + " string " + s +
			" (" + err.Error() + ")")
	}
	return &TaoRPC{protorpc.NewClient(ms, taoMux{})}, nil
}

type expectedResponse int

const (
	wantNothing expectedResponse = 0
	wantData    expectedResponse = 1 << iota
	wantPolicy
)

var ErrMalformedResponse = errors.New("taorpc: malformed response")

func (t *TaoRPC) call(method string, r *TaoRPCRequest, e expectedResponse) (data []byte, policy string, err error) {
	s := new(TaoRPCResponse)
	err = t.rpc.Call(method, r, s)
	if err != nil {
		return
	}
	if s.Error != nil {
		err = errors.New(*s.Error)
		return
	}
	if (s.Data != nil) != (e&wantData != 0) ||
		(s.Policy != nil) != (e&wantPolicy != 0) {
		err = ErrMalformedResponse
		return
	}
	if s.Data != nil {
		data = s.Data
	}
	if s.Policy != nil {
		policy = *s.Policy
	}
	return
}

func (t *TaoRPC) GetTaoName() (string, error) {
	r := &TaoRPCRequest{}
	data, _, err := t.call("Tao.GetTaoName", r, wantData)
	return string(data), err
}

func (t *TaoRPC) ExtendTaoName(subprin string) error {
	r := &TaoRPCRequest{Data: []byte(subprin)}
	_, _, err := t.call("Tao.ExtendTaoName", r, wantNothing)
	return err
}

type taoRandReader TaoRPC

func (t *taoRandReader) Read(p []byte) (n int, err error) {
	bytes, err := (*TaoRPC)(t).GetRandomBytes(len(p))
	if err != nil {
		return 0, err
	}
	copy(p, bytes)
	return len(p), nil
}

// TODO(kwalsh) Can Rand be made generic, or does it need to be defined for the
// concrete type TaoRPC?
func (t *TaoRPC) Rand() io.Reader {
	return (*taoRandReader)(t)
}

func (t *TaoRPC) GetRandomBytes(n int) ([]byte, error) {
	if n > math.MaxUint32 {
		return nil, errors.New("taorpc: request for too many random bytes")
	}
	r := &TaoRPCRequest{Size: proto.Int32(int32(n))}
	bytes, _, err := t.call("Tao.GetRandomBytes", r, wantData)
	return bytes, err
}

func (t *TaoRPC) GetSharedSecret(n int, policy string) ([]byte, error) {
	if n > math.MaxUint32 {
		return nil, errors.New("taorpc: request for too many secret bytes")
	}
	r := &TaoRPCRequest{Size: proto.Int32(int32(n)), Policy: proto.String(policy)}
	bytes, _, err := t.call("Tao.GetSharedSecret", r, wantData)
	return bytes, err
}

func (t *TaoRPC) Attest(stmt *Statement) (*Attestation, error) {
	data, err := proto.Marshal(stmt)
	if _, ok := err.(*proto.RequiredNotSetError); err != nil && !ok {
		return nil, err
	}
	r := &TaoRPCRequest{Data: data}
	bytes, _, err := t.call("Tao.Attest", r, wantData)
	if err != nil {
		return nil, err
	}
	var a Attestation
	err = proto.Unmarshal(bytes, &a)
	if err != nil {
		return nil, err
	}
	return &a, nil
}

func (t *TaoRPC) Seal(data []byte, policy string) (sealed []byte, err error) {
	r := &TaoRPCRequest{Data: data, Policy: proto.String(policy)}
	sealed, _, err = t.call("Tao.Seal", r, wantData)
	return
}

func (t *TaoRPC) Unseal(sealed []byte) (data []byte, policy string, err error) {
	r := &TaoRPCRequest{Data: sealed}
	data, policy, err = t.call("Tao.Unseal", r, wantData|wantPolicy)
	return
}
<|endoftext|>


<commit_before>package main

import (
	"bytes"
	"flag"
	"gopkg.in/fsnotify.v1"
	"io"
	"log"
	"net/http"
	"os"
	"os/exec"
	"path/filepath"
	"sync"
)

func main() {
	d, err := os.Getwd()
	if err != nil {
		d = "./"
	}

	port := flag.String("http", ":8081", "the listening port")
	dir := flag.String("dir", d, "the working folder")
	flag.Parse()

	compiler := filepath.Join(*dir, "compiler.jar")
	_, err = os.Stat(compiler)
	compilerMissing := os.IsNotExist(err)

	jsDir := filepath.Join(*dir, "src")
	_, err = os.Stat(jsDir)
	jsDirMissing := os.IsNotExist(err)

	lock := sync.Mutex{}

	if compilerMissing {
		log.Println("Not found 'compiler.jar', please download it here: http://dl.google.com/closure-compiler/compiler-latest.zip")
	} else if jsDirMissing {
		log.Println("Not found 'src' dir")
	} else {
		watcher, err := fsnotify.NewWatcher()
		if err != nil {
			log.Fatal(err)
		}
		defer watcher.Close()

		go func() {
			for {
				select {
				case event := <-watcher.Events:
					// call compile only for script change
					if event.Op != fsnotify.Chmod && event.Op != fsnotify.Rename {
						if ext := filepath.Ext(event.Name); ext != ".js" && ext != ".json" {
							break
						}

						lock.Lock()
						log.Println("Compiling script...")
						err := compileScript(compiler, jsDir)
						if err != nil {
							log.Println("Error compiling script:", err)
						}
						log.Println("Done!")
						lock.Unlock()
					}
				case err := <-watcher.Errors:
					log.Println("watcher error:", err)
				}
			}
		}()

		// watch all sub-dir
		err = filepath.Walk(jsDir, func(path string, info os.FileInfo, err error) error {
			if info.IsDir() {
				err := watcher.Add(path)
				if err != nil {
					return err
				}
			}
			return nil
		})
		if err != nil {
			log.Fatal(err)
		}
	}

	lock.Lock()
	err = compileScript(compiler, jsDir)
	if err != nil {
		log.Println("Error compiling script:", err)
	} else {
		log.Println("Compiling script at", jsDir)
	}
	lock.Unlock()

	log.Println("Serving client at http://localhost" + *port)
	panic(http.ListenAndServe(*port, http.FileServer(http.Dir(*dir))))
}

func compileScript(compiler, jsDir string) error {
	args := []string{"-jar", compiler, "--language_in", "ECMASCRIPT5"}
	err := filepath.Walk(filepath.Join(jsDir, "app"), func(path string, info os.FileInfo, err error) error {
		if filepath.Ext(path) == ".js" {
			args = append(args, "--js", path)
		}
		return nil
	})
	if err != nil {
		return err
	}

	cmd := exec.Command("java", args...)
	var out bytes.Buffer
	cmd.Stdout = &out
	var er bytes.Buffer
	cmd.Stderr = &er
	err = cmd.Run()
	if err != nil {
		return err
	}

	// merge with lib file
	libs := []string{}
	filepath.Walk(filepath.Join(jsDir, "lib"), func(path string, info os.FileInfo, err error) error {
		if filepath.Ext(path) == ".js" {
			libs = append(libs, path)
		}
		return nil
	})
	if err != nil {
		return err
	}

	f, err := os.OpenFile(filepath.Join(jsDir, "../js/app.min.js"), os.O_WRONLY, 0600)
	if err != nil {
		return err
	}
	defer f.Close()

	for _, l := range libs {
		f2, err := os.Open(l)
		if err != nil {
			log.Println("Error when open lib file:", err)
			continue
		}
		defer f2.Close()

		_, err = io.Copy(f, f2)
		if err != nil {
			log.Println("Error when copy lib file content:", err)
		} else {
			f.Write([]byte{10, 13})
		}
	}

	_, err = io.Copy(f, &out)

	return err
}
<commit_msg>add debug flag<commit_after>package main

import (
	"bytes"
	"flag"
	"gopkg.in/fsnotify.v1"
	"io"
	"log"
	"net/http"
	"os"
	"os/exec"
	"path/filepath"
	"sync"
)

func main() {
	d, err := os.Getwd()
	if err != nil {
		d = "./"
	}

	port := flag.String("http", ":8081", "the listening port")
	dir := flag.String("dir", d, "the working folder")
	debug := flag.Bool("debug", false, "enable/disable debuging mode")
	flag.Parse()

	compiler := filepath.Join(*dir, "compiler.jar")
	_, err = os.Stat(compiler)
	compilerMissing := os.IsNotExist(err)

	jsDir := filepath.Join(*dir, "src")
	_, err = os.Stat(jsDir)
	jsDirMissing := os.IsNotExist(err)

	lock := sync.Mutex{}

	if compilerMissing {
		log.Println("Not found 'compiler.jar', please download it here: http://dl.google.com/closure-compiler/compiler-latest.zip")
	} else if jsDirMissing {
		log.Println("Not found 'src' dir")
	} else {
		watcher, err := fsnotify.NewWatcher()
		if err != nil {
			log.Fatal(err)
		}
		defer watcher.Close()

		go func() {
			for {
				select {
				case event := <-watcher.Events:
					// call compile only for script change
					if event.Op != fsnotify.Chmod && event.Op != fsnotify.Rename {
						if ext := filepath.Ext(event.Name); ext != ".js" && ext != ".json" {
							break
						}

						lock.Lock()
						log.Println("Compiling script...")
						err := compileScript(compiler, jsDir, *debug)
						if err != nil {
							log.Println("Error compiling script:", err)
						}
						log.Println("Done!")
						lock.Unlock()
					}
				case err := <-watcher.Errors:
					log.Println("watcher error:", err)
				}
			}
		}()

		// watch all sub-dir
		err = filepath.Walk(jsDir, func(path string, info os.FileInfo, err error) error {
			if info.IsDir() {
				err := watcher.Add(path)
				if err != nil {
					return err
				}
			}
			return nil
		})
		if err != nil {
			log.Fatal(err)
		}
	}

	lock.Lock()
	err = compileScript(compiler, jsDir, *debug)
	if err != nil {
		log.Println("Error compiling script:", err)
	} else {
		log.Println("Compiling script at", jsDir)
	}
	lock.Unlock()

	log.Println("Serving client at http://localhost" + *port)
	panic(http.ListenAndServe(*port, http.FileServer(http.Dir(*dir))))
}

func compileScript(compiler, jsDir string, debug bool) error {
	args := []string{"-jar", compiler, "--language_in", "ECMASCRIPT5"}
	if debug {
		args = []string{"-jar", compiler, "--language_in", "ECMASCRIPT5", "--formatting", "PRETTY_PRINT", "--compilation_level", "SIMPLE"}
	}
	err := filepath.Walk(filepath.Join(jsDir, "app"), func(path string, info os.FileInfo, err error) error {
		if filepath.Ext(path) == ".js" {
			args = append(args, "--js", path)
		}
		return nil
	})
	if err != nil {
		return err
	}

	cmd := exec.Command("java", args...)
	var out bytes.Buffer
	cmd.Stdout = &out
	var er bytes.Buffer
	cmd.Stderr = &er
	err = cmd.Run()
	if err != nil {
		return err
	}

	// merge with lib file
	libs := []string{}
	filepath.Walk(filepath.Join(jsDir, "lib"), func(path string, info os.FileInfo, err error) error {
		if filepath.Ext(path) == ".js" {
			libs = append(libs, path)
		}
		return nil
	})
	if err != nil {
		return err
	}

	f, err := os.OpenFile(filepath.Join(jsDir, "../js/app.min.js"), os.O_WRONLY, 0600)
	if err != nil {
		return err
	}
	defer f.Close()

	for _, l := range libs {
		f2, err := os.Open(l)
		if err != nil {
			log.Println("Error when open lib file:", err)
			continue
		}
		defer f2.Close()

		_, err = io.Copy(f, f2)
		if err != nil {
			log.Println("Error when copy lib file content:", err)
		} else {
			f.Write([]byte{10, 13})
		}
	}

	_, err = io.Copy(f, &out)

	return err
}
<|endoftext|>


<commit_before>package ergo

type Ergo struct {
	*Route
}

func New(path string) *Ergo {
	return &Ergo{
		Route: NewRoute(path),
	}
}

func (e *Ergo) Schemes(s ...string) *Ergo {
	schemes(e, s)
	return e
}

func (e *Ergo) Consumes(mimes ...string) *Ergo {
	consumes(e, mimes)
	return e
}

func (e *Ergo) Produces(mimes ...string) *Ergo {
	produces(e, mimes)
	return e
}

func (e *Ergo) Params(params ...*Param) *Ergo {
	addParams(e, params...)
	return e
}

func (e *Ergo) ResetParams(params ...*Param) *Ergo {
	e.setParamsSlice(params...)
	return e
}

func (e *Ergo) SetParams(params map[string]*Param) *Ergo {
	e.setParams(params)
	return e
}

func (e *Ergo) IgnoreParams(params ...string) *Ergo {
	ignoreParams(e, params...)
	return e
}

func (e *Ergo) IgnoreParamsBut(params ...string) *Ergo {
	ignoreParamsBut(e, params...)
	return e
}

<commit_msg>NotFoundHandler setter for Ergo<commit_after>package ergo

type Ergo struct {
	*Route
}

func New(path string) *Ergo {
	return &Ergo{
		Route: NewRoute(path),
	}
}

func (e *Ergo) Schemes(s ...string) *Ergo {
	schemes(e, s)
	return e
}

func (e *Ergo) Consumes(mimes ...string) *Ergo {
	consumes(e, mimes)
	return e
}

func (e *Ergo) Produces(mimes ...string) *Ergo {
	produces(e, mimes)
	return e
}

func (e *Ergo) Params(params ...*Param) *Ergo {
	addParams(e, params...)
	return e
}

func (e *Ergo) ResetParams(params ...*Param) *Ergo {
	e.setParamsSlice(params...)
	return e
}

func (e *Ergo) SetParams(params map[string]*Param) *Ergo {
	e.setParams(params)
	return e
}

func (e *Ergo) IgnoreParams(params ...string) *Ergo {
	ignoreParams(e, params...)
	return e
}

func (e *Ergo) IgnoreParamsBut(params ...string) *Ergo {
	ignoreParamsBut(e, params...)
	return e
}

func (e *Ergo) NotFoundHandler(h Handler) *Ergo {
	e.Route.notFoundHandler = h
	return e
}

<|endoftext|>


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

package kvstore

import (
	"encoding/base64"
	"errors"
	"fmt"
	"strings"
	"time"

	"github.com/cilium/cilium/pkg/backoff"
	"github.com/cilium/cilium/pkg/controller"
	"github.com/cilium/cilium/pkg/logging/logfields"

	consulAPI "github.com/hashicorp/consul/api"
	"github.com/sirupsen/logrus"
)

const (
	consulName = "consul"

	// optAddress is the string representing the key mapping to the value of the
	// address for Consul.
	optAddress = "consul.address"

	// maxLockRetries is the number of retries attempted when acquiring a lock
	maxLockRetries = 10
)

type consulModule struct {
	opts   backendOptions
	config *consulAPI.Config
}

var (
	//consulDummyAddress can be overwritten from test invokers using ldflags
	consulDummyAddress = "127.0.0.1:8501"

	module = &consulModule{
		opts: backendOptions{
			optAddress: &backendOption{
				description: "Addresses of consul cluster",
			},
		},
	}
)

func init() {
	// register consul module for use
	registerBackend(consulName, module)
}

func (c *consulModule) createInstance() backendModule {
	cpy := *module
	return &cpy
}

func (c *consulModule) getName() string {
	return consulName
}

func (c *consulModule) setConfigDummy() {
	c.config = consulAPI.DefaultConfig()
	c.config.Address = consulDummyAddress
}

func (c *consulModule) setConfig(opts map[string]string) error {
	return setOpts(opts, c.opts)
}

func (c *consulModule) getConfig() map[string]string {
	return getOpts(c.opts)
}

func (c *consulModule) newClient() (BackendOperations, error) {
	if c.config == nil {
		consulAddr, ok := c.opts[optAddress]
		if !ok {
			return nil, fmt.Errorf("invalid consul configuration, please specify %s option", optAddress)
		}

		addr := consulAddr.value
		consulSplitAddr := strings.Split(addr, "://")
		if len(consulSplitAddr) == 2 {
			addr = consulSplitAddr[1]
		} else if len(consulSplitAddr) == 1 {
			addr = consulSplitAddr[0]
		}

		c.config = consulAPI.DefaultConfig()
		c.config.Address = addr
	}

	client, err := newConsulClient(c.config)
	if err != nil {
		return nil, err
	}

	return client, nil
}

var (
	maxRetries = 30
)

type consulClient struct {
	*consulAPI.Client
	lease       string
	controllers *controller.Manager
}

func newConsulClient(config *consulAPI.Config) (BackendOperations, error) {
	var (
		c   *consulAPI.Client
		err error
	)
	if config != nil {
		c, err = consulAPI.NewClient(config)
	} else {
		c, err = consulAPI.NewClient(consulAPI.DefaultConfig())
	}
	if err != nil {
		return nil, err
	}

	boff := backoff.Exponential{Min: time.Duration(100) * time.Millisecond}
	log.Info("Waiting for consul to elect a leader")

	for i := 0; i < maxRetries; i++ {
		var leader string
		leader, err = c.Status().Leader()

		if err == nil {
			if leader != "" {
				// happy path
				break
			} else {
				err = errors.New("timeout while waiting for leader to be elected")
			}
		}

		boff.Wait()
	}

	if err != nil {
		log.WithError(err).Fatal("Unable to contact consul server")
	}

	entry := &consulAPI.SessionEntry{
		TTL:      fmt.Sprintf("%ds", int(LeaseTTL.Seconds())),
		Behavior: consulAPI.SessionBehaviorDelete,
	}

	lease, _, err := c.Session().Create(entry, nil)
	if err != nil {
		return nil, fmt.Errorf("unable to create default lease: %s", err)
	}

	client := &consulClient{
		Client:      c,
		lease:       lease,
		controllers: controller.NewManager(),
	}

	client.controllers.UpdateController(fmt.Sprintf("consul-lease-keepalive-%p", c),
		controller.ControllerParams{
			DoFunc: func() error {
				_, _, err := c.Session().Renew(lease, nil)
				return err
			},
			RunInterval: KeepAliveInterval,
		},
	)

	return client, nil
}

func (c *consulClient) LockPath(path string) (kvLocker, error) {
	lockKey, err := c.LockOpts(&consulAPI.LockOptions{Key: getLockPath(path)})
	if err != nil {
		return nil, err
	}

	for retries := 0; retries < maxLockRetries; retries++ {
		ch, err := lockKey.Lock(nil)
		switch {
		case err != nil:
			return nil, err
		case ch == nil && err == nil:
			Trace("Acquiring lock timed out, retrying", nil, logrus.Fields{fieldKey: path, logfields.Attempt: retries})
		default:
			return lockKey, err
		}
	}

	return nil, fmt.Errorf("maximum retries (%d) reached", maxLockRetries)
}

// Watch starts watching for changes in a prefix
func (c *consulClient) Watch(w *Watcher) {
	// Last known state of all KVPairs matching the prefix
	localState := map[string]consulAPI.KVPair{}
	nextIndex := uint64(0)

	qo := consulAPI.QueryOptions{
		WaitTime: time.Second,
	}

	for {
		// Initialize sleep time to a millisecond as we don't
		// want to sleep in between successful watch cycles
		sleepTime := 1 * time.Millisecond

		qo.WaitIndex = nextIndex
		pairs, q, err := c.KV().List(w.prefix, &qo)
		if err != nil {
			sleepTime = 5 * time.Second
			Trace("List of Watch failed", err, logrus.Fields{fieldPrefix: w.prefix, fieldWatcher: w.name})
		}

		if q != nil {
			nextIndex = q.LastIndex
		}

		// timeout while watching for changes, re-schedule
		if qo.WaitIndex != 0 && (q == nil || q.LastIndex == qo.WaitIndex) {
			goto wait
		}

		for _, newPair := range pairs {
			oldPair, ok := localState[newPair.Key]

			// Keys reported for the first time must be new
			if !ok {
				if newPair.CreateIndex != newPair.ModifyIndex {
					log.Debugf("consul: Previously unknown key %s received with CreateIndex(%d) != ModifyIndex(%d)",
						newPair.Key, newPair.CreateIndex, newPair.ModifyIndex)
				}

				w.Events <- KeyValueEvent{
					Typ:   EventTypeCreate,
					Key:   newPair.Key,
					Value: newPair.Value,
				}
			} else if oldPair.ModifyIndex != newPair.ModifyIndex {
				w.Events <- KeyValueEvent{
					Typ:   EventTypeModify,
					Key:   newPair.Key,
					Value: newPair.Value,
				}
			}

			// Everything left on localState will be assumed to
			// have been deleted, therefore remove all keys in
			// localState that still exist in the kvstore
			delete(localState, newPair.Key)
		}

		for k, deletedPair := range localState {
			w.Events <- KeyValueEvent{
				Typ:   EventTypeDelete,
				Key:   deletedPair.Key,
				Value: deletedPair.Value,
			}
			delete(localState, k)
		}

		for _, newPair := range pairs {
			localState[newPair.Key] = *newPair

		}

		// Initial list operation has been completed, signal this
		if qo.WaitIndex == 0 {
			w.Events <- KeyValueEvent{Typ: EventTypeListDone}
		}

	wait:
		select {
		case <-time.After(sleepTime):
		case <-w.stopWatch:
			close(w.Events)
			w.stopWait.Done()
			return
		}
	}
}

func (c *consulClient) Status() (string, error) {
	leader, err := c.Client.Status().Leader()
	return "Consul: " + leader, err
}

func (c *consulClient) DeletePrefix(path string) error {
	increaseMetric(path, metricDelete, "DeletePrefix")
	_, err := c.Client.KV().DeleteTree(path, nil)
	return err
}

// Set sets value of key
func (c *consulClient) Set(key string, value []byte) error {
	increaseMetric(key, metricSet, "Set")
	_, err := c.KV().Put(&consulAPI.KVPair{Key: key, Value: value}, nil)
	return err
}

// Delete deletes a key
func (c *consulClient) Delete(key string) error {
	increaseMetric(key, metricDelete, "Delete")
	_, err := c.KV().Delete(key, nil)
	return err
}

// Get returns value of key
func (c *consulClient) Get(key string) ([]byte, error) {
	increaseMetric(key, metricRead, "Get")
	pair, _, err := c.KV().Get(key, nil)
	if err != nil {
		return nil, err
	}
	if pair == nil {
		return nil, nil
	}
	return pair.Value, nil
}

// GetPrefix returns the first key which matches the prefix
func (c *consulClient) GetPrefix(prefix string) ([]byte, error) {
	increaseMetric(prefix, metricRead, "GetPrefix")
	pairs, _, err := c.KV().List(prefix, nil)
	if err != nil {
		return nil, err
	}

	if len(pairs) == 0 {
		return nil, nil
	}

	return pairs[0].Value, nil
}

// Update creates or updates a key with the value
func (c *consulClient) Update(key string, value []byte, lease bool) error {
	increaseMetric(key, metricSet, "Update")
	k := &consulAPI.KVPair{Key: key, Value: value}

	if lease {
		k.Session = c.lease
	}

	_, err := c.KV().Put(k, nil)
	return err
}

// CreateOnly creates a key with the value and will fail if the key already exists
func (c *consulClient) CreateOnly(key string, value []byte, lease bool) error {
	increaseMetric(key, metricSet, "CreateOnly")
	k := &consulAPI.KVPair{
		Key:         key,
		Value:       value,
		CreateIndex: 0,
	}

	if lease {
		k.Session = c.lease
	}

	success, _, err := c.KV().CAS(k, nil)
	if err != nil {
		return fmt.Errorf("unable to compare-and-swap: %s", err)
	}
	if !success {
		return fmt.Errorf("compare-and-swap unsuccessful")
	}

	return nil
}

// CreateIfExists creates a key with the value only if key condKey exists
func (c *consulClient) CreateIfExists(condKey, key string, value []byte, lease bool) error {
	// Consul does not support transactions which would allow to check for
	// the presence of a conditional key if the key is not the key being
	// manipulated
	//
	// Lock the conditional key to serialize all CreateIfExists() calls

	increaseMetric(key, metricSet, "CreateIfExists")
	l, err := LockPath(condKey)
	if err != nil {
		return fmt.Errorf("unable to lock condKey for CreateIfExists: %s", err)
	}

	defer l.Unlock()

	// Create the key if it does not exist
	if err := c.CreateOnly(key, value, lease); err != nil {
		return err
	}

	// Consul does not support transactions which would allow to check for
	// the presence of another key
	masterKey, err := c.Get(condKey)
	if err != nil || masterKey == nil {
		c.Delete(key)
		return fmt.Errorf("conditional key not present")
	}

	return nil
}

// ListPrefix returns a map of matching keys
func (c *consulClient) ListPrefix(prefix string) (KeyValuePairs, error) {
	increaseMetric(prefix, metricRead, "ListPrefix")
	pairs, _, err := c.KV().List(prefix, nil)
	if err != nil {
		return nil, err
	}

	p := KeyValuePairs(make(map[string][]byte, len(pairs)))
	for i := 0; i < len(pairs); i++ {
		p[pairs[i].Key] = pairs[i].Value
	}

	return p, nil
}

// Close closes the consul session
func (c *consulClient) Close() {
	if c.controllers != nil {
		c.controllers.RemoveAll()
	}
	if c.lease != "" {
		c.Session().Destroy(c.lease, nil)
	}
}

// GetCapabilities returns the capabilities of the backend
func (c *consulClient) GetCapabilities() Capabilities {
	return Capabilities(0)
}

// Encode encodes a binary slice into a character set that the backend supports
func (c *consulClient) Encode(in []byte) string {
	return base64.URLEncoding.EncodeToString([]byte(in))
}

// Decode decodes a key previously encoded back into the original binary slice
func (c *consulClient) Decode(in string) ([]byte, error) {
	return base64.URLEncoding.DecodeString(in)
}

// ListAndWatch implements the BackendOperations.ListAndWatch using consul
func (c *consulClient) ListAndWatch(name, prefix string, chanSize int) *Watcher {
	w := newWatcher(name, prefix, chanSize)

	log.WithField(fieldWatcher, w).Debug("Starting watcher...")

	go c.Watch(w)

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

package kvstore

import (
	"encoding/base64"
	"errors"
	"fmt"
	"io/ioutil"
	"time"

	"github.com/cilium/cilium/pkg/backoff"
	"github.com/cilium/cilium/pkg/controller"
	"github.com/cilium/cilium/pkg/logging/logfields"

	consulAPI "github.com/hashicorp/consul/api"
	"github.com/sirupsen/logrus"
	"gopkg.in/yaml.v2"
)

const (
	consulName = "consul"

	// optAddress is the string representing the key mapping to the value of the
	// address for Consul.
	optAddress         = "consul.address"
	consulOptionConfig = "consul.tlsconfig"

	// maxLockRetries is the number of retries attempted when acquiring a lock
	maxLockRetries = 10
)

type consulModule struct {
	opts   backendOptions
	config *consulAPI.Config
}

var (
	//consulDummyAddress can be overwritten from test invokers using ldflags
	consulDummyAddress = "127.0.0.1:8501"

	module = &consulModule{
		opts: backendOptions{
			optAddress: &backendOption{
				description: "Addresses of consul cluster",
			},
			consulOptionConfig: &backendOption{
				description: "Path to consul tls configuration file",
			},
		},
	}
)

func init() {
	// register consul module for use
	registerBackend(consulName, module)
}

func (c *consulModule) createInstance() backendModule {
	cpy := *module
	return &cpy
}

func (c *consulModule) getName() string {
	return consulName
}

func (c *consulModule) setConfigDummy() {
	c.config = consulAPI.DefaultConfig()
	c.config.Address = consulDummyAddress
}

func (c *consulModule) setConfig(opts map[string]string) error {
	return setOpts(opts, c.opts)
}

func (c *consulModule) getConfig() map[string]string {
	return getOpts(c.opts)
}

func (c *consulModule) newClient() (BackendOperations, error) {
	if c.config == nil {
		consulAddr, consulAddrSet := c.opts[optAddress]
		configPathOpt, configPathOptSet := c.opts[consulOptionConfig]
		if !consulAddrSet {
			return nil, fmt.Errorf("invalid consul configuration, please specify %s option", optAddress)
		}

		if consulAddr.value == "" {
			return nil, fmt.Errorf("invalid consul configuration, please specify %s option", optAddress)
		}

		addr := consulAddr.value
		c.config = consulAPI.DefaultConfig()
		if configPathOptSet && configPathOpt.value != "" {
			b, err := ioutil.ReadFile(configPathOpt.value)
			if err != nil {
				return nil, fmt.Errorf("unable to read consul tls configuration file %s: %s", configPathOpt.value, err)
			}
			yc := consulAPI.TLSConfig{}
			err = yaml.Unmarshal(b, &yc)
			if err != nil {
				return nil, fmt.Errorf("invalid consul tls configuration in %s: %s", configPathOpt.value, err)
			}
			c.config.TLSConfig = yc
		}

		c.config.Address = addr

	}
	client, err := newConsulClient(c.config)
	if err != nil {
		return nil, err
	}

	return client, nil
}

var (
	maxRetries = 30
)

type consulClient struct {
	*consulAPI.Client
	lease       string
	controllers *controller.Manager
}

func newConsulClient(config *consulAPI.Config) (BackendOperations, error) {
	var (
		c   *consulAPI.Client
		err error
	)
	if config != nil {
		c, err = consulAPI.NewClient(config)
	} else {
		c, err = consulAPI.NewClient(consulAPI.DefaultConfig())
	}
	if err != nil {
		return nil, err
	}

	boff := backoff.Exponential{Min: time.Duration(100) * time.Millisecond}
	log.Info("Waiting for consul to elect a leader")

	for i := 0; i < maxRetries; i++ {
		var leader string
		leader, err = c.Status().Leader()

		if err == nil {
			if leader != "" {
				// happy path
				break
			} else {
				err = errors.New("timeout while waiting for leader to be elected")
			}
		}

		boff.Wait()
	}

	if err != nil {
		log.WithError(err).Fatal("Unable to contact consul server")
	}

	entry := &consulAPI.SessionEntry{
		TTL:      fmt.Sprintf("%ds", int(LeaseTTL.Seconds())),
		Behavior: consulAPI.SessionBehaviorDelete,
	}

	lease, _, err := c.Session().Create(entry, nil)
	if err != nil {
		return nil, fmt.Errorf("unable to create default lease: %s", err)
	}

	client := &consulClient{
		Client:      c,
		lease:       lease,
		controllers: controller.NewManager(),
	}

	client.controllers.UpdateController(fmt.Sprintf("consul-lease-keepalive-%p", c),
		controller.ControllerParams{
			DoFunc: func() error {
				_, _, err := c.Session().Renew(lease, nil)
				return err
			},
			RunInterval: KeepAliveInterval,
		},
	)

	return client, nil
}

func (c *consulClient) LockPath(path string) (kvLocker, error) {
	lockKey, err := c.LockOpts(&consulAPI.LockOptions{Key: getLockPath(path)})
	if err != nil {
		return nil, err
	}

	for retries := 0; retries < maxLockRetries; retries++ {
		ch, err := lockKey.Lock(nil)
		switch {
		case err != nil:
			return nil, err
		case ch == nil && err == nil:
			Trace("Acquiring lock timed out, retrying", nil, logrus.Fields{fieldKey: path, logfields.Attempt: retries})
		default:
			return lockKey, err
		}
	}

	return nil, fmt.Errorf("maximum retries (%d) reached", maxLockRetries)
}

// Watch starts watching for changes in a prefix
func (c *consulClient) Watch(w *Watcher) {
	// Last known state of all KVPairs matching the prefix
	localState := map[string]consulAPI.KVPair{}
	nextIndex := uint64(0)

	qo := consulAPI.QueryOptions{
		WaitTime: time.Second,
	}

	for {
		// Initialize sleep time to a millisecond as we don't
		// want to sleep in between successful watch cycles
		sleepTime := 1 * time.Millisecond

		qo.WaitIndex = nextIndex
		pairs, q, err := c.KV().List(w.prefix, &qo)
		if err != nil {
			sleepTime = 5 * time.Second
			Trace("List of Watch failed", err, logrus.Fields{fieldPrefix: w.prefix, fieldWatcher: w.name})
		}

		if q != nil {
			nextIndex = q.LastIndex
		}

		// timeout while watching for changes, re-schedule
		if qo.WaitIndex != 0 && (q == nil || q.LastIndex == qo.WaitIndex) {
			goto wait
		}

		for _, newPair := range pairs {
			oldPair, ok := localState[newPair.Key]

			// Keys reported for the first time must be new
			if !ok {
				if newPair.CreateIndex != newPair.ModifyIndex {
					log.Debugf("consul: Previously unknown key %s received with CreateIndex(%d) != ModifyIndex(%d)",
						newPair.Key, newPair.CreateIndex, newPair.ModifyIndex)
				}

				w.Events <- KeyValueEvent{
					Typ:   EventTypeCreate,
					Key:   newPair.Key,
					Value: newPair.Value,
				}
			} else if oldPair.ModifyIndex != newPair.ModifyIndex {
				w.Events <- KeyValueEvent{
					Typ:   EventTypeModify,
					Key:   newPair.Key,
					Value: newPair.Value,
				}
			}

			// Everything left on localState will be assumed to
			// have been deleted, therefore remove all keys in
			// localState that still exist in the kvstore
			delete(localState, newPair.Key)
		}

		for k, deletedPair := range localState {
			w.Events <- KeyValueEvent{
				Typ:   EventTypeDelete,
				Key:   deletedPair.Key,
				Value: deletedPair.Value,
			}
			delete(localState, k)
		}

		for _, newPair := range pairs {
			localState[newPair.Key] = *newPair

		}

		// Initial list operation has been completed, signal this
		if qo.WaitIndex == 0 {
			w.Events <- KeyValueEvent{Typ: EventTypeListDone}
		}

	wait:
		select {
		case <-time.After(sleepTime):
		case <-w.stopWatch:
			close(w.Events)
			w.stopWait.Done()
			return
		}
	}
}

func (c *consulClient) Status() (string, error) {
	leader, err := c.Client.Status().Leader()
	return "Consul: " + leader, err
}

func (c *consulClient) DeletePrefix(path string) error {
	increaseMetric(path, metricDelete, "DeletePrefix")
	_, err := c.Client.KV().DeleteTree(path, nil)
	return err
}

// Set sets value of key
func (c *consulClient) Set(key string, value []byte) error {
	increaseMetric(key, metricSet, "Set")
	_, err := c.KV().Put(&consulAPI.KVPair{Key: key, Value: value}, nil)
	return err
}

// Delete deletes a key
func (c *consulClient) Delete(key string) error {
	increaseMetric(key, metricDelete, "Delete")
	_, err := c.KV().Delete(key, nil)
	return err
}

// Get returns value of key
func (c *consulClient) Get(key string) ([]byte, error) {
	increaseMetric(key, metricRead, "Get")
	pair, _, err := c.KV().Get(key, nil)
	if err != nil {
		return nil, err
	}
	if pair == nil {
		return nil, nil
	}
	return pair.Value, nil
}

// GetPrefix returns the first key which matches the prefix
func (c *consulClient) GetPrefix(prefix string) ([]byte, error) {
	increaseMetric(prefix, metricRead, "GetPrefix")
	pairs, _, err := c.KV().List(prefix, nil)
	if err != nil {
		return nil, err
	}

	if len(pairs) == 0 {
		return nil, nil
	}

	return pairs[0].Value, nil
}

// Update creates or updates a key with the value
func (c *consulClient) Update(key string, value []byte, lease bool) error {
	increaseMetric(key, metricSet, "Update")
	k := &consulAPI.KVPair{Key: key, Value: value}

	if lease {
		k.Session = c.lease
	}

	_, err := c.KV().Put(k, nil)
	return err
}

// CreateOnly creates a key with the value and will fail if the key already exists
func (c *consulClient) CreateOnly(key string, value []byte, lease bool) error {
	increaseMetric(key, metricSet, "CreateOnly")
	k := &consulAPI.KVPair{
		Key:         key,
		Value:       value,
		CreateIndex: 0,
	}

	if lease {
		k.Session = c.lease
	}

	success, _, err := c.KV().CAS(k, nil)
	if err != nil {
		return fmt.Errorf("unable to compare-and-swap: %s", err)
	}
	if !success {
		return fmt.Errorf("compare-and-swap unsuccessful")
	}

	return nil
}

// CreateIfExists creates a key with the value only if key condKey exists
func (c *consulClient) CreateIfExists(condKey, key string, value []byte, lease bool) error {
	// Consul does not support transactions which would allow to check for
	// the presence of a conditional key if the key is not the key being
	// manipulated
	//
	// Lock the conditional key to serialize all CreateIfExists() calls

	increaseMetric(key, metricSet, "CreateIfExists")
	l, err := LockPath(condKey)
	if err != nil {
		return fmt.Errorf("unable to lock condKey for CreateIfExists: %s", err)
	}

	defer l.Unlock()

	// Create the key if it does not exist
	if err := c.CreateOnly(key, value, lease); err != nil {
		return err
	}

	// Consul does not support transactions which would allow to check for
	// the presence of another key
	masterKey, err := c.Get(condKey)
	if err != nil || masterKey == nil {
		c.Delete(key)
		return fmt.Errorf("conditional key not present")
	}

	return nil
}

// ListPrefix returns a map of matching keys
func (c *consulClient) ListPrefix(prefix string) (KeyValuePairs, error) {
	increaseMetric(prefix, metricRead, "ListPrefix")
	pairs, _, err := c.KV().List(prefix, nil)
	if err != nil {
		return nil, err
	}

	p := KeyValuePairs(make(map[string][]byte, len(pairs)))
	for i := 0; i < len(pairs); i++ {
		p[pairs[i].Key] = pairs[i].Value
	}

	return p, nil
}

// Close closes the consul session
func (c *consulClient) Close() {
	if c.controllers != nil {
		c.controllers.RemoveAll()
	}
	if c.lease != "" {
		c.Session().Destroy(c.lease, nil)
	}
}

// GetCapabilities returns the capabilities of the backend
func (c *consulClient) GetCapabilities() Capabilities {
	return Capabilities(0)
}

// Encode encodes a binary slice into a character set that the backend supports
func (c *consulClient) Encode(in []byte) string {
	return base64.URLEncoding.EncodeToString([]byte(in))
}

// Decode decodes a key previously encoded back into the original binary slice
func (c *consulClient) Decode(in string) ([]byte, error) {
	return base64.URLEncoding.DecodeString(in)
}

// ListAndWatch implements the BackendOperations.ListAndWatch using consul
func (c *consulClient) ListAndWatch(name, prefix string, chanSize int) *Watcher {
	w := newWatcher(name, prefix, chanSize)

	log.WithField(fieldWatcher, w).Debug("Starting watcher...")

	go c.Watch(w)

	return w
}
<|endoftext|>