<commit_before>package git

import (
	"crypto/sha1"
	"encoding/binary"
	"errors"
	"fmt"
	"io"
	"io/ioutil"
	"os"
	"sort"
)

var InvalidIndex error = errors.New("Invalid index")

// index file is defined as Network byte Order (Big Endian)

// 12 byte header:
// 4 bytes: D I R C (stands for "Dir cache")
// 4-byte version number (can be 2, 3 or 4)
// 32bit number of index entries
type fixedGitIndex struct {
	Signature          [4]byte // 4
	Version            uint32  // 8
	NumberIndexEntries uint32  // 12
}

type Index struct {
	fixedGitIndex // 12
	Objects       []*IndexEntry
}
type IndexEntry struct {
	FixedIndexEntry

	PathName IndexPath
}

func (ie IndexEntry) Stage() Stage {
	return Stage((ie.Flags >> 12) & 0x3)
}
func NewIndex() *Index {
	return &Index{
		fixedGitIndex: fixedGitIndex{
			Signature:          [4]byte{'D', 'I', 'R', 'C'},
			Version:            2,
			NumberIndexEntries: 0,
		},
		Objects: make([]*IndexEntry, 0),
	}
}

type FixedIndexEntry struct {
	Ctime     uint32 // 16
	Ctimenano uint32 // 20

	Mtime     uint32 // 24
	Mtimenano uint32 // 28

	Dev uint32 // 32
	Ino uint32 // 36

	Mode EntryMode // 40

	Uid uint32 // 44
	Gid uint32 // 48

	Fsize uint32 // 52

	Sha1 Sha1 // 72

	Flags uint16 // 74
}

func (d GitDir) ReadIndex() (*Index, error) {
	file, err := d.Open("index")
	if err != nil {
		return &Index{
			fixedGitIndex{
				[4]byte{'D', 'I', 'R', 'C'},
				2, // version 2
				0, // no entries
			},
			make([]*IndexEntry, 0),
		}, err
	}
	defer file.Close()

	var i fixedGitIndex
	binary.Read(file, binary.BigEndian, &i)
	if i.Signature != [4]byte{'D', 'I', 'R', 'C'} {
		return nil, InvalidIndex
	}
	if i.Version < 2 || i.Version > 4 {
		return nil, InvalidIndex
	}

	var idx uint32
	indexes := make([]*IndexEntry, i.NumberIndexEntries, i.NumberIndexEntries)
	for idx = 0; idx < i.NumberIndexEntries; idx += 1 {
		if index, err := ReadIndexEntry(file); err == nil {
			indexes[idx] = index
		}
	}
	return &Index{i, indexes}, nil
}

func ReadIndexEntry(file *os.File) (*IndexEntry, error) {
	var f FixedIndexEntry
	var name []byte
	binary.Read(file, binary.BigEndian, &f)

	var nameLength uint16
	nameLength = f.Flags & 0x0FFF

	if nameLength&0xFFF != 0xFFF {
		name = make([]byte, nameLength, nameLength)
		n, err := file.Read(name)
		if err != nil {
			panic("I don't know what to do")
		}
		if n != int(nameLength) {
			panic("Error reading the name")
		}

		// I don't understand where this +4 comes from, but it seems to work
		// out with how the c git implementation calculates the padding..
		//
		// The definition of the index file format at:
		// https://github.com/git/git/blob/master/Documentation/technical/index-format.txt
		// claims that there should be "1-8 nul bytes as necessary to pad the entry to a multiple of eight
		// bytes while keeping the name NUL-terminated."
		//
		// The fixed size of the header is 82 bytes if you add up all the types.
		// the length of the name is nameLength bytes, so according to the spec
		// this *should* be 8 - ((82 + nameLength) % 8) bytes of padding.
		// But reading existant index files, there seems to be an extra 4 bytes
		// incorporated into the index size calculation.
		expectedOffset := 8 - ((82 + nameLength + 4) % 8)
		file.Seek(int64(expectedOffset), 1)
		/*
			This was used to verify that the offset is correct, but it causes problems if the data following
			the offset is empty..
			whitespace := make([]byte, 1, 1)
			var w uint16
			// Read all the whitespace that git uses for alignment.
			for _, _ = file.Read(whitespace); whitespace[0] == 0; _, _ = file.Read(whitespace) {
				w += 1
			}

			if w % 8 != expectedOffset {
				panic(fmt.Sprintf("Read incorrect number of whitespace characters %d vs %d", w, expectedOffset))
			}
			if w == 0 {
				panic("Name was not null terminated in index")
			}

			// Undo the last read, which wasn't whitespace..
			file.Seek(-1, 1)
		*/

	} else {
		panic("TODO: I can't handle such long names yet")
	}
	return &IndexEntry{f, IndexPath(name)}, nil
}

// A Stage represents a git merge stage in the index.
type Stage uint8

// Valid merge stages.
const (
	Stage0 = Stage(iota)
	Stage1
	Stage2
	Stage3
)

// Adds an entry to the index with Sha1 s and stage stage during a merge.
// If an entry already exists for this pathname/stage, it will be overwritten,
// otherwise it will be added if createEntry is true, and return an error if not.
//
// As a special case, if something is added as Stage0, then Stage1-3 entries
// will be removed.
func (g *Index) AddStage(c *Client, path IndexPath, s Sha1, stage Stage, mtime, mtimenano, size uint32, createEntry bool) error {
	if stage == Stage0 {
		defer g.RemoveUnmergedStages(c, path)
	}

	// Update the existing stage, if it exists.
	for _, entry := range g.Objects {
		if entry.PathName == path && entry.Stage() == stage {
			entry.Sha1 = s
			entry.Mtime = mtime
			entry.Mtimenano = mtimenano
			entry.Fsize = size

			// We found and updated the entry, no need to continue
			return nil
		}
	}

	if !createEntry {
		return fmt.Errorf("%v not found in index", path)
	}
	// There was no path/stage combo already in the index. Add it.

	// According to the git documentation:
	// Flags is
	//    A 16-bit 'flags' field split into (high to low bits)
	//
	//       1-bit assume-valid flag
	//
	//       1-bit extended flag (must be zero in version 2)
	//
	//       2-bit stage (during merge)
	//       12-bit name length if the length is less than 0xFFF; otherwise 0xFFF
	//     is stored in this field.

	// So we'll construct the flags based on what we know.

	var flags = uint16(stage) << 12 // start with the stage.
	// Add the name length.
	if len(path) >= 0x0FFF {
		flags |= 0x0FFF
	} else {
		flags |= (uint16(len(path)) & 0x0FFF)
	}

	g.Objects = append(g.Objects, &IndexEntry{
		FixedIndexEntry{
			0, //uint32(csec),
			0, //uint32(cnano),
			mtime,
			mtimenano,
			0,        //uint32(stat.Dev),
			0,        //uint32(stat.Ino),
			ModeBlob, // Directories are never added, only their files, so assume blob
			0,        //stat.Uid,
			0,        //stat.Gid,
			size,
			s,
			flags,
		},
		path,
	})
	g.NumberIndexEntries += 1
	sort.Sort(ByPath(g.Objects))
	return nil

}

// Remove any unmerged (non-stage 0) stage from the index for the given path
func (g *Index) RemoveUnmergedStages(c *Client, path IndexPath) error {
	// There are likely 3 things being deleted, so make a new slice
	newobjects := make([]*IndexEntry, 0, len(g.Objects))
	for _, entry := range g.Objects {
		stage := entry.Stage()
		if entry.PathName == path && stage == Stage0 {
			newobjects = append(newobjects, entry)
		} else if entry.PathName == path && stage != Stage0 {
			// do not add it, it's the wrong stage.
		} else {
			// It's a different Pathname, keep it.
			newobjects = append(newobjects, entry)
		}
	}
	g.Objects = newobjects
	g.NumberIndexEntries = uint32(len(newobjects))
	return nil
}

// Adds a file to the index, without writing it to disk.
// To write it to disk after calling this, use GitIndex.WriteIndex
//
// This will do the following:
// 1. write git object blob of file contents to .git/objects
// 2. normalize os.File name to path relative to gitRoot
// 3. search GitIndex for normalized name
//	if GitIndexEntry found
//		update GitIndexEntry to point to the new object blob
// 	else
// 		add new GitIndexEntry if not found and createEntry is true, error otherwise
//
func (g *Index) AddFile(c *Client, file *os.File, createEntry bool) error {
	contents, err := ioutil.ReadAll(file)
	if err != nil {
		return err
	}
	hash, err := c.WriteObject("blob", contents)
	if err != nil && err != ObjectExists {
		fmt.Fprintf(os.Stderr, "Error storing object: %s", err)
		return err
	}
	name, err := File(file.Name()).IndexPath(c)
	if err != nil {
		return err
	}

	fstat, err := file.Stat()
	if err != nil {
		return err
	}
	if fstat.IsDir() {
		// This should really recursively call add for each file in the directory.
		return fmt.Errorf("Add can't handle directories. yet.")
	}

	modTime := fstat.ModTime()
	return g.AddStage(
		c,
		name,
		hash,
		Stage0,
		uint32(modTime.Unix()),
		uint32(modTime.Nanosecond()),
		uint32(fstat.Size()),
		createEntry,
	)
}

type IndexStageEntry struct {
	IndexPath
	Stage
}

func (i *Index) GetStageMap() map[IndexStageEntry]*IndexEntry {
	r := make(map[IndexStageEntry]*IndexEntry)
	for _, entry := range i.Objects {
		r[IndexStageEntry{entry.PathName, entry.Stage()}] = entry
	}
	return r
}

type UnmergedPath struct {
	Stage1, Stage2, Stage3 *IndexEntry
}

func (i *Index) GetUnmerged() map[IndexPath]*UnmergedPath {
	r := make(map[IndexPath]*UnmergedPath)
	for _, entry := range i.Objects {
		if entry.Stage() != Stage0 {
			e, ok := r[entry.PathName]
			if !ok {
				e = &UnmergedPath{}
				r[entry.PathName] = e
			}
			switch entry.Stage() {
			case Stage1:
				e.Stage1 = entry
			case Stage2:
				e.Stage2 = entry
			case Stage3:
				e.Stage3 = entry
			}
		}
	}
	return r
}
func (i *Index) GetMap() map[IndexPath]*IndexEntry {
	r := make(map[IndexPath]*IndexEntry)
	for _, entry := range i.Objects {
		r[entry.PathName] = entry
	}
	return r
}

// Remove the first instance of file from the index. (This will usually
// be stage 0.)
func (g *Index) RemoveFile(file IndexPath) {
	for i, entry := range g.Objects {
		if entry.PathName == file {
			//println("Should remove ", i)
			g.Objects = append(g.Objects[:i], g.Objects[i+1:]...)
			g.NumberIndexEntries -= 1
			return
		}
	}
}

// This will write a new index file to w by doing the following:
// 1. Sort the objects in g.Index to ascending order based on name
// 2. Write g.fixedGitIndex to w
// 3. for each entry in g.Objects, write it to w.
// 4. Write the Sha1 of the contents of what was written
func (g Index) WriteIndex(file io.Writer) error {
	sort.Sort(ByPath(g.Objects))
	s := sha1.New()
	w := io.MultiWriter(file, s)
	binary.Write(w, binary.BigEndian, g.fixedGitIndex)
	for _, entry := range g.Objects {
		binary.Write(w, binary.BigEndian, entry.FixedIndexEntry)
		binary.Write(w, binary.BigEndian, []byte(entry.PathName))
		padding := 8 - ((82 + len(entry.PathName) + 4) % 8)
		p := make([]byte, padding)
		binary.Write(w, binary.BigEndian, p)
	}
	binary.Write(w, binary.BigEndian, s.Sum(nil))
	return nil
}

// Looks up the Sha1 of path currently stored in the index.
// Will return the 0 Sha1 if not found.
func (g Index) GetSha1(path IndexPath) Sha1 {
	for _, entry := range g.Objects {
		if entry.PathName == path {
			return entry.Sha1
		}
	}
	return Sha1{}
}

// Implement the sort interface on *GitIndexEntry, so that
// it's easy to sort by name.
type ByPath []*IndexEntry

func (g ByPath) Len() int      { return len(g) }
func (g ByPath) Swap(i, j int) { g[i], g[j] = g[j], g[i] }
func (g ByPath) Less(i, j int) bool {
	if g[i].PathName < g[j].PathName {
		return true
	} else if g[i].PathName == g[j].PathName {
		return g[i].Stage() < g[j].Stage()
	} else {
		return false
	}
}

// Replaces the index of Client with the the tree from the provided Treeish.
// if PreserveStatInfo is true, the stat information in the index won't be
// modified for existing entries.
func (g *Index) ResetIndex(c *Client, tree Treeish) error {
	newEntries, err := ExpandGitTreeIntoIndexes(c, tree, true, false)
	if err != nil {
		return err
	}
	g.NumberIndexEntries = uint32(len(newEntries))
	g.Objects = newEntries
	return nil
}

func (g Index) String() string {
	ret := ""

	for _, i := range g.Objects {
		ret += fmt.Sprintf("%v %v %v\n", i.Mode, i.Sha1, i.PathName)
	}
	return ret
}
<commit_msg>Fix bug where "git add" and "git update-index" failed before first git add<commit_after>package git

import (
	"crypto/sha1"
	"encoding/binary"
	"errors"
	"fmt"
	"io"
	"io/ioutil"
	"os"
	"sort"
)

var InvalidIndex error = errors.New("Invalid index")

// index file is defined as Network byte Order (Big Endian)

// 12 byte header:
// 4 bytes: D I R C (stands for "Dir cache")
// 4-byte version number (can be 2, 3 or 4)
// 32bit number of index entries
type fixedGitIndex struct {
	Signature          [4]byte // 4
	Version            uint32  // 8
	NumberIndexEntries uint32  // 12
}

type Index struct {
	fixedGitIndex // 12
	Objects       []*IndexEntry
}
type IndexEntry struct {
	FixedIndexEntry

	PathName IndexPath
}

func (ie IndexEntry) Stage() Stage {
	return Stage((ie.Flags >> 12) & 0x3)
}
func NewIndex() *Index {
	return &Index{
		fixedGitIndex: fixedGitIndex{
			Signature:          [4]byte{'D', 'I', 'R', 'C'},
			Version:            2,
			NumberIndexEntries: 0,
		},
		Objects: make([]*IndexEntry, 0),
	}
}

type FixedIndexEntry struct {
	Ctime     uint32 // 16
	Ctimenano uint32 // 20

	Mtime     uint32 // 24
	Mtimenano uint32 // 28

	Dev uint32 // 32
	Ino uint32 // 36

	Mode EntryMode // 40

	Uid uint32 // 44
	Gid uint32 // 48

	Fsize uint32 // 52

	Sha1 Sha1 // 72

	Flags uint16 // 74
}

// Reads the index file from the GitDir and returns a Index object.
// If the index file does not exist, it will return a new empty Index.
func (d GitDir) ReadIndex() (*Index, error) {
	file, err := d.Open("index")
	if err != nil {
		if os.IsNotExist(err) {
			// Is the file doesn't exist, treat it
			// as a new empty index.
			return &Index{
				fixedGitIndex{
					[4]byte{'D', 'I', 'R', 'C'},
					2, // version 2
					0, // no entries
				},
				make([]*IndexEntry, 0),
			}, nil
		}
		return nil, err
	}
	defer file.Close()

	var i fixedGitIndex
	binary.Read(file, binary.BigEndian, &i)
	if i.Signature != [4]byte{'D', 'I', 'R', 'C'} {
		return nil, InvalidIndex
	}
	if i.Version < 2 || i.Version > 4 {
		return nil, InvalidIndex
	}

	var idx uint32
	indexes := make([]*IndexEntry, i.NumberIndexEntries, i.NumberIndexEntries)
	for idx = 0; idx < i.NumberIndexEntries; idx += 1 {
		if index, err := ReadIndexEntry(file); err == nil {
			indexes[idx] = index
		}
	}
	return &Index{i, indexes}, nil
}

func ReadIndexEntry(file *os.File) (*IndexEntry, error) {
	var f FixedIndexEntry
	var name []byte
	binary.Read(file, binary.BigEndian, &f)

	var nameLength uint16
	nameLength = f.Flags & 0x0FFF

	if nameLength&0xFFF != 0xFFF {
		name = make([]byte, nameLength, nameLength)
		n, err := file.Read(name)
		if err != nil {
			panic("I don't know what to do")
		}
		if n != int(nameLength) {
			panic("Error reading the name")
		}

		// I don't understand where this +4 comes from, but it seems to work
		// out with how the c git implementation calculates the padding..
		//
		// The definition of the index file format at:
		// https://github.com/git/git/blob/master/Documentation/technical/index-format.txt
		// claims that there should be "1-8 nul bytes as necessary to pad the entry to a multiple of eight
		// bytes while keeping the name NUL-terminated."
		//
		// The fixed size of the header is 82 bytes if you add up all the types.
		// the length of the name is nameLength bytes, so according to the spec
		// this *should* be 8 - ((82 + nameLength) % 8) bytes of padding.
		// But reading existant index files, there seems to be an extra 4 bytes
		// incorporated into the index size calculation.
		expectedOffset := 8 - ((82 + nameLength + 4) % 8)
		file.Seek(int64(expectedOffset), 1)
		/*
			This was used to verify that the offset is correct, but it causes problems if the data following
			the offset is empty..
			whitespace := make([]byte, 1, 1)
			var w uint16
			// Read all the whitespace that git uses for alignment.
			for _, _ = file.Read(whitespace); whitespace[0] == 0; _, _ = file.Read(whitespace) {
				w += 1
			}

			if w % 8 != expectedOffset {
				panic(fmt.Sprintf("Read incorrect number of whitespace characters %d vs %d", w, expectedOffset))
			}
			if w == 0 {
				panic("Name was not null terminated in index")
			}

			// Undo the last read, which wasn't whitespace..
			file.Seek(-1, 1)
		*/

	} else {
		panic("TODO: I can't handle such long names yet")
	}
	return &IndexEntry{f, IndexPath(name)}, nil
}

// A Stage represents a git merge stage in the index.
type Stage uint8

// Valid merge stages.
const (
	Stage0 = Stage(iota)
	Stage1
	Stage2
	Stage3
)

// Adds an entry to the index with Sha1 s and stage stage during a merge.
// If an entry already exists for this pathname/stage, it will be overwritten,
// otherwise it will be added if createEntry is true, and return an error if not.
//
// As a special case, if something is added as Stage0, then Stage1-3 entries
// will be removed.
func (g *Index) AddStage(c *Client, path IndexPath, s Sha1, stage Stage, mtime, mtimenano, size uint32, createEntry bool) error {
	if stage == Stage0 {
		defer g.RemoveUnmergedStages(c, path)
	}

	// Update the existing stage, if it exists.
	for _, entry := range g.Objects {
		if entry.PathName == path && entry.Stage() == stage {
			entry.Sha1 = s
			entry.Mtime = mtime
			entry.Mtimenano = mtimenano
			entry.Fsize = size

			// We found and updated the entry, no need to continue
			return nil
		}
	}

	if !createEntry {
		return fmt.Errorf("%v not found in index", path)
	}
	// There was no path/stage combo already in the index. Add it.

	// According to the git documentation:
	// Flags is
	//    A 16-bit 'flags' field split into (high to low bits)
	//
	//       1-bit assume-valid flag
	//
	//       1-bit extended flag (must be zero in version 2)
	//
	//       2-bit stage (during merge)
	//       12-bit name length if the length is less than 0xFFF; otherwise 0xFFF
	//     is stored in this field.

	// So we'll construct the flags based on what we know.

	var flags = uint16(stage) << 12 // start with the stage.
	// Add the name length.
	if len(path) >= 0x0FFF {
		flags |= 0x0FFF
	} else {
		flags |= (uint16(len(path)) & 0x0FFF)
	}

	g.Objects = append(g.Objects, &IndexEntry{
		FixedIndexEntry{
			0, //uint32(csec),
			0, //uint32(cnano),
			mtime,
			mtimenano,
			0,        //uint32(stat.Dev),
			0,        //uint32(stat.Ino),
			ModeBlob, // Directories are never added, only their files, so assume blob
			0,        //stat.Uid,
			0,        //stat.Gid,
			size,
			s,
			flags,
		},
		path,
	})
	g.NumberIndexEntries += 1
	sort.Sort(ByPath(g.Objects))
	return nil

}

// Remove any unmerged (non-stage 0) stage from the index for the given path
func (g *Index) RemoveUnmergedStages(c *Client, path IndexPath) error {
	// There are likely 3 things being deleted, so make a new slice
	newobjects := make([]*IndexEntry, 0, len(g.Objects))
	for _, entry := range g.Objects {
		stage := entry.Stage()
		if entry.PathName == path && stage == Stage0 {
			newobjects = append(newobjects, entry)
		} else if entry.PathName == path && stage != Stage0 {
			// do not add it, it's the wrong stage.
		} else {
			// It's a different Pathname, keep it.
			newobjects = append(newobjects, entry)
		}
	}
	g.Objects = newobjects
	g.NumberIndexEntries = uint32(len(newobjects))
	return nil
}

// Adds a file to the index, without writing it to disk.
// To write it to disk after calling this, use GitIndex.WriteIndex
//
// This will do the following:
// 1. write git object blob of file contents to .git/objects
// 2. normalize os.File name to path relative to gitRoot
// 3. search GitIndex for normalized name
//	if GitIndexEntry found
//		update GitIndexEntry to point to the new object blob
// 	else
// 		add new GitIndexEntry if not found and createEntry is true, error otherwise
//
func (g *Index) AddFile(c *Client, file *os.File, createEntry bool) error {
	contents, err := ioutil.ReadAll(file)
	if err != nil {
		return err
	}
	hash, err := c.WriteObject("blob", contents)
	if err != nil && err != ObjectExists {
		fmt.Fprintf(os.Stderr, "Error storing object: %s", err)
		return err
	}
	name, err := File(file.Name()).IndexPath(c)
	if err != nil {
		return err
	}

	fstat, err := file.Stat()
	if err != nil {
		return err
	}
	if fstat.IsDir() {
		// This should really recursively call add for each file in the directory.
		return fmt.Errorf("Add can't handle directories. yet.")
	}

	modTime := fstat.ModTime()
	return g.AddStage(
		c,
		name,
		hash,
		Stage0,
		uint32(modTime.Unix()),
		uint32(modTime.Nanosecond()),
		uint32(fstat.Size()),
		createEntry,
	)
}

type IndexStageEntry struct {
	IndexPath
	Stage
}

func (i *Index) GetStageMap() map[IndexStageEntry]*IndexEntry {
	r := make(map[IndexStageEntry]*IndexEntry)
	for _, entry := range i.Objects {
		r[IndexStageEntry{entry.PathName, entry.Stage()}] = entry
	}
	return r
}

type UnmergedPath struct {
	Stage1, Stage2, Stage3 *IndexEntry
}

func (i *Index) GetUnmerged() map[IndexPath]*UnmergedPath {
	r := make(map[IndexPath]*UnmergedPath)
	for _, entry := range i.Objects {
		if entry.Stage() != Stage0 {
			e, ok := r[entry.PathName]
			if !ok {
				e = &UnmergedPath{}
				r[entry.PathName] = e
			}
			switch entry.Stage() {
			case Stage1:
				e.Stage1 = entry
			case Stage2:
				e.Stage2 = entry
			case Stage3:
				e.Stage3 = entry
			}
		}
	}
	return r
}
func (i *Index) GetMap() map[IndexPath]*IndexEntry {
	r := make(map[IndexPath]*IndexEntry)
	for _, entry := range i.Objects {
		r[entry.PathName] = entry
	}
	return r
}

// Remove the first instance of file from the index. (This will usually
// be stage 0.)
func (g *Index) RemoveFile(file IndexPath) {
	for i, entry := range g.Objects {
		if entry.PathName == file {
			//println("Should remove ", i)
			g.Objects = append(g.Objects[:i], g.Objects[i+1:]...)
			g.NumberIndexEntries -= 1
			return
		}
	}
}

// This will write a new index file to w by doing the following:
// 1. Sort the objects in g.Index to ascending order based on name
// 2. Write g.fixedGitIndex to w
// 3. for each entry in g.Objects, write it to w.
// 4. Write the Sha1 of the contents of what was written
func (g Index) WriteIndex(file io.Writer) error {
	sort.Sort(ByPath(g.Objects))
	s := sha1.New()
	w := io.MultiWriter(file, s)
	binary.Write(w, binary.BigEndian, g.fixedGitIndex)
	for _, entry := range g.Objects {
		binary.Write(w, binary.BigEndian, entry.FixedIndexEntry)
		binary.Write(w, binary.BigEndian, []byte(entry.PathName))
		padding := 8 - ((82 + len(entry.PathName) + 4) % 8)
		p := make([]byte, padding)
		binary.Write(w, binary.BigEndian, p)
	}
	binary.Write(w, binary.BigEndian, s.Sum(nil))
	return nil
}

// Looks up the Sha1 of path currently stored in the index.
// Will return the 0 Sha1 if not found.
func (g Index) GetSha1(path IndexPath) Sha1 {
	for _, entry := range g.Objects {
		if entry.PathName == path {
			return entry.Sha1
		}
	}
	return Sha1{}
}

// Implement the sort interface on *GitIndexEntry, so that
// it's easy to sort by name.
type ByPath []*IndexEntry

func (g ByPath) Len() int      { return len(g) }
func (g ByPath) Swap(i, j int) { g[i], g[j] = g[j], g[i] }
func (g ByPath) Less(i, j int) bool {
	if g[i].PathName < g[j].PathName {
		return true
	} else if g[i].PathName == g[j].PathName {
		return g[i].Stage() < g[j].Stage()
	} else {
		return false
	}
}

// Replaces the index of Client with the the tree from the provided Treeish.
// if PreserveStatInfo is true, the stat information in the index won't be
// modified for existing entries.
func (g *Index) ResetIndex(c *Client, tree Treeish) error {
	newEntries, err := ExpandGitTreeIntoIndexes(c, tree, true, false)
	if err != nil {
		return err
	}
	g.NumberIndexEntries = uint32(len(newEntries))
	g.Objects = newEntries
	return nil
}

func (g Index) String() string {
	ret := ""

	for _, i := range g.Objects {
		ret += fmt.Sprintf("%v %v %v\n", i.Mode, i.Sha1, i.PathName)
	}
	return ret
}
<|endoftext|>


<commit_before>// +build !windows

/*
Copyright 2017 The Kubernetes Authors.

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

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

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

package initsystem

import (
	"fmt"
	"os/exec"
	"strings"
)

// OpenRCInitSystem defines openrc
type OpenRCInitSystem struct{}

// ServiceStart tries to start a specific service
func (openrc OpenRCInitSystem) ServiceStart(service string) error {
	args := []string{service, "start"}
	return exec.Command("rc-service", args...).Run()
}

// ServiceStop tries to stop a specific service
func (openrc OpenRCInitSystem) ServiceStop(service string) error {
	args := []string{service, "stop"}
	return exec.Command("rc-service", args...).Run()
}

// ServiceRestart tries to reload the environment and restart the specific service
func (openrc OpenRCInitSystem) ServiceRestart(service string) error {
	args := []string{service, "restart"}
	return exec.Command("rc-service", args...).Run()
}

// ServiceExists ensures the service is defined for this init system.
// openrc writes to stderr if a service is not found or not enabled
// this is in contrast to systemd which only writes to stdout.
// Hence, we use the Combinedoutput, and ignore the error.
func (openrc OpenRCInitSystem) ServiceExists(service string) bool {
	args := []string{service, "status"}
	outBytes, _ := exec.Command("rc-service", args...).CombinedOutput()
	return !strings.Contains(string(outBytes), "does not exist")
}

// ServiceIsEnabled ensures the service is enabled to start on each boot.
func (openrc OpenRCInitSystem) ServiceIsEnabled(service string) bool {
	args := []string{"show", "default"}
	outBytes, _ := exec.Command("rc-update", args...).Output()
	return strings.Contains(string(outBytes), service)
}

// ServiceIsActive ensures the service is running, or attempting to run. (crash looping in the case of kubelet)
func (openrc OpenRCInitSystem) ServiceIsActive(service string) bool {
	args := []string{service, "status"}
	outBytes, _ := exec.Command("rc-service", args...).Output()
	return !strings.Contains(string(outBytes), "stopped")
}

// EnableCommand return a string describing how to enable a service
func (openrc OpenRCInitSystem) EnableCommand(service string) string {
	return fmt.Sprintf("rc-update add %s default", service)
}

// SystemdInitSystem defines systemd
type SystemdInitSystem struct{}

// EnableCommand return a string describing how to enable a service
func (sysd SystemdInitSystem) EnableCommand(service string) string {
	return fmt.Sprintf("systemctl enable %s.service", service)
}

// reloadSystemd reloeads the systemd daemon
func (sysd SystemdInitSystem) reloadSystemd() error {
	if err := exec.Command("systemctl", "daemon-reload").Run(); err != nil {
		return fmt.Errorf("failed to reload systemd: %v", err)
	}
	return nil
}

// ServiceStart tries to start a specific service
func (sysd SystemdInitSystem) ServiceStart(service string) error {
	// Before we try to start any service, make sure that systemd is ready
	if err := sysd.reloadSystemd(); err != nil {
		return err
	}
	args := []string{"start", service}
	return exec.Command("systemctl", args...).Run()
}

// ServiceRestart tries to reload the environment and restart the specific service
func (sysd SystemdInitSystem) ServiceRestart(service string) error {
	// Before we try to restart any service, make sure that systemd is ready
	if err := sysd.reloadSystemd(); err != nil {
		return err
	}
	args := []string{"restart", service}
	return exec.Command("systemctl", args...).Run()
}

// ServiceStop tries to stop a specific service
func (sysd SystemdInitSystem) ServiceStop(service string) error {
	args := []string{"stop", service}
	return exec.Command("systemctl", args...).Run()
}

// ServiceExists ensures the service is defined for this init system.
func (sysd SystemdInitSystem) ServiceExists(service string) bool {
	args := []string{"status", service}
	outBytes, _ := exec.Command("systemctl", args...).Output()
	output := string(outBytes)
	return !strings.Contains(output, "Loaded: not-found")
}

// ServiceIsEnabled ensures the service is enabled to start on each boot.
func (sysd SystemdInitSystem) ServiceIsEnabled(service string) bool {
	args := []string{"is-enabled", service}
	err := exec.Command("systemctl", args...).Run()
	return err == nil
}

// ServiceIsActive will check is the service is "active". In the case of
// crash looping services (kubelet in our case) status will return as
// "activating", so we will consider this active as well.
func (sysd SystemdInitSystem) ServiceIsActive(service string) bool {
	args := []string{"is-active", service}
	// Ignoring error here, command returns non-0 if in "activating" status:
	outBytes, _ := exec.Command("systemctl", args...).Output()
	output := strings.TrimSpace(string(outBytes))
	if output == "active" || output == "activating" {
		return true
	}
	return false
}

// GetInitSystem returns an InitSystem for the current system, or nil
// if we cannot detect a supported init system.
// This indicates we will skip init system checks, not an error.
func GetInitSystem() (InitSystem, error) {
	// Assume existence of systemctl in path implies this is a systemd system:
	_, err := exec.LookPath("systemctl")
	if err == nil {
		return &SystemdInitSystem{}, nil
	}
	_, err = exec.LookPath("openrc")
	if err == nil {
		return &OpenRCInitSystem{}, nil
	}

	return nil, fmt.Errorf("no supported init system detected, skipping checking for services")
}
<commit_msg>On OpenRC ServiceIsActive should not report true if no such service exists<commit_after>// +build !windows

/*
Copyright 2017 The Kubernetes Authors.

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

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

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

package initsystem

import (
	"fmt"
	"os/exec"
	"strings"
)

// OpenRCInitSystem defines openrc
type OpenRCInitSystem struct{}

// ServiceStart tries to start a specific service
func (openrc OpenRCInitSystem) ServiceStart(service string) error {
	args := []string{service, "start"}
	return exec.Command("rc-service", args...).Run()
}

// ServiceStop tries to stop a specific service
func (openrc OpenRCInitSystem) ServiceStop(service string) error {
	args := []string{service, "stop"}
	return exec.Command("rc-service", args...).Run()
}

// ServiceRestart tries to reload the environment and restart the specific service
func (openrc OpenRCInitSystem) ServiceRestart(service string) error {
	args := []string{service, "restart"}
	return exec.Command("rc-service", args...).Run()
}

// ServiceExists ensures the service is defined for this init system.
// openrc writes to stderr if a service is not found or not enabled
// this is in contrast to systemd which only writes to stdout.
// Hence, we use the Combinedoutput, and ignore the error.
func (openrc OpenRCInitSystem) ServiceExists(service string) bool {
	args := []string{service, "status"}
	outBytes, _ := exec.Command("rc-service", args...).CombinedOutput()
	return !strings.Contains(string(outBytes), "does not exist")
}

// ServiceIsEnabled ensures the service is enabled to start on each boot.
func (openrc OpenRCInitSystem) ServiceIsEnabled(service string) bool {
	args := []string{"show", "default"}
	outBytes, _ := exec.Command("rc-update", args...).Output()
	return strings.Contains(string(outBytes), service)
}

// ServiceIsActive ensures the service is running, or attempting to run. (crash looping in the case of kubelet)
func (openrc OpenRCInitSystem) ServiceIsActive(service string) bool {
	args := []string{service, "status"}
	outBytes, _ := exec.Command("rc-service", args...).CombinedOutput()
	outStr := string(outBytes)
	return !strings.Contains(outStr, "stopped") && !strings.Contains(outStr, "does not exist")
}

// EnableCommand return a string describing how to enable a service
func (openrc OpenRCInitSystem) EnableCommand(service string) string {
	return fmt.Sprintf("rc-update add %s default", service)
}

// SystemdInitSystem defines systemd
type SystemdInitSystem struct{}

// EnableCommand return a string describing how to enable a service
func (sysd SystemdInitSystem) EnableCommand(service string) string {
	return fmt.Sprintf("systemctl enable %s.service", service)
}

// reloadSystemd reloeads the systemd daemon
func (sysd SystemdInitSystem) reloadSystemd() error {
	if err := exec.Command("systemctl", "daemon-reload").Run(); err != nil {
		return fmt.Errorf("failed to reload systemd: %v", err)
	}
	return nil
}

// ServiceStart tries to start a specific service
func (sysd SystemdInitSystem) ServiceStart(service string) error {
	// Before we try to start any service, make sure that systemd is ready
	if err := sysd.reloadSystemd(); err != nil {
		return err
	}
	args := []string{"start", service}
	return exec.Command("systemctl", args...).Run()
}

// ServiceRestart tries to reload the environment and restart the specific service
func (sysd SystemdInitSystem) ServiceRestart(service string) error {
	// Before we try to restart any service, make sure that systemd is ready
	if err := sysd.reloadSystemd(); err != nil {
		return err
	}
	args := []string{"restart", service}
	return exec.Command("systemctl", args...).Run()
}

// ServiceStop tries to stop a specific service
func (sysd SystemdInitSystem) ServiceStop(service string) error {
	args := []string{"stop", service}
	return exec.Command("systemctl", args...).Run()
}

// ServiceExists ensures the service is defined for this init system.
func (sysd SystemdInitSystem) ServiceExists(service string) bool {
	args := []string{"status", service}
	outBytes, _ := exec.Command("systemctl", args...).Output()
	output := string(outBytes)
	return !strings.Contains(output, "Loaded: not-found")
}

// ServiceIsEnabled ensures the service is enabled to start on each boot.
func (sysd SystemdInitSystem) ServiceIsEnabled(service string) bool {
	args := []string{"is-enabled", service}
	err := exec.Command("systemctl", args...).Run()
	return err == nil
}

// ServiceIsActive will check is the service is "active". In the case of
// crash looping services (kubelet in our case) status will return as
// "activating", so we will consider this active as well.
func (sysd SystemdInitSystem) ServiceIsActive(service string) bool {
	args := []string{"is-active", service}
	// Ignoring error here, command returns non-0 if in "activating" status:
	outBytes, _ := exec.Command("systemctl", args...).Output()
	output := strings.TrimSpace(string(outBytes))
	if output == "active" || output == "activating" {
		return true
	}
	return false
}

// GetInitSystem returns an InitSystem for the current system, or nil
// if we cannot detect a supported init system.
// This indicates we will skip init system checks, not an error.
func GetInitSystem() (InitSystem, error) {
	// Assume existence of systemctl in path implies this is a systemd system:
	_, err := exec.LookPath("systemctl")
	if err == nil {
		return &SystemdInitSystem{}, nil
	}
	_, err = exec.LookPath("openrc")
	if err == nil {
		return &OpenRCInitSystem{}, nil
	}

	return nil, fmt.Errorf("no supported init system detected, skipping checking for services")
}
<|endoftext|>


<commit_before>/*
Copyright IBM Corp. All Rights Reserved.

SPDX-License-Identifier: Apache-2.0
*/

package leveldbhelper

import (
	"fmt"
	"sync"

	"github.com/hyperledger/fabric/common/flogging"
	"github.com/hyperledger/fabric/common/ledger/util"
	"github.com/pkg/errors"
	"github.com/syndtr/goleveldb/leveldb"
	"github.com/syndtr/goleveldb/leveldb/iterator"
	"github.com/syndtr/goleveldb/leveldb/opt"
	goleveldbutil "github.com/syndtr/goleveldb/leveldb/util"
)

var logger = flogging.MustGetLogger("leveldbhelper")

type dbState int32

const (
	closed dbState = iota
	opened
)

// Conf configuration for `DB`
type Conf struct {
	DBPath string
}

// DB - a wrapper on an actual store
type DB struct {
	conf    *Conf
	db      *leveldb.DB
	dbState dbState
	mux     sync.RWMutex

	readOpts        *opt.ReadOptions
	writeOptsNoSync *opt.WriteOptions
	writeOptsSync   *opt.WriteOptions
}

// CreateDB constructs a `DB`
func CreateDB(conf *Conf) *DB {
	readOpts := &opt.ReadOptions{}
	writeOptsNoSync := &opt.WriteOptions{}
	writeOptsSync := &opt.WriteOptions{}
	writeOptsSync.Sync = true

	return &DB{
		conf:            conf,
		dbState:         closed,
		readOpts:        readOpts,
		writeOptsNoSync: writeOptsNoSync,
		writeOptsSync:   writeOptsSync}
}

// Open opens the underlying db
func (dbInst *DB) Open() {
	dbInst.mux.Lock()
	defer dbInst.mux.Unlock()
	if dbInst.dbState == opened {
		return
	}
	dbOpts := &opt.Options{}
	dbPath := dbInst.conf.DBPath
	var err error
	var dirEmpty bool
	if dirEmpty, err = util.CreateDirIfMissing(dbPath); err != nil {
		panic(fmt.Sprintf("Error creating dir if missing: %s", err))
	}
	dbOpts.ErrorIfMissing = !dirEmpty
	if dbInst.db, err = leveldb.OpenFile(dbPath, dbOpts); err != nil {
		panic(fmt.Sprintf("Error opening leveldb: %s", err))
	}
	dbInst.dbState = opened
}

// Close closes the underlying db
func (dbInst *DB) Close() {
	dbInst.mux.Lock()
	defer dbInst.mux.Unlock()
	if dbInst.dbState == closed {
		return
	}
	if err := dbInst.db.Close(); err != nil {
		logger.Errorf("Error closing leveldb: %s", err)
	}
	dbInst.dbState = closed
}

// Get returns the value for the given key
func (dbInst *DB) Get(key []byte) ([]byte, error) {
	dbInst.mux.RLock()
	defer dbInst.mux.RUnlock()
	value, err := dbInst.db.Get(key, dbInst.readOpts)
	if err == leveldb.ErrNotFound {
		value = nil
		err = nil
	}
	if err != nil {
		logger.Errorf("Error retrieving leveldb key [%#v]: %s", key, err)
		return nil, errors.Wrapf(err, "error retrieving leveldb key [%#v]", key)
	}
	return value, nil
}

// Put saves the key/value
func (dbInst *DB) Put(key []byte, value []byte, sync bool) error {
	dbInst.mux.RLock()
	defer dbInst.mux.RUnlock()
	wo := dbInst.writeOptsNoSync
	if sync {
		wo = dbInst.writeOptsSync
	}
	err := dbInst.db.Put(key, value, wo)
	if err != nil {
		logger.Errorf("Error writing leveldb key [%#v]", key)
		return errors.Wrapf(err, "error writing leveldb key [%#v]", key)
	}
	return nil
}

// Delete deletes the given key
func (dbInst *DB) Delete(key []byte, sync bool) error {
	dbInst.mux.RLock()
	defer dbInst.mux.RUnlock()
	wo := dbInst.writeOptsNoSync
	if sync {
		wo = dbInst.writeOptsSync
	}
	err := dbInst.db.Delete(key, wo)
	if err != nil {
		logger.Errorf("Error deleting leveldb key [%#v]", key)
		return errors.Wrapf(err, "error deleting leveldb key [%#v]", key)
	}
	return nil
}

// GetIterator returns an iterator over key-value store. The iterator should be released after the use.
// The resultset contains all the keys that are present in the db between the startKey (inclusive) and the endKey (exclusive).
// A nil startKey represents the first available key and a nil endKey represent a logical key after the last available key
func (dbInst *DB) GetIterator(startKey []byte, endKey []byte) iterator.Iterator {
	dbInst.mux.RLock()
	defer dbInst.mux.RUnlock()
	return dbInst.db.NewIterator(&goleveldbutil.Range{Start: startKey, Limit: endKey}, dbInst.readOpts)
}

// WriteBatch writes a batch
func (dbInst *DB) WriteBatch(batch *leveldb.Batch, sync bool) error {
	dbInst.mux.RLock()
	defer dbInst.mux.RUnlock()
	wo := dbInst.writeOptsNoSync
	if sync {
		wo = dbInst.writeOptsSync
	}
	if err := dbInst.db.Write(batch, wo); err != nil {
		return errors.Wrap(err, "error writing batch to leveldb")
	}
	return nil
}
<commit_msg>code hygiene: rename mux to mutex<commit_after>/*
Copyright IBM Corp. All Rights Reserved.

SPDX-License-Identifier: Apache-2.0
*/

package leveldbhelper

import (
	"fmt"
	"sync"

	"github.com/hyperledger/fabric/common/flogging"
	"github.com/hyperledger/fabric/common/ledger/util"
	"github.com/pkg/errors"
	"github.com/syndtr/goleveldb/leveldb"
	"github.com/syndtr/goleveldb/leveldb/iterator"
	"github.com/syndtr/goleveldb/leveldb/opt"
	goleveldbutil "github.com/syndtr/goleveldb/leveldb/util"
)

var logger = flogging.MustGetLogger("leveldbhelper")

type dbState int32

const (
	closed dbState = iota
	opened
)

// Conf configuration for `DB`
type Conf struct {
	DBPath string
}

// DB - a wrapper on an actual store
type DB struct {
	conf    *Conf
	db      *leveldb.DB
	dbState dbState
	mutex   sync.RWMutex

	readOpts        *opt.ReadOptions
	writeOptsNoSync *opt.WriteOptions
	writeOptsSync   *opt.WriteOptions
}

// CreateDB constructs a `DB`
func CreateDB(conf *Conf) *DB {
	readOpts := &opt.ReadOptions{}
	writeOptsNoSync := &opt.WriteOptions{}
	writeOptsSync := &opt.WriteOptions{}
	writeOptsSync.Sync = true

	return &DB{
		conf:            conf,
		dbState:         closed,
		readOpts:        readOpts,
		writeOptsNoSync: writeOptsNoSync,
		writeOptsSync:   writeOptsSync}
}

// Open opens the underlying db
func (dbInst *DB) Open() {
	dbInst.mutex.Lock()
	defer dbInst.mutex.Unlock()
	if dbInst.dbState == opened {
		return
	}
	dbOpts := &opt.Options{}
	dbPath := dbInst.conf.DBPath
	var err error
	var dirEmpty bool
	if dirEmpty, err = util.CreateDirIfMissing(dbPath); err != nil {
		panic(fmt.Sprintf("Error creating dir if missing: %s", err))
	}
	dbOpts.ErrorIfMissing = !dirEmpty
	if dbInst.db, err = leveldb.OpenFile(dbPath, dbOpts); err != nil {
		panic(fmt.Sprintf("Error opening leveldb: %s", err))
	}
	dbInst.dbState = opened
}

// Close closes the underlying db
func (dbInst *DB) Close() {
	dbInst.mutex.Lock()
	defer dbInst.mutex.Unlock()
	if dbInst.dbState == closed {
		return
	}
	if err := dbInst.db.Close(); err != nil {
		logger.Errorf("Error closing leveldb: %s", err)
	}
	dbInst.dbState = closed
}

// Get returns the value for the given key
func (dbInst *DB) Get(key []byte) ([]byte, error) {
	dbInst.mutex.RLock()
	defer dbInst.mutex.RUnlock()
	value, err := dbInst.db.Get(key, dbInst.readOpts)
	if err == leveldb.ErrNotFound {
		value = nil
		err = nil
	}
	if err != nil {
		logger.Errorf("Error retrieving leveldb key [%#v]: %s", key, err)
		return nil, errors.Wrapf(err, "error retrieving leveldb key [%#v]", key)
	}
	return value, nil
}

// Put saves the key/value
func (dbInst *DB) Put(key []byte, value []byte, sync bool) error {
	dbInst.mutex.RLock()
	defer dbInst.mutex.RUnlock()
	wo := dbInst.writeOptsNoSync
	if sync {
		wo = dbInst.writeOptsSync
	}
	err := dbInst.db.Put(key, value, wo)
	if err != nil {
		logger.Errorf("Error writing leveldb key [%#v]", key)
		return errors.Wrapf(err, "error writing leveldb key [%#v]", key)
	}
	return nil
}

// Delete deletes the given key
func (dbInst *DB) Delete(key []byte, sync bool) error {
	dbInst.mutex.RLock()
	defer dbInst.mutex.RUnlock()
	wo := dbInst.writeOptsNoSync
	if sync {
		wo = dbInst.writeOptsSync
	}
	err := dbInst.db.Delete(key, wo)
	if err != nil {
		logger.Errorf("Error deleting leveldb key [%#v]", key)
		return errors.Wrapf(err, "error deleting leveldb key [%#v]", key)
	}
	return nil
}

// GetIterator returns an iterator over key-value store. The iterator should be released after the use.
// The resultset contains all the keys that are present in the db between the startKey (inclusive) and the endKey (exclusive).
// A nil startKey represents the first available key and a nil endKey represent a logical key after the last available key
func (dbInst *DB) GetIterator(startKey []byte, endKey []byte) iterator.Iterator {
	dbInst.mutex.RLock()
	defer dbInst.mutex.RUnlock()
	return dbInst.db.NewIterator(&goleveldbutil.Range{Start: startKey, Limit: endKey}, dbInst.readOpts)
}

// WriteBatch writes a batch
func (dbInst *DB) WriteBatch(batch *leveldb.Batch, sync bool) error {
	dbInst.mutex.RLock()
	defer dbInst.mutex.RUnlock()
	wo := dbInst.writeOptsNoSync
	if sync {
		wo = dbInst.writeOptsSync
	}
	if err := dbInst.db.Write(batch, wo); err != nil {
		return errors.Wrap(err, "error writing batch to leveldb")
	}
	return nil
}
<|endoftext|>


<commit_before>package restful

import (
	"net/http"
	"net/http/httptest"
	"testing"
)

const (
	pathGetFriends = "/get/{userId}/friends"
)

func TestParameter(t *testing.T) {
	p := &Parameter{&ParameterData{Name: "name", Description: "desc"}}
	p.AllowMultiple(true)
	p.DataType("int")
	p.Required(true)
	values := map[string]string{"a": "b"}
	p.AllowableValues(values)
	p.bePath()

	ws := new(WebService)
	ws.Param(p)
	if ws.pathParameters[0].Data().Name != "name" {
		t.Error("path parameter (or name) invalid")
	}
}
func TestWebService_CanCreateParameterKinds(t *testing.T) {
	ws := new(WebService)
	if ws.BodyParameter("b", "b").Kind() != BodyParameterKind {
		t.Error("body parameter expected")
	}
	if ws.PathParameter("p", "p").Kind() != PathParameterKind {
		t.Error("path parameter expected")
	}
	if ws.QueryParameter("q", "q").Kind() != QueryParameterKind {
		t.Error("query parameter expected")
	}
}

func TestCapturePanic(t *testing.T) {
	tearDown()
	Add(newPanicingService())
	httpRequest, _ := http.NewRequest("GET", "http://here.com/fire", nil)
	httpRequest.Header.Set("Accept", "*/*")
	httpWriter := httptest.NewRecorder()
	DefaultContainer.dispatch(httpWriter, httpRequest)
	if 500 != httpWriter.Code {
		t.Error("500 expected on fire")
	}
}

func TestNotFound(t *testing.T) {
	tearDown()
	httpRequest, _ := http.NewRequest("GET", "http://here.com/missing", nil)
	httpRequest.Header.Set("Accept", "*/*")
	httpWriter := httptest.NewRecorder()
	DefaultContainer.dispatch(httpWriter, httpRequest)
	if 404 != httpWriter.Code {
		t.Error("404 expected on missing")
	}
}

func TestMethodNotAllowed(t *testing.T) {
	tearDown()
	Add(newGetOnlyService())
	httpRequest, _ := http.NewRequest("POST", "http://here.com/get", nil)
	httpRequest.Header.Set("Accept", "*/*")
	httpWriter := httptest.NewRecorder()
	DefaultContainer.dispatch(httpWriter, httpRequest)
	if 405 != httpWriter.Code {
		t.Error("405 expected method not allowed")
	}
}

func TestSelectedRoutePath_Issue100(t *testing.T) {
	tearDown()
	Add(newSelectedRouteTestingService())
	httpRequest, _ := http.NewRequest("GET", "http://here.com/get/232452/friends", nil)
	httpRequest.Header.Set("Accept", "*/*")
	httpWriter := httptest.NewRecorder()
	DefaultContainer.dispatch(httpWriter, httpRequest)
	if http.StatusOK != httpWriter.Code {
		t.Error(http.StatusOK, "expected,", httpWriter.Code, "received.")
	}
}

func TestContentType415_Issue170(t *testing.T) {
	tearDown()
	Add(newGetOnlyJsonOnlyService())
	httpRequest, _ := http.NewRequest("GET", "http://here.com/get", nil)
	httpWriter := httptest.NewRecorder()
	DefaultContainer.dispatch(httpWriter, httpRequest)
	if 200 != httpWriter.Code {
		t.Errorf("Expected 200, got %d", httpWriter.Code)
	}
}

func TestContentType415_POST_Issue170(t *testing.T) {
	tearDown()
	Add(newPostOnlyJsonOnlyService())
	httpRequest, _ := http.NewRequest("POST", "http://here.com/post", nil)
	httpRequest.Header.Set("Content-Type", "application/json")
	httpWriter := httptest.NewRecorder()
	DefaultContainer.dispatch(httpWriter, httpRequest)
	if 200 != httpWriter.Code {
		t.Errorf("Expected 200, got %d", httpWriter.Code)
	}
}

// go test -v -test.run TestContentType406PlainJson ...restful
func TestContentType406PlainJson(t *testing.T) {
	tearDown()
	TraceLogger(testLogger{t})
	Add(newGetPlainTextOrJsonService())
	httpRequest, _ := http.NewRequest("GET", "http://here.com/get", nil)
	httpRequest.Header.Set("Accept", "text/plain")
	httpWriter := httptest.NewRecorder()
	DefaultContainer.dispatch(httpWriter, httpRequest)
	if got, want := httpWriter.Code, 200; got != want {
		t.Errorf("got %v, want %v", got, want)
	}
}

// go test -v -test.run TestContentTypeOctet_Issue170 ...restful
func TestContentTypeOctet_Issue170(t *testing.T) {
	tearDown()
	Add(newGetConsumingOctetStreamService())
	// with content-type
	httpRequest, _ := http.NewRequest("GET", "http://here.com/get", nil)
	httpRequest.Header.Set("Content-Type", MIME_OCTET)
	httpWriter := httptest.NewRecorder()
	DefaultContainer.dispatch(httpWriter, httpRequest)
	if 200 != httpWriter.Code {
		t.Errorf("Expected 200, got %d", httpWriter.Code)
	}
	// without content-type
	httpRequest, _ = http.NewRequest("GET", "http://here.com/get", nil)
	httpWriter = httptest.NewRecorder()
	DefaultContainer.dispatch(httpWriter, httpRequest)
	if 200 != httpWriter.Code {
		t.Errorf("Expected 200, got %d", httpWriter.Code)
	}
}

func newPanicingService() *WebService {
	ws := new(WebService).Path("")
	ws.Route(ws.GET("/fire").To(doPanic))
	return ws
}

func newGetOnlyService() *WebService {
	ws := new(WebService).Path("")
	ws.Route(ws.GET("/get").To(doPanic))
	return ws
}

func newPostOnlyJsonOnlyService() *WebService {
	ws := new(WebService).Path("")
	ws.Consumes("application/json")
	ws.Route(ws.POST("/post").To(doNothing))
	return ws
}

func newGetOnlyJsonOnlyService() *WebService {
	ws := new(WebService).Path("")
	ws.Consumes("application/json")
	ws.Route(ws.GET("/get").To(doNothing))
	return ws
}

func newGetPlainTextOrJsonService() *WebService {
	ws := new(WebService).Path("")
	ws.Produces("text/plain", "application/json")
	ws.Route(ws.GET("/get").To(doNothing))
	return ws
}

func newGetConsumingOctetStreamService() *WebService {
	ws := new(WebService).Path("")
	ws.Consumes("application/octet-stream")
	ws.Route(ws.GET("/get").To(doNothing))
	return ws
}

func newSelectedRouteTestingService() *WebService {
	ws := new(WebService).Path("")
	ws.Route(ws.GET(pathGetFriends).To(selectedRouteChecker))
	return ws
}

func selectedRouteChecker(req *Request, resp *Response) {
	if req.SelectedRoutePath() != pathGetFriends {
		resp.InternalServerError()
	}
}

func doPanic(req *Request, resp *Response) {
	println("lightning...")
	panic("fire")
}

func doNothing(req *Request, resp *Response) {
}
<commit_msg>Add TestCapturePanicWithEncoded test<commit_after>package restful

import (
	"net/http"
	"net/http/httptest"
	"testing"
)

const (
	pathGetFriends = "/get/{userId}/friends"
)

func TestParameter(t *testing.T) {
	p := &Parameter{&ParameterData{Name: "name", Description: "desc"}}
	p.AllowMultiple(true)
	p.DataType("int")
	p.Required(true)
	values := map[string]string{"a": "b"}
	p.AllowableValues(values)
	p.bePath()

	ws := new(WebService)
	ws.Param(p)
	if ws.pathParameters[0].Data().Name != "name" {
		t.Error("path parameter (or name) invalid")
	}
}
func TestWebService_CanCreateParameterKinds(t *testing.T) {
	ws := new(WebService)
	if ws.BodyParameter("b", "b").Kind() != BodyParameterKind {
		t.Error("body parameter expected")
	}
	if ws.PathParameter("p", "p").Kind() != PathParameterKind {
		t.Error("path parameter expected")
	}
	if ws.QueryParameter("q", "q").Kind() != QueryParameterKind {
		t.Error("query parameter expected")
	}
}

func TestCapturePanic(t *testing.T) {
	tearDown()
	Add(newPanicingService())
	httpRequest, _ := http.NewRequest("GET", "http://here.com/fire", nil)
	httpRequest.Header.Set("Accept", "*/*")
	httpWriter := httptest.NewRecorder()
	DefaultContainer.dispatch(httpWriter, httpRequest)
	if 500 != httpWriter.Code {
		t.Error("500 expected on fire")
	}
}

func TestCapturePanicWithEncoded(t *testing.T) {
	tearDown()
	Add(newPanicingService())
	DefaultContainer.EnableContentEncoding(true)
	httpRequest, _ := http.NewRequest("GET", "http://here.com/fire", nil)
	httpRequest.Header.Set("Accept", "*/*")
	httpRequest.Header.Set("Accept-Encoding", "gzip")
	httpWriter := httptest.NewRecorder()
	DefaultContainer.dispatch(httpWriter, httpRequest)
	if 500 != httpWriter.Code {
		t.Error("500 expected on fire, got", httpWriter.Code)
	}
}

func TestNotFound(t *testing.T) {
	tearDown()
	httpRequest, _ := http.NewRequest("GET", "http://here.com/missing", nil)
	httpRequest.Header.Set("Accept", "*/*")
	httpWriter := httptest.NewRecorder()
	DefaultContainer.dispatch(httpWriter, httpRequest)
	if 404 != httpWriter.Code {
		t.Error("404 expected on missing")
	}
}

func TestMethodNotAllowed(t *testing.T) {
	tearDown()
	Add(newGetOnlyService())
	httpRequest, _ := http.NewRequest("POST", "http://here.com/get", nil)
	httpRequest.Header.Set("Accept", "*/*")
	httpWriter := httptest.NewRecorder()
	DefaultContainer.dispatch(httpWriter, httpRequest)
	if 405 != httpWriter.Code {
		t.Error("405 expected method not allowed")
	}
}

func TestSelectedRoutePath_Issue100(t *testing.T) {
	tearDown()
	Add(newSelectedRouteTestingService())
	httpRequest, _ := http.NewRequest("GET", "http://here.com/get/232452/friends", nil)
	httpRequest.Header.Set("Accept", "*/*")
	httpWriter := httptest.NewRecorder()
	DefaultContainer.dispatch(httpWriter, httpRequest)
	if http.StatusOK != httpWriter.Code {
		t.Error(http.StatusOK, "expected,", httpWriter.Code, "received.")
	}
}

func TestContentType415_Issue170(t *testing.T) {
	tearDown()
	Add(newGetOnlyJsonOnlyService())
	httpRequest, _ := http.NewRequest("GET", "http://here.com/get", nil)
	httpWriter := httptest.NewRecorder()
	DefaultContainer.dispatch(httpWriter, httpRequest)
	if 200 != httpWriter.Code {
		t.Errorf("Expected 200, got %d", httpWriter.Code)
	}
}

func TestContentType415_POST_Issue170(t *testing.T) {
	tearDown()
	Add(newPostOnlyJsonOnlyService())
	httpRequest, _ := http.NewRequest("POST", "http://here.com/post", nil)
	httpRequest.Header.Set("Content-Type", "application/json")
	httpWriter := httptest.NewRecorder()
	DefaultContainer.dispatch(httpWriter, httpRequest)
	if 200 != httpWriter.Code {
		t.Errorf("Expected 200, got %d", httpWriter.Code)
	}
}

// go test -v -test.run TestContentType406PlainJson ...restful
func TestContentType406PlainJson(t *testing.T) {
	tearDown()
	TraceLogger(testLogger{t})
	Add(newGetPlainTextOrJsonService())
	httpRequest, _ := http.NewRequest("GET", "http://here.com/get", nil)
	httpRequest.Header.Set("Accept", "text/plain")
	httpWriter := httptest.NewRecorder()
	DefaultContainer.dispatch(httpWriter, httpRequest)
	if got, want := httpWriter.Code, 200; got != want {
		t.Errorf("got %v, want %v", got, want)
	}
}

// go test -v -test.run TestContentTypeOctet_Issue170 ...restful
func TestContentTypeOctet_Issue170(t *testing.T) {
	tearDown()
	Add(newGetConsumingOctetStreamService())
	// with content-type
	httpRequest, _ := http.NewRequest("GET", "http://here.com/get", nil)
	httpRequest.Header.Set("Content-Type", MIME_OCTET)
	httpWriter := httptest.NewRecorder()
	DefaultContainer.dispatch(httpWriter, httpRequest)
	if 200 != httpWriter.Code {
		t.Errorf("Expected 200, got %d", httpWriter.Code)
	}
	// without content-type
	httpRequest, _ = http.NewRequest("GET", "http://here.com/get", nil)
	httpWriter = httptest.NewRecorder()
	DefaultContainer.dispatch(httpWriter, httpRequest)
	if 200 != httpWriter.Code {
		t.Errorf("Expected 200, got %d", httpWriter.Code)
	}
}

func newPanicingService() *WebService {
	ws := new(WebService).Path("")
	ws.Route(ws.GET("/fire").To(doPanic))
	return ws
}

func newGetOnlyService() *WebService {
	ws := new(WebService).Path("")
	ws.Route(ws.GET("/get").To(doPanic))
	return ws
}

func newPostOnlyJsonOnlyService() *WebService {
	ws := new(WebService).Path("")
	ws.Consumes("application/json")
	ws.Route(ws.POST("/post").To(doNothing))
	return ws
}

func newGetOnlyJsonOnlyService() *WebService {
	ws := new(WebService).Path("")
	ws.Consumes("application/json")
	ws.Route(ws.GET("/get").To(doNothing))
	return ws
}

func newGetPlainTextOrJsonService() *WebService {
	ws := new(WebService).Path("")
	ws.Produces("text/plain", "application/json")
	ws.Route(ws.GET("/get").To(doNothing))
	return ws
}

func newGetConsumingOctetStreamService() *WebService {
	ws := new(WebService).Path("")
	ws.Consumes("application/octet-stream")
	ws.Route(ws.GET("/get").To(doNothing))
	return ws
}

func newSelectedRouteTestingService() *WebService {
	ws := new(WebService).Path("")
	ws.Route(ws.GET(pathGetFriends).To(selectedRouteChecker))
	return ws
}

func selectedRouteChecker(req *Request, resp *Response) {
	if req.SelectedRoutePath() != pathGetFriends {
		resp.InternalServerError()
	}
}

func doPanic(req *Request, resp *Response) {
	println("lightning...")
	panic("fire")
}

func doNothing(req *Request, resp *Response) {
}
<|endoftext|>


<commit_before>package gode

import (
	"io/ioutil"
	"os"
	"path/filepath"
	"runtime"
	"sort"
	"strings"
)

const registry = "https://cli-npm.heroku.com"

var rootPath string
var lockPath string
var nodePath string
var npmPath string

// SetRootPath sets the root for gode
func SetRootPath(root string) {
	rootPath = root
	base := filepath.Join(root, getLatestInstalledNode())
	nodePath = nodePathFromBase(base)
	herokuNodePath := os.Getenv("HEROKU_NODE_PATH")
	if herokuNodePath != "" {
		nodePath = herokuNodePath
	}
	npmPath = npmPathFromBase(base)
	lockPath = filepath.Join(rootPath, "node.lock")
}

func getLatestInstalledNode() string {
	nodes := getNodeInstalls()
	if len(nodes) == 0 {
		return ""
	}
	latest := nodes[len(nodes)-1]
	// ignore ancient versions
	if strings.HasPrefix(latest, "node-v0") {
		return ""
	}
	return latest
}

func getNodeInstalls() []string {
	nodes := []string{}
	files, _ := ioutil.ReadDir(rootPath)
	for _, f := range files {
		name := f.Name()
		if f.IsDir() && strings.HasPrefix(name, "node-v") {
			nodes = append(nodes, name)
		}
	}
	sort.Strings(nodes)
	return nodes
}

func nodePathFromBase(base string) string {
	path := filepath.Join(base, "bin", "node")
	if runtime.GOOS == "windows" {
		return path + ".exe"
	}
	return path
}

func npmPathFromBase(base string) string {
	return filepath.Join(base, "lib", "node_modules", "npm", "cli.js")
}

func modulesDir() string {
	return filepath.Join(rootPath, "node_modules")
}
<commit_msg>allow setting HEROKU_NPM_REGISTRY env var<commit_after>package gode

import (
	"io/ioutil"
	"os"
	"path/filepath"
	"runtime"
	"sort"
	"strings"
)

var registry string

var rootPath string
var lockPath string
var nodePath string
var npmPath string

func init() {
	registry = os.Getenv("HEROKU_NPM_REGISTRY")
	if registry == "" {
		registry = "https://cli-npm.heroku.com"
	}
}

// SetRootPath sets the root for gode
func SetRootPath(root string) {
	rootPath = root
	base := filepath.Join(root, getLatestInstalledNode())
	nodePath = nodePathFromBase(base)
	herokuNodePath := os.Getenv("HEROKU_NODE_PATH")
	if herokuNodePath != "" {
		nodePath = herokuNodePath
	}
	npmPath = npmPathFromBase(base)
	lockPath = filepath.Join(rootPath, "node.lock")
}

func getLatestInstalledNode() string {
	nodes := getNodeInstalls()
	if len(nodes) == 0 {
		return ""
	}
	latest := nodes[len(nodes)-1]
	// ignore ancient versions
	if strings.HasPrefix(latest, "node-v0") {
		return ""
	}
	return latest
}

func getNodeInstalls() []string {
	nodes := []string{}
	files, _ := ioutil.ReadDir(rootPath)
	for _, f := range files {
		name := f.Name()
		if f.IsDir() && strings.HasPrefix(name, "node-v") {
			nodes = append(nodes, name)
		}
	}
	sort.Strings(nodes)
	return nodes
}

func nodePathFromBase(base string) string {
	path := filepath.Join(base, "bin", "node")
	if runtime.GOOS == "windows" {
		return path + ".exe"
	}
	return path
}

func npmPathFromBase(base string) string {
	return filepath.Join(base, "lib", "node_modules", "npm", "cli.js")
}

func modulesDir() string {
	return filepath.Join(rootPath, "node_modules")
}
<|endoftext|>


<commit_before>package main

import (
	. "polydawn.net/pogo/gosh"
	"os"
)

func main() {
	echo := Sh("echo")("-n", "-e").
		BakeOpts(DefaultIO)

	echo("wat\n", "\t\033[0;31mred and indented\033[0m\n")()

	cat := Sh("cat")
	catIn := cat.BakeArgs("-")
	catIn(Opts{In:"piping in things is easy!\n", Out: os.Stdout})()

	shell := Sh("bash")("-c")

	shell("echo $TERM > testlag")()

	shell(ClearEnv{})("echo $TERM > testlag2")()

	shell(Env{"VAR": "59"})(Opts{OkExit:[]int{59}})("exit $VAR")()
}
<commit_msg>extending the examples of gosh usage<commit_after>package main

import (
	"os"
	. "polydawn.net/pogo/gosh"
)

func main() {
	echo := Sh("echo")("-n", "-e")(DefaultIO)
	echo("wat\n", "\t\033[0;31mred and indented\033[0m\n")()

	// basic commands
	Sh("head")("--bytes=20", "/dev/zero")()

	// making your own shorthand
	aptget := Sh("apt-get")("install")
	aptyes := aptget("-y")
	// now `aptyes("git")` will install git and the `-y` argument means it will proceed automatically

	// we'll do shorthand for a shell, so we can use the shell's echo for the rest of the demo.
	shell := Sh("bash")("-c")

	// output rules
	shell("echo 'default output is quiet'")()    // gosh defaults to NullIO.
	shell(DefaultIO)("echo 'such a message!'")() // *now* your termnial gets it!
	shell = shell(DefaultIO)                     // let's always use DefaultIO now :)

	// environment rules
	shell("echo path=$PATH")()                           // your default env is retained
	shell(ClearEnv{})("echo path=$PATH")()               // unless you don't want that, of course
	shell(Env{"VAR": "59"})("echo some_var=$SOME_VAR")() // you can easily set an env var for just this command
	shell = shell(Env{"VAR": "59"})                      // or make any env var part of the command template (try it with $PATH -- no more bashrc needed every again!)

	// directing input and output
	cat := Sh("cat")
	catIn := cat.BakeArgs("-")
	catIn(Opts{In: "piping in things is easy!\n", Out: os.Stdout})()

	// collecting output

	// channelling output

	// changing the working directory

	// failed command automatically gets you the heck outta here

	// accepting exit codes
	shell(Env{"VAR": "59"})(Opts{OkExit: []int{59}})("exit $VAR")()
}
<|endoftext|>


<commit_before>package gracegrpc

import (
	"bytes"
	"flag"
	"fmt"
	"log"
	"net"
	"os"
	"os/signal"
	"strconv"
	"sync"
	"syscall"

	"github.com/facebookgo/grace/gracenet"
	"google.golang.org/grpc"
)

var (
	verbose    = flag.Bool("gracelog", true, "Enable logging.")
	gracePid   = os.Getenv("GRACE_PID_FILE")
	didInherit = os.Getenv("LISTEN_FDS") != ""
	ppid       = os.Getppid()
)

type graceGrpc struct {
	server   *grpc.Server
	net      *gracenet.Net
	listener net.Listener
	errors   chan error
	pidfile  string
}

func (gr *graceGrpc) InitPidFile(pid_file string) {
	gr.pidfile = pid_file
}

func NewGraceGrpc(s *grpc.Server, net, addr string) *graceGrpc {
	gr := &graceGrpc{
		server: s,
		net:    &gracenet.Net{},

		//for  StartProcess error.
		errors:  make(chan error),
		pidfile: gracePid,
	}
	l, err := gr.net.Listen(net, addr)
	if err != nil {
		panic(err)
	}
	gr.listener = l
	return gr
}

func (gr *graceGrpc) serve() {
	go gr.server.Serve(gr.listener)
}

func (gr *graceGrpc) wait() {
	var wg sync.WaitGroup
	wg.Add(1)
	go gr.signalHandler(&wg)
	wg.Wait()
}

func (gr *graceGrpc) signalHandler(wg *sync.WaitGroup) {
	ch := make(chan os.Signal, 10)
	signal.Notify(ch, syscall.SIGINT, syscall.SIGTERM, syscall.SIGUSR2)
	for {
		sig := <-ch
		log.Println(sig, "signal has received")
		switch sig {
		case syscall.SIGINT, syscall.SIGTERM:
			signal.Stop(ch)
			gr.server.GracefulStop()
			defer wg.Done()
			return
		case syscall.SIGUSR2:
			if _, err := gr.net.StartProcess(); err != nil {
				gr.errors <- err
			}
		}
	}
}

func (gr *graceGrpc) doWritePid(pid int) (err error) {
	if gr.pidfile == "" {
		return nil
	}

	pf, err := os.Create(gr.pidfile)
	defer pf.Close()
	if err != nil {
		log.Println(err)
	}

	_, err = pf.WriteString(strconv.Itoa(pid))
	if err != nil {
		log.Println(err)
	}
	return err
}

func (gr *graceGrpc) Serve() error {

	if *verbose {
		if didInherit {
			if ppid == 1 {
				log.Printf("Listening on init activated %s\n", pprintAddr(gr.listener))
			} else {
				const msg = "Graceful handoff of %s with new pid %d replace old pid %d"
				log.Printf(msg, pprintAddr(gr.listener), os.Getpid(), ppid)
			}
		} else {
			const msg = "Serving %s with pid %d\n"
			log.Printf(msg, pprintAddr(gr.listener), os.Getpid())
		}
	}

	err := gr.doWritePid(os.Getpid())
	if err != nil {
		log.Println(err)
	}

	gr.serve()

	if didInherit && ppid != 1 {
		if err := syscall.Kill(ppid, syscall.SIGTERM); err != nil {
			return fmt.Errorf("failed to close parent: %s", err)
		}
	}

	waitdone := make(chan struct{})
	go func() {
		defer close(waitdone)
		gr.wait()
	}()

	select {
	case err := <-gr.errors:
		if err == nil {
			panic("unexpected nil error")
		}
		return err
	case <-waitdone:
		if *verbose {
			log.Printf("Exiting pid %d.", os.Getpid())
		}
		return nil
	}
}

func pprintAddr(l net.Listener) []byte {
	var out bytes.Buffer
	fmt.Fprint(&out, l.Addr())
	return out.Bytes()
}
<commit_msg>调整顺序 格式化输出<commit_after>package gracegrpc

import (
	"bytes"
	"flag"
	"fmt"
	"log"
	"net"
	"os"
	"os/signal"
	"strconv"
	"sync"
	"syscall"

	"github.com/facebookgo/grace/gracenet"
	"google.golang.org/grpc"
)

var (
	verbose    = flag.Bool("gracelog", true, "Enable logging.")
	gracePid   = os.Getenv("GRACE_PID_FILE")
	didInherit = os.Getenv("LISTEN_FDS") != ""
	ppid       = os.Getppid()
)

type graceGrpc struct {
	server   *grpc.Server
	net      *gracenet.Net
	listener net.Listener
	errors   chan error
	pidfile  string
}

func (gr *graceGrpc) InitPidFile(pid_file string) {
	gr.pidfile = pid_file
}

func NewGraceGrpc(s *grpc.Server, net, addr string) *graceGrpc {
	gr := &graceGrpc{
		server: s,
		net:    &gracenet.Net{},

		//for  StartProcess error.
		errors:  make(chan error),
		pidfile: gracePid,
	}
	l, err := gr.net.Listen(net, addr)
	if err != nil {
		panic(err)
	}
	gr.listener = l
	return gr
}

func (gr *graceGrpc) serve() {
	go gr.server.Serve(gr.listener)
}

func (gr *graceGrpc) wait() {
	var wg sync.WaitGroup
	wg.Add(1)
	go gr.signalHandler(&wg)
	wg.Wait()
}

func (gr *graceGrpc) signalHandler(wg *sync.WaitGroup) {
	ch := make(chan os.Signal, 10)
	signal.Notify(ch, syscall.SIGINT, syscall.SIGTERM, syscall.SIGUSR2)
	for {
		sig := <-ch
		log.Printf("signal: %s has received", sig)
		switch sig {
		case syscall.SIGINT, syscall.SIGTERM:
			defer wg.Done()
			signal.Stop(ch)
			gr.server.GracefulStop()
			return
		case syscall.SIGUSR2:
			if _, err := gr.net.StartProcess(); err != nil {
				gr.errors <- err
			}
		}
	}
}

func (gr *graceGrpc) doWritePid(pid int) (err error) {
	if gr.pidfile == "" {
		return nil
	}

	pf, err := os.Create(gr.pidfile)
	defer pf.Close()
	if err != nil {
		log.Println(err)
	}

	_, err = pf.WriteString(strconv.Itoa(pid))
	if err != nil {
		log.Println(err)
	}
	return err
}

func (gr *graceGrpc) Serve() error {

	if *verbose {
		if didInherit {
			if ppid == 1 {
				log.Printf("Listening on init activated %s\n", pprintAddr(gr.listener))
			} else {
				const msg = "Graceful handoff of %s with new pid %d replace old pid %d"
				log.Printf(msg, pprintAddr(gr.listener), os.Getpid(), ppid)
			}
		} else {
			const msg = "Serving %s with pid %d\n"
			log.Printf(msg, pprintAddr(gr.listener), os.Getpid())
		}
	}

	err := gr.doWritePid(os.Getpid())
	if err != nil {
		log.Println(err)
	}

	gr.serve()

	if didInherit && ppid != 1 {
		if err := syscall.Kill(ppid, syscall.SIGTERM); err != nil {
			return fmt.Errorf("failed to close parent: %s", err)
		}
	}

	waitdone := make(chan struct{})
	go func() {
		defer close(waitdone)
		gr.wait()
	}()

	select {
	case err := <-gr.errors:
		if err == nil {
			panic("unexpected nil error")
		}
		return err
	case <-waitdone:
		if *verbose {
			log.Printf("Exiting pid %d.", os.Getpid())
		}
		return nil
	}
}

func pprintAddr(l net.Listener) []byte {
	var out bytes.Buffer
	fmt.Fprint(&out, l.Addr())
	return out.Bytes()
}
<|endoftext|>


<commit_before>package libkb

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

	humanize "github.com/dustin/go-humanize"
	keybase1 "github.com/keybase/client/go/protocol/keybase1"
	"github.com/keybase/go-framed-msgpack-rpc/rpc"
	context "golang.org/x/net/context"
	"golang.org/x/sync/errgroup"
)

type NetworkStatsJSON struct {
	Local  []keybase1.InstrumentationStat `json:"remote"`
	Remote []keybase1.InstrumentationStat `json:"local"`
}

var internalHosts = map[string]struct{}{
	DevelServerURI:      {},
	StagingServerURI:    {},
	ProductionServerURI: {},
	ProductionSiteURI:   {},
}

func InstrumentationTagFromRequest(req *http.Request) string {
	if req.URL == nil {
		return ""
	}
	host := req.URL.Host
	path := req.URL.Path
	if _, ok := internalHosts[fmt.Sprintf("%s://%s", req.URL.Scheme, host)]; ok {
		host = ""
		path = strings.TrimPrefix(req.URL.Path, APIURIPathPrefix)
		path = strings.TrimPrefix(path, "/")
	}
	return fmt.Sprintf("%s %s%s", req.Method, host, path)
}

func AddRPCRecord(tag string, stat keybase1.InstrumentationStat, record rpc.InstrumentationRecord) keybase1.InstrumentationStat {
	if stat.NumCalls == 0 {
		stat.Ctime = keybase1.ToTime(time.Now())
	}
	if stat.Tag == "" {
		stat.Tag = tag
	}
	stat.Mtime = keybase1.ToTime(time.Now())
	stat.NumCalls++
	dur := keybase1.ToDurationMsec(record.Dur)
	stat.TotalDur += dur
	if dur > stat.MaxDur {
		stat.MaxDur = dur
	}
	if dur < stat.MinDur || stat.NumCalls == 1 {
		stat.MinDur = dur
	}

	stat.TotalSize += record.Size
	if record.Size > stat.MaxSize {
		stat.MaxSize = record.Size
	}
	if record.Size < stat.MinSize || stat.NumCalls == 1 {
		stat.MinSize = record.Size
	}

	stat.AvgDur = stat.TotalDur / keybase1.DurationMsec(stat.NumCalls)
	stat.AvgSize = stat.TotalSize / int64(stat.NumCalls)
	return stat
}

type DiskInstrumentationStorage struct {
	Contextified
	sync.Mutex
	src     keybase1.NetworkSource
	storage map[string]keybase1.InstrumentationStat

	eg     errgroup.Group
	stopCh chan struct{}
}

var _ rpc.NetworkInstrumenterStorage = (*DiskInstrumentationStorage)(nil)

func NewDiskInstrumentationStorage(g *GlobalContext, src keybase1.NetworkSource) *DiskInstrumentationStorage {
	return &DiskInstrumentationStorage{
		Contextified: NewContextified(g),
		src:          src,
		storage:      make(map[string]keybase1.InstrumentationStat),
	}
}

func (s *DiskInstrumentationStorage) Start(ctx context.Context) {
	defer s.G().CTraceTimed(ctx, "DiskInstrumentationStorage: Start", func() error { return nil })()
	s.Lock()
	defer s.Unlock()
	s.stopCh = make(chan struct{})
	s.eg.Go(func() error { return s.flushLoop(s.stopCh) })
}

func (s *DiskInstrumentationStorage) Stop(ctx context.Context) chan struct{} {
	defer s.G().CTraceTimed(ctx, "DiskInstrumentationStorage: Stop", func() error { return nil })()
	s.Lock()
	defer s.Unlock()
	ch := make(chan struct{})
	if s.stopCh != nil {
		close(s.stopCh)
		s.stopCh = nil
		go func() {
			if err := s.eg.Wait(); err != nil {
				s.G().Log.Debug("DiskInstrumentationStorage: flush: unable to wait for shutdown: %v", err)
			}
			close(ch)
		}()
	} else {
		close(ch)
	}
	return ch
}

func (s *DiskInstrumentationStorage) flushLoop(stopCh chan struct{}) error {
	ctx := context.Background()
	for {
		select {
		case <-stopCh:
			return s.Flush(ctx)
		case <-time.After(5 * time.Minute):
			if err := s.Flush(ctx); err != nil {
				s.G().Log.CDebugf(ctx, "DiskInstrumentationStorage: flushLoop: unable to flush: %v", err)
			}
		}
	}
}

func (s *DiskInstrumentationStorage) Flush(ctx context.Context) (err error) {
	s.Lock()
	storage := s.storage
	s.storage = make(map[string]keybase1.InstrumentationStat)
	s.Unlock()
	return s.flush(ctx, storage)
}

func (s *DiskInstrumentationStorage) flush(ctx context.Context, storage map[string]keybase1.InstrumentationStat) (err error) {
	defer s.G().CTraceTimed(ctx, "DiskInstrumentationStorage: flush", func() error { return err })()
	for tag, record := range storage {
		var existing keybase1.InstrumentationStat
		found, err := s.G().LocalDb.GetIntoMsgpack(&existing, s.dbKey(tag))
		if err != nil {
			return err
		}
		// Keep only window of the past month
		if found && time.Since(existing.Ctime.Time()) <= time.Hour*24*30 {
			record = existing.AppendStat(record)
		}
		if err := s.G().LocalDb.PutObjMsgpack(s.dbKey(tag), nil, record); err != nil {
			return err
		}
	}
	return nil
}

func (s *DiskInstrumentationStorage) keyPrefix() string {
	return fmt.Sprintf("src:%d", s.src)
}

func (s *DiskInstrumentationStorage) dbKey(tag string) DbKey {
	return DbKey{
		Typ: DBNetworkInstrumentation,
		Key: fmt.Sprintf("%s|%s", s.keyPrefix(), tag),
	}
}

func (s *DiskInstrumentationStorage) getAllKeysLocked() (keys []DbKey, err error) {
	prefix := DbKey{
		Typ: DBNetworkInstrumentation,
	}.ToBytes()
	dbKeys, err := s.G().LocalDb.KeysWithPrefixes(prefix)
	if err != nil {
		return nil, fmt.Errorf("could not get KeysWithPrefixes: %v", err)
	}
	keys = make([]DbKey, 0, len(dbKeys))
	for dbKey := range dbKeys {
		if dbKey.Typ == DBNetworkInstrumentation {
			keys = append(keys, dbKey)
		}
	}
	return keys, nil
}

func (s *DiskInstrumentationStorage) GetAll(ctx context.Context) (res []keybase1.InstrumentationStat, err error) {
	defer s.G().CTraceTimed(ctx, "DiskInstrumentationStorage: GetAll", func() error { return err })()
	s.Lock()
	defer s.Unlock()

	if err := s.flush(ctx, s.storage); err != nil {
		return nil, err
	}

	dbKeys, err := s.getAllKeysLocked()
	if err != nil {
		return nil, err
	}
	keyPrefix := s.keyPrefix()
	for _, dbKey := range dbKeys {
		// ensure key matches expected format
		keyParts := strings.Split(dbKey.Key, "|")
		if len(keyParts) < 2 || keyParts[0] != keyPrefix {
			continue
		}
		var record keybase1.InstrumentationStat
		ok, err := s.G().LocalDb.GetIntoMsgpack(&record, dbKey)
		if err != nil {
			return nil, err
		} else if !ok {
			continue
		}
		res = append(res, record)
	}
	return res, nil
}

func (s *DiskInstrumentationStorage) Stats(ctx context.Context) (res []keybase1.InstrumentationStat, err error) {
	defer s.G().CTraceTimed(ctx, "DiskInstrumentationStorage: Stats", func() error { return err })()
	return s.GetAll(ctx)
}

var tagLogBlacklist = map[string]struct{}{
	"Call gregor.1.incoming.ping": struct{}{},
}

func (s *DiskInstrumentationStorage) logRecord(ctx context.Context, tag string, record rpc.InstrumentationRecord) {
	if s.src == keybase1.NetworkSource_LOCAL {
		return
	}
	if _, ok := tagLogBlacklist[tag]; !ok {
		s.G().PerfLog.CDebugf(ctx, "%s %v %s", tag, record.Dur, humanize.Bytes(uint64(record.Size)))
		s.G().RuntimeStats.PushPerfEvent(keybase1.PerfEvent{
			EventType: keybase1.PerfEventType_NETWORK,
			Message:   tag,
			Ctime:     keybase1.ToTime(record.Ctime),
		})
	}
}

func (s *DiskInstrumentationStorage) Put(ctx context.Context, tag string, record rpc.InstrumentationRecord) error {
	s.Lock()
	defer s.Unlock()
	s.storage[tag] = AddRPCRecord(tag, s.storage[tag], record)
	s.logRecord(ctx, tag, record)
	return nil
}

func NetworkInstrumenterStorageFromSrc(g *GlobalContext, src keybase1.NetworkSource) rpc.NetworkInstrumenterStorage {
	switch src {
	case keybase1.NetworkSource_LOCAL:
		return g.LocalNetworkInstrumenterStorage
	case keybase1.NetworkSource_REMOTE:
		return g.RemoteNetworkInstrumenterStorage
	default:
		return rpc.NewDummyInstrumentationStorage()
	}
}
<commit_msg>fix instrumenter json (#22896)<commit_after>package libkb

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

	humanize "github.com/dustin/go-humanize"
	keybase1 "github.com/keybase/client/go/protocol/keybase1"
	"github.com/keybase/go-framed-msgpack-rpc/rpc"
	context "golang.org/x/net/context"
	"golang.org/x/sync/errgroup"
)

type NetworkStatsJSON struct {
	Local  []keybase1.InstrumentationStat `json:"local"`
	Remote []keybase1.InstrumentationStat `json:"remote"`
}

var internalHosts = map[string]struct{}{
	DevelServerURI:      {},
	StagingServerURI:    {},
	ProductionServerURI: {},
	ProductionSiteURI:   {},
}

func InstrumentationTagFromRequest(req *http.Request) string {
	if req.URL == nil {
		return ""
	}
	host := req.URL.Host
	path := req.URL.Path
	if _, ok := internalHosts[fmt.Sprintf("%s://%s", req.URL.Scheme, host)]; ok {
		host = ""
		path = strings.TrimPrefix(req.URL.Path, APIURIPathPrefix)
		path = strings.TrimPrefix(path, "/")
	}
	return fmt.Sprintf("%s %s%s", req.Method, host, path)
}

func AddRPCRecord(tag string, stat keybase1.InstrumentationStat, record rpc.InstrumentationRecord) keybase1.InstrumentationStat {
	if stat.NumCalls == 0 {
		stat.Ctime = keybase1.ToTime(time.Now())
	}
	if stat.Tag == "" {
		stat.Tag = tag
	}
	stat.Mtime = keybase1.ToTime(time.Now())
	stat.NumCalls++
	dur := keybase1.ToDurationMsec(record.Dur)
	stat.TotalDur += dur
	if dur > stat.MaxDur {
		stat.MaxDur = dur
	}
	if dur < stat.MinDur || stat.NumCalls == 1 {
		stat.MinDur = dur
	}

	stat.TotalSize += record.Size
	if record.Size > stat.MaxSize {
		stat.MaxSize = record.Size
	}
	if record.Size < stat.MinSize || stat.NumCalls == 1 {
		stat.MinSize = record.Size
	}

	stat.AvgDur = stat.TotalDur / keybase1.DurationMsec(stat.NumCalls)
	stat.AvgSize = stat.TotalSize / int64(stat.NumCalls)
	return stat
}

type DiskInstrumentationStorage struct {
	Contextified
	sync.Mutex
	src     keybase1.NetworkSource
	storage map[string]keybase1.InstrumentationStat

	eg     errgroup.Group
	stopCh chan struct{}
}

var _ rpc.NetworkInstrumenterStorage = (*DiskInstrumentationStorage)(nil)

func NewDiskInstrumentationStorage(g *GlobalContext, src keybase1.NetworkSource) *DiskInstrumentationStorage {
	return &DiskInstrumentationStorage{
		Contextified: NewContextified(g),
		src:          src,
		storage:      make(map[string]keybase1.InstrumentationStat),
	}
}

func (s *DiskInstrumentationStorage) Start(ctx context.Context) {
	defer s.G().CTraceTimed(ctx, "DiskInstrumentationStorage: Start", func() error { return nil })()
	s.Lock()
	defer s.Unlock()
	s.stopCh = make(chan struct{})
	s.eg.Go(func() error { return s.flushLoop(s.stopCh) })
}

func (s *DiskInstrumentationStorage) Stop(ctx context.Context) chan struct{} {
	defer s.G().CTraceTimed(ctx, "DiskInstrumentationStorage: Stop", func() error { return nil })()
	s.Lock()
	defer s.Unlock()
	ch := make(chan struct{})
	if s.stopCh != nil {
		close(s.stopCh)
		s.stopCh = nil
		go func() {
			if err := s.eg.Wait(); err != nil {
				s.G().Log.Debug("DiskInstrumentationStorage: flush: unable to wait for shutdown: %v", err)
			}
			close(ch)
		}()
	} else {
		close(ch)
	}
	return ch
}

func (s *DiskInstrumentationStorage) flushLoop(stopCh chan struct{}) error {
	ctx := context.Background()
	for {
		select {
		case <-stopCh:
			return s.Flush(ctx)
		case <-time.After(5 * time.Minute):
			if err := s.Flush(ctx); err != nil {
				s.G().Log.CDebugf(ctx, "DiskInstrumentationStorage: flushLoop: unable to flush: %v", err)
			}
		}
	}
}

func (s *DiskInstrumentationStorage) Flush(ctx context.Context) (err error) {
	s.Lock()
	storage := s.storage
	s.storage = make(map[string]keybase1.InstrumentationStat)
	s.Unlock()
	return s.flush(ctx, storage)
}

func (s *DiskInstrumentationStorage) flush(ctx context.Context, storage map[string]keybase1.InstrumentationStat) (err error) {
	defer s.G().CTraceTimed(ctx, "DiskInstrumentationStorage: flush", func() error { return err })()
	for tag, record := range storage {
		var existing keybase1.InstrumentationStat
		found, err := s.G().LocalDb.GetIntoMsgpack(&existing, s.dbKey(tag))
		if err != nil {
			return err
		}
		// Keep only window of the past month
		if found && time.Since(existing.Ctime.Time()) <= time.Hour*24*30 {
			record = existing.AppendStat(record)
		}
		if err := s.G().LocalDb.PutObjMsgpack(s.dbKey(tag), nil, record); err != nil {
			return err
		}
	}
	return nil
}

func (s *DiskInstrumentationStorage) keyPrefix() string {
	return fmt.Sprintf("src:%d", s.src)
}

func (s *DiskInstrumentationStorage) dbKey(tag string) DbKey {
	return DbKey{
		Typ: DBNetworkInstrumentation,
		Key: fmt.Sprintf("%s|%s", s.keyPrefix(), tag),
	}
}

func (s *DiskInstrumentationStorage) getAllKeysLocked() (keys []DbKey, err error) {
	prefix := DbKey{
		Typ: DBNetworkInstrumentation,
	}.ToBytes()
	dbKeys, err := s.G().LocalDb.KeysWithPrefixes(prefix)
	if err != nil {
		return nil, fmt.Errorf("could not get KeysWithPrefixes: %v", err)
	}
	keys = make([]DbKey, 0, len(dbKeys))
	for dbKey := range dbKeys {
		if dbKey.Typ == DBNetworkInstrumentation {
			keys = append(keys, dbKey)
		}
	}
	return keys, nil
}

func (s *DiskInstrumentationStorage) GetAll(ctx context.Context) (res []keybase1.InstrumentationStat, err error) {
	defer s.G().CTraceTimed(ctx, "DiskInstrumentationStorage: GetAll", func() error { return err })()
	s.Lock()
	defer s.Unlock()

	if err := s.flush(ctx, s.storage); err != nil {
		return nil, err
	}

	dbKeys, err := s.getAllKeysLocked()
	if err != nil {
		return nil, err
	}
	keyPrefix := s.keyPrefix()
	for _, dbKey := range dbKeys {
		// ensure key matches expected format
		keyParts := strings.Split(dbKey.Key, "|")
		if len(keyParts) < 2 || keyParts[0] != keyPrefix {
			continue
		}
		var record keybase1.InstrumentationStat
		ok, err := s.G().LocalDb.GetIntoMsgpack(&record, dbKey)
		if err != nil {
			return nil, err
		} else if !ok {
			continue
		}
		res = append(res, record)
	}
	return res, nil
}

func (s *DiskInstrumentationStorage) Stats(ctx context.Context) (res []keybase1.InstrumentationStat, err error) {
	defer s.G().CTraceTimed(ctx, "DiskInstrumentationStorage: Stats", func() error { return err })()
	return s.GetAll(ctx)
}

var tagLogBlacklist = map[string]struct{}{
	"Call gregor.1.incoming.ping": struct{}{},
}

func (s *DiskInstrumentationStorage) logRecord(ctx context.Context, tag string, record rpc.InstrumentationRecord) {
	if s.src == keybase1.NetworkSource_LOCAL {
		return
	}
	if _, ok := tagLogBlacklist[tag]; !ok {
		s.G().PerfLog.CDebugf(ctx, "%s %v %s", tag, record.Dur, humanize.Bytes(uint64(record.Size)))
		s.G().RuntimeStats.PushPerfEvent(keybase1.PerfEvent{
			EventType: keybase1.PerfEventType_NETWORK,
			Message:   tag,
			Ctime:     keybase1.ToTime(record.Ctime),
		})
	}
}

func (s *DiskInstrumentationStorage) Put(ctx context.Context, tag string, record rpc.InstrumentationRecord) error {
	s.Lock()
	defer s.Unlock()
	s.storage[tag] = AddRPCRecord(tag, s.storage[tag], record)
	s.logRecord(ctx, tag, record)
	return nil
}

func NetworkInstrumenterStorageFromSrc(g *GlobalContext, src keybase1.NetworkSource) rpc.NetworkInstrumenterStorage {
	switch src {
	case keybase1.NetworkSource_LOCAL:
		return g.LocalNetworkInstrumenterStorage
	case keybase1.NetworkSource_REMOTE:
		return g.RemoteNetworkInstrumenterStorage
	default:
		return rpc.NewDummyInstrumentationStorage()
	}
}
<|endoftext|>


<commit_before>// Copyright 2012, Google Inc. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package tabletserver

import (
	"fmt"
	"strconv"
	"strings"

	log "github.com/golang/glog"
	"github.com/youtube/vitess/go/sqltypes"
	"github.com/youtube/vitess/go/sync2"
	"github.com/youtube/vitess/go/vt/dbconnpool"
	"github.com/youtube/vitess/go/vt/schema"
)

type TableInfo struct {
	*schema.Table
	Cache *RowCache
	// stats updated by sqlquery.go
	hits, absent, misses, invalidations sync2.AtomicInt64
}

func NewTableInfo(conn dbconnpool.PoolConnection, tableName string, tableType string, createTime sqltypes.Value, comment string, cachePool *CachePool) (ti *TableInfo, err error) {
	ti, err = loadTableInfo(conn, tableName)
	if err != nil {
		return nil, err
	}
	ti.initRowCache(conn, tableType, createTime, comment, cachePool)
	return ti, nil
}

func loadTableInfo(conn dbconnpool.PoolConnection, tableName string) (ti *TableInfo, err error) {
	ti = &TableInfo{Table: schema.NewTable(tableName)}
	if err = ti.fetchColumns(conn); err != nil {
		return nil, err
	}
	if err = ti.fetchIndexes(conn); err != nil {
		return nil, err
	}
	return ti, nil
}

func (ti *TableInfo) fetchColumns(conn dbconnpool.PoolConnection) error {
	columns, err := conn.ExecuteFetch(fmt.Sprintf("describe %s", ti.Name), 10000, false)
	if err != nil {
		return err
	}
	for _, row := range columns.Rows {
		ti.AddColumn(row[0].String(), row[1].String(), row[4], row[5].String())
	}
	return nil
}

func (ti *TableInfo) SetPK(colnames []string) error {
	pkIndex := schema.NewIndex("PRIMARY")
	colnums := make([]int, len(colnames))
	for i, colname := range colnames {
		colnums[i] = ti.FindColumn(colname)
		if colnums[i] == -1 {
			return fmt.Errorf("column %s not found", colname)
		}
		pkIndex.AddColumn(colname, 1)
	}
	for _, col := range ti.Columns {
		pkIndex.DataColumns = append(pkIndex.DataColumns, col.Name)
	}
	if len(ti.Indexes) == 0 {
		ti.Indexes = make([]*schema.Index, 1)
	} else if ti.Indexes[0].Name != "PRIMARY" {
		ti.Indexes = append(ti.Indexes, nil)
		copy(ti.Indexes[1:], ti.Indexes[:len(ti.Indexes)-1])
	} // else we replace the currunt primary key
	ti.Indexes[0] = pkIndex
	ti.PKColumns = colnums
	return nil
}

func (ti *TableInfo) fetchIndexes(conn dbconnpool.PoolConnection) error {
	indexes, err := conn.ExecuteFetch(fmt.Sprintf("show index from %s", ti.Name), 10000, false)
	if err != nil {
		return err
	}
	var currentIndex *schema.Index
	currentName := ""
	for _, row := range indexes.Rows {
		indexName := row[2].String()
		if currentName != indexName {
			currentIndex = ti.AddIndex(indexName)
			currentName = indexName
		}
		var cardinality uint64
		if !row[6].IsNull() {
			cardinality, err = strconv.ParseUint(row[6].String(), 0, 64)
			if err != nil {
				log.Warningf("%s", err)
			}
		}
		currentIndex.AddColumn(row[4].String(), cardinality)
	}
	if len(ti.Indexes) == 0 {
		return nil
	}
	pkIndex := ti.Indexes[0]
	if pkIndex.Name != "PRIMARY" {
		return nil
	}
	ti.PKColumns = make([]int, len(pkIndex.Columns))
	for i, pkCol := range pkIndex.Columns {
		ti.PKColumns[i] = ti.FindColumn(pkCol)
	}
	// Primary key contains all table columns
	for _, col := range ti.Columns {
		pkIndex.DataColumns = append(pkIndex.DataColumns, col.Name)
	}
	// Secondary indices contain all primary key columns
	for i := 1; i < len(ti.Indexes); i++ {
		ti.Indexes[i].DataColumns = pkIndex.Columns
	}
	return nil
}

func (ti *TableInfo) initRowCache(conn dbconnpool.PoolConnection, tableType string, createTime sqltypes.Value, comment string, cachePool *CachePool) {
	if cachePool.IsClosed() {
		return
	}

	if strings.Contains(comment, "vtocc_nocache") {
		log.Infof("%s commented as vtocc_nocache. Will not be cached.", ti.Name)
		return
	}

	if tableType == "VIEW" {
		log.Infof("%s is a view. Will not be cached.", ti.Name)
		return
	}

	if ti.PKColumns == nil {
		log.Infof("Table %s has no primary key. Will not be cached.", ti.Name)
		return
	}
	for _, col := range ti.PKColumns {
		if ti.Columns[col].Category == schema.CAT_OTHER {
			log.Infof("Table %s pk has unsupported column types. Will not be cached.", ti.Name)
			return
		}
	}

	ti.CacheType = schema.CACHE_RW
	ti.Cache = NewRowCache(ti, cachePool)
}

func (ti *TableInfo) StatsJSON() string {
	if ti.Cache == nil {
		return fmt.Sprintf("null")
	}
	h, a, m, i := ti.Stats()
	return fmt.Sprintf("{\"Hits\": %v, \"Absent\": %v, \"Misses\": %v, \"Invalidations\": %v}", h, a, m, i)
}

func (ti *TableInfo) Stats() (hits, absent, misses, invalidations int64) {
	return ti.hits.Get(), ti.absent.Get(), ti.misses.Get(), ti.invalidations.Get()
}
<commit_msg>query service: load schema safety<commit_after>// Copyright 2012, Google Inc. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package tabletserver

import (
	"fmt"
	"strconv"
	"strings"

	log "github.com/golang/glog"
	"github.com/youtube/vitess/go/sqltypes"
	"github.com/youtube/vitess/go/sync2"
	"github.com/youtube/vitess/go/vt/dbconnpool"
	"github.com/youtube/vitess/go/vt/schema"
)

type TableInfo struct {
	*schema.Table
	Cache *RowCache
	// stats updated by sqlquery.go
	hits, absent, misses, invalidations sync2.AtomicInt64
}

func NewTableInfo(conn dbconnpool.PoolConnection, tableName string, tableType string, createTime sqltypes.Value, comment string, cachePool *CachePool) (ti *TableInfo, err error) {
	ti, err = loadTableInfo(conn, tableName)
	if err != nil {
		return nil, err
	}
	ti.initRowCache(conn, tableType, createTime, comment, cachePool)
	return ti, nil
}

func loadTableInfo(conn dbconnpool.PoolConnection, tableName string) (ti *TableInfo, err error) {
	ti = &TableInfo{Table: schema.NewTable(tableName)}
	if err = ti.fetchColumns(conn); err != nil {
		return nil, err
	}
	if err = ti.fetchIndexes(conn); err != nil {
		return nil, err
	}
	return ti, nil
}

func (ti *TableInfo) fetchColumns(conn dbconnpool.PoolConnection) error {
	columns, err := conn.ExecuteFetch(fmt.Sprintf("describe `%s`", ti.Name), 10000, false)
	if err != nil {
		return err
	}
	for _, row := range columns.Rows {
		ti.AddColumn(row[0].String(), row[1].String(), row[4], row[5].String())
	}
	return nil
}

func (ti *TableInfo) SetPK(colnames []string) error {
	pkIndex := schema.NewIndex("PRIMARY")
	colnums := make([]int, len(colnames))
	for i, colname := range colnames {
		colnums[i] = ti.FindColumn(colname)
		if colnums[i] == -1 {
			return fmt.Errorf("column %s not found", colname)
		}
		pkIndex.AddColumn(colname, 1)
	}
	for _, col := range ti.Columns {
		pkIndex.DataColumns = append(pkIndex.DataColumns, col.Name)
	}
	if len(ti.Indexes) == 0 {
		ti.Indexes = make([]*schema.Index, 1)
	} else if ti.Indexes[0].Name != "PRIMARY" {
		ti.Indexes = append(ti.Indexes, nil)
		copy(ti.Indexes[1:], ti.Indexes[:len(ti.Indexes)-1])
	} // else we replace the currunt primary key
	ti.Indexes[0] = pkIndex
	ti.PKColumns = colnums
	return nil
}

func (ti *TableInfo) fetchIndexes(conn dbconnpool.PoolConnection) error {
	indexes, err := conn.ExecuteFetch(fmt.Sprintf("show index from `%s`", ti.Name), 10000, false)
	if err != nil {
		return err
	}
	var currentIndex *schema.Index
	currentName := ""
	for _, row := range indexes.Rows {
		indexName := row[2].String()
		if currentName != indexName {
			currentIndex = ti.AddIndex(indexName)
			currentName = indexName
		}
		var cardinality uint64
		if !row[6].IsNull() {
			cardinality, err = strconv.ParseUint(row[6].String(), 0, 64)
			if err != nil {
				log.Warningf("%s", err)
			}
		}
		currentIndex.AddColumn(row[4].String(), cardinality)
	}
	if len(ti.Indexes) == 0 {
		return nil
	}
	pkIndex := ti.Indexes[0]
	if pkIndex.Name != "PRIMARY" {
		return nil
	}
	ti.PKColumns = make([]int, len(pkIndex.Columns))
	for i, pkCol := range pkIndex.Columns {
		ti.PKColumns[i] = ti.FindColumn(pkCol)
	}
	// Primary key contains all table columns
	for _, col := range ti.Columns {
		pkIndex.DataColumns = append(pkIndex.DataColumns, col.Name)
	}
	// Secondary indices contain all primary key columns
	for i := 1; i < len(ti.Indexes); i++ {
		ti.Indexes[i].DataColumns = pkIndex.Columns
	}
	return nil
}

func (ti *TableInfo) initRowCache(conn dbconnpool.PoolConnection, tableType string, createTime sqltypes.Value, comment string, cachePool *CachePool) {
	if cachePool.IsClosed() {
		return
	}

	if strings.Contains(comment, "vtocc_nocache") {
		log.Infof("%s commented as vtocc_nocache. Will not be cached.", ti.Name)
		return
	}

	if tableType == "VIEW" {
		log.Infof("%s is a view. Will not be cached.", ti.Name)
		return
	}

	if ti.PKColumns == nil {
		log.Infof("Table %s has no primary key. Will not be cached.", ti.Name)
		return
	}
	for _, col := range ti.PKColumns {
		if ti.Columns[col].Category == schema.CAT_OTHER {
			log.Infof("Table %s pk has unsupported column types. Will not be cached.", ti.Name)
			return
		}
	}

	ti.CacheType = schema.CACHE_RW
	ti.Cache = NewRowCache(ti, cachePool)
}

func (ti *TableInfo) StatsJSON() string {
	if ti.Cache == nil {
		return fmt.Sprintf("null")
	}
	h, a, m, i := ti.Stats()
	return fmt.Sprintf("{\"Hits\": %v, \"Absent\": %v, \"Misses\": %v, \"Invalidations\": %v}", h, a, m, i)
}

func (ti *TableInfo) Stats() (hits, absent, misses, invalidations int64) {
	return ti.hits.Get(), ti.absent.Get(), ti.misses.Get(), ti.invalidations.Get()
}
<|endoftext|>


<commit_before>// TO RUN:
//  go run ./examples/mqtt_driver_ping.go tcp://test.mosquitto.org:1883
// OR
//	go run ./examples/mqtt_driver_ping.go ssl://iot.eclipse.org:8883

package main

import (
	"fmt"
	"os"
	"time"

	"gobot.io/x/gobot"
	"gobot.io/x/gobot/platforms/mqtt"
)

func main() {
	mqttAdaptor := mqtt.NewAdaptor(os.Args[1], "pinger")
	mqttAdaptor.SetAutoReconnect(true)

	holaDriver := mqtt.NewDriver(mqttAdaptor, "hola")
	helloDriver := mqtt.NewDriver(mqttAdaptor, "hello")

	work := func() {
		helloDriver.On(mqtt.Data, func(data interface{}) {
			fmt.Println("hello")
		})

		holaDriver.On(mqtt.Data, func(data interface{}) {
			fmt.Println("hola")
		})

		data := []byte("o")
		gobot.Every(1*time.Second, func() {
			helloDriver.Publish(data)
		})

		gobot.Every(5*time.Second, func() {
			holaDriver.Publish(data)
		})
	}

	robot := gobot.NewRobot("mqttBot",
		[]gobot.Connection{mqttAdaptor},
		work,
	)

	robot.Start()
}
<commit_msg>docs: add header to example with execution instructions<commit_after>// TO RUN:
//  go run ./examples/mqtt_driver_ping.go <SERVER>
//
// EXAMPLE:
//	go run ./examples/mqtt_driver_ping.go ssl://iot.eclipse.org:8883
//
package main

import (
	"fmt"
	"os"
	"time"

	"gobot.io/x/gobot"
	"gobot.io/x/gobot/platforms/mqtt"
)

func main() {
	mqttAdaptor := mqtt.NewAdaptor(os.Args[1], "pinger")
	mqttAdaptor.SetAutoReconnect(true)

	holaDriver := mqtt.NewDriver(mqttAdaptor, "hola")
	helloDriver := mqtt.NewDriver(mqttAdaptor, "hello")

	work := func() {
		helloDriver.On(mqtt.Data, func(data interface{}) {
			fmt.Println("hello")
		})

		holaDriver.On(mqtt.Data, func(data interface{}) {
			fmt.Println("hola")
		})

		data := []byte("o")
		gobot.Every(1*time.Second, func() {
			helloDriver.Publish(data)
		})

		gobot.Every(5*time.Second, func() {
			holaDriver.Publish(data)
		})
	}

	robot := gobot.NewRobot("mqttBot",
		[]gobot.Connection{mqttAdaptor},
		[]gobot.Devices{helloDriver, holaDriver},
		work,
	)

	robot.Start()
}
<|endoftext|>