<commit_before>package decide

import (
	"fmt"
	"github.com/pkg/errors"
	exp "github.com/sjhitchner/go-decide/expression"
	"github.com/sjhitchner/go-decide/lexer"
	"github.com/sjhitchner/go-decide/parser"
	"io"
)

type Tree struct {
	root *Node
}

func NewTree(objects map[string][]string) (*Tree, error) {

	sorter := NewFrequencySorter()

	// Build up frequency table
	for _, list := range objects {
		for _, expString := range list {
			sorter.AddToFrequencies(expString)
		}
	}

	fmt.Println("FL", sorter.FrequencyList())

	var root *Node
	var err error
	for object, expressions := range objects {
		sorter.Sort(expressions)

		root, err = addNode(root, expressions, object)
		if err != nil {
			return nil, err
		}
	}

	return &Tree{root}, nil
}

func addNode(node *Node, expressions []string, object string) (*Node, error) {
	expression, err := NewExpression(expressions[0])
	if err != nil {
		return nil, err
	}

	if node == nil {
		node = NewNode(expression)
	}

	if node.Expression.String() == expression.String() {
		if len(expressions) > 1 {
			node.True, err = addNode(node.True, expressions[1:], object)
			return node, err
		}

		node.Payload = append((*node).Payload, object)
		return node, nil
	}

	node.False, err = addNode(node.False, expressions, object)
	return node, err
}

func (t Tree) Evaluate(ctx exp.Context) ([]string, error) {
	payloadMap := make(map[string]struct{})
	if err := t.root.Evaluate(ctx, payloadMap); err != nil {
		return nil, errors.Wrap(err, "Error evaluating tree")
	}

	list := make([]string, 0, len(payloadMap))
	for key := range payloadMap {
		list = append(list, key)
	}

	return list, nil
}

func (t Tree) String() string {
	return t.root.String()
}

func (t Tree) Graph(w io.Writer) error {
	return Graph(w, t.root)
}

func NewExpression(str string) (exp.Expression, error) {

	lex := lexer.NewLexer([]byte(str))
	p := parser.NewParser()
	expression, err := p.Parse(lex)
	if err != nil {
		return nil, errors.Wrap(err, "Unable to build AST")
	}

	return expression.(exp.Expression), nil
}
<commit_msg>removing fmt<commit_after>package decide

import (
	"github.com/pkg/errors"
	exp "github.com/sjhitchner/go-decide/expression"
	"github.com/sjhitchner/go-decide/lexer"
	"github.com/sjhitchner/go-decide/parser"
	"io"
)

type Tree struct {
	root *Node
}

func NewTree(objects map[string][]string) (*Tree, error) {

	sorter := NewFrequencySorter()

	// Build up frequency table
	for _, list := range objects {
		for _, expString := range list {
			sorter.AddToFrequencies(expString)
		}
	}

	var root *Node
	var err error
	for object, expressions := range objects {
		sorter.Sort(expressions)

		root, err = addNode(root, expressions, object)
		if err != nil {
			return nil, err
		}
	}

	return &Tree{root}, nil
}

func addNode(node *Node, expressions []string, object string) (*Node, error) {
	expression, err := NewExpression(expressions[0])
	if err != nil {
		return nil, err
	}

	if node == nil {
		node = NewNode(expression)
	}

	if node.Expression.String() == expression.String() {
		if len(expressions) > 1 {
			node.True, err = addNode(node.True, expressions[1:], object)
			return node, err
		}

		node.Payload = append((*node).Payload, object)
		return node, nil
	}

	node.False, err = addNode(node.False, expressions, object)
	return node, err
}

func (t Tree) Evaluate(ctx exp.Context) ([]string, error) {
	payloadMap := make(map[string]struct{})
	if err := t.root.Evaluate(ctx, payloadMap); err != nil {
		return nil, errors.Wrap(err, "Error evaluating tree")
	}

	list := make([]string, 0, len(payloadMap))
	for key := range payloadMap {
		list = append(list, key)
	}

	return list, nil
}

func (t Tree) String() string {
	return t.root.String()
}

func (t Tree) Graph(w io.Writer) error {
	return Graph(w, t.root)
}

func NewExpression(str string) (exp.Expression, error) {

	lex := lexer.NewLexer([]byte(str))
	p := parser.NewParser()
	expression, err := p.Parse(lex)
	if err != nil {
		return nil, errors.Wrap(err, "Unable to build AST")
	}

	return expression.(exp.Expression), nil
}
<|endoftext|>


<commit_before>package tga

import (
  "bytes"
  "encoding/binary"
  "errors"
  "image"
  "image/color"
  "io"
)

type rawHeader struct {
  IdLength      uint8
  PaletteType   uint8
  ImageType     uint8
  PaletteFirst  uint16
  PaletteLength uint16
  PaletteBPP    uint8
  OriginX       uint16
  PriginY       uint16
  Width         uint16
  Height        uint16
  BPP           uint8
  Flags         uint8
}

type rawFooter struct {
  ExtAreaOffset uint32
  DevDirOffset  uint32
  Signature     [18]byte // tgaSignature
}

type tga struct {
  r             *bytes.Reader
  raw           rawHeader
  isPaletted    bool
  hasAlpha      bool
  width         int
  height        int
  pixelSize     int
  palette       []byte
  paletteLength int
  ColorModel    color.Model
  outPixelSize  int
  tmp           [4]byte
  decode        func(tga *tga, out []byte) (err error)
  flip          func(tga *tga, out []byte)
}

const (
  flagOriginRight = uint8(1 << 4)
  flagOriginTop   = 1 << 5
  flagMask        = 7
)

const (
  imageTypePaletted   = uint8(1)
  imageTypeTrueColor  = 2
  imageTypeMonoChrome = 3
  imageTypeRLE        = 1 << 3
  imageTypeMask       = 3
)

const (
  tgaRawHeaderSize = 18
  tgaRawFooterSize = 26
  tgaSignature     = "TRUEVISION-XFILE.\x00"
)

func decode(r io.Reader) (outImage image.Image, err error) {
  var tga tga
  var data bytes.Buffer

  if _, err = data.ReadFrom(r); err != nil {
    return
  }

  tga.r = bytes.NewReader(data.Bytes())

  if err = tga.getHeader(); err != nil {
    return
  }

  // skip header
  if _, err = tga.r.Seek(int64(tgaRawHeaderSize+tga.raw.IdLength), 0); err != nil {
    return
  }

  if tga.isPaletted {
    // read palette
    entrySize := int((tga.raw.PaletteBPP + 1) >> 3)
    tga.paletteLength = int(tga.raw.PaletteLength - tga.raw.PaletteFirst)
    tga.palette = make([]byte, entrySize*tga.paletteLength)

    // skip to colormap
    if _, err = tga.r.Seek(int64(entrySize)*int64(tga.raw.PaletteFirst), 1); err != nil {
      return
    }

    if _, err = io.ReadFull(tga.r, tga.palette); err != nil {
      return
    }
  }

  var pixels []byte
  rect := image.Rect(0, 0, tga.width, tga.height)

  // choose a right color model
  switch tga.ColorModel {
  case color.NRGBAModel:
    im := image.NewNRGBA(rect)
    outImage = im
    pixels = im.Pix

  case color.RGBAModel:
    im := image.NewRGBA(rect)
    outImage = im
    pixels = im.Pix

  case color.GrayModel:
    im := image.NewGray(rect)
    outImage = im
    pixels = im.Pix
  }

  if err = tga.decode(&tga, pixels); err == nil {
    tga.flip(&tga, pixels)
  }

  return
}

func decodeConfig(r io.Reader) (cfg image.Config, err error) {
  var tga tga
  var data bytes.Buffer

  if _, err = data.ReadFrom(r); err != nil {
    return
  }

  tga.r = bytes.NewReader(data.Bytes())

  if err = tga.getHeader(); err == nil {
    cfg = image.Config{
      ColorModel: tga.ColorModel,
      Width:      tga.width,
      Height:     tga.height,
    }
  }

  return
}

func init() {
  image.RegisterFormat("tga", "", decode, decodeConfig)
}

// applyExtensions reads extensions section (if it exists) and parses attribute type.
func (tga *tga) applyExtensions() (err error) {
  var rawFooter rawFooter

  if _, err = tga.r.Seek(int64(-tgaRawFooterSize), 2); err != nil {
    return
  } else if err = binary.Read(tga.r, binary.LittleEndian, &rawFooter); err != nil {
    return
  } else if bytes.Equal(rawFooter.Signature[:], []byte(tgaSignature)) && rawFooter.ExtAreaOffset != 0 {
    offset := int64(rawFooter.ExtAreaOffset + 0x1ee)

    var n int64
    var t byte

    if n, err = tga.r.Seek(offset, 0); err != nil || n != offset {
      return
    } else if t, err = tga.r.ReadByte(); err != nil {
      return
    } else if t == 3 {
      // alpha
      tga.hasAlpha = true
      tga.ColorModel = color.NRGBAModel
    } else if t == 4 {
      // premultiplied alpha
      tga.hasAlpha = true
      tga.ColorModel = color.RGBAModel
    } else {
      // attribute is not an alpha channel value, ignore it
      tga.hasAlpha = false
    }
  }

  return
}

// decodeRaw decodes a raw (uncompressed) data.
func decodeRaw(tga *tga, out []byte) (err error) {
  for i := 0; i < len(out) && err == nil; i += tga.outPixelSize {
    err = tga.getPixel(out[i:])
  }

  return
}

// decodeRLE decodes run-length encoded data.
func decodeRLE(tga *tga, out []byte) (err error) {
  size := tga.width * tga.height * tga.outPixelSize

  for i := 0; i < size && err == nil; {
    var b byte

    if b, err = tga.r.ReadByte(); err != nil {
      break
    }

    count := uint(b)

    if count&(1<<7) != 0 {
      // encoded packet
      count &= ^uint(1 << 7)

      if err = tga.getPixel(tga.tmp[:]); err == nil {
        for count++; count > 0 && i < size; count-- {
          copy(out[i:], tga.tmp[:tga.outPixelSize])
          i += tga.outPixelSize
        }
      }
    } else {
      // raw packet
      for count++; count > 0 && i < size && err == nil; count-- {
        err = tga.getPixel(out[i:])
        i += tga.outPixelSize
      }
    }
  }

  return
}

// flip32 flips pixels of image based on its origin.
func flip32(tga *tga, out []byte) {
  flipH := tga.raw.Flags&flagOriginRight != 0
  flipV := tga.raw.Flags&flagOriginTop == 0
  rowSize := tga.width * 4

  if flipH {
    for y := 0; y < tga.height; y++ {
      for x, offset := 0, y*rowSize; x < tga.width/2; x++ {
        a := out[offset+x*4:]
        b := out[offset+(tga.width-x-1)*4:]

        a[0], a[1], a[2], a[3], b[0], b[1], b[2], b[3] = b[0], b[1], b[2], b[3], a[0], a[1], a[2], a[3]
      }
    }
  }

  if flipV {
    for y := 0; y < tga.height/2; y++ {
      for x := 0; x < tga.width; x++ {
        a := out[y*rowSize+x*4:]
        b := out[(tga.height-y-1)*rowSize+x*4:]

        a[0], a[1], a[2], a[3], b[0], b[1], b[2], b[3] = b[0], b[1], b[2], b[3], a[0], a[1], a[2], a[3]
      }
    }
  }
}

// flip8 flips pixels of image based on its origin.
func flip8(tga *tga, out []byte) {
  flipH := tga.raw.Flags&flagOriginRight != 0
  flipV := tga.raw.Flags&flagOriginTop == 0
  rowSize := tga.width

  if flipH {
    for y := 0; y < tga.height; y++ {
      for x, offset := 0, y*rowSize; x < tga.width/2; x++ {
        ai, bi := offset+x, offset+(tga.width-x-1)
        out[ai], out[bi] = out[bi], out[ai]
      }
    }
  }

  if flipV {
    for y := 0; y < tga.height/2; y++ {
      for x := 0; x < tga.width; x++ {
        ai, bi := y*rowSize+x, (tga.height-y-1)*rowSize+x
        out[ai], out[bi] = out[bi], out[ai]
      }
    }
  }
}

// getHeader reads and validates TGA header.
func (tga *tga) getHeader() (err error) {
  if err = binary.Read(tga.r, binary.LittleEndian, &tga.raw); err != nil {
    return
  }

  if tga.raw.ImageType&imageTypeRLE != 0 {
    tga.decode = decodeRLE
  } else {
    tga.decode = decodeRaw
  }

  tga.raw.ImageType &= imageTypeMask
  flags := tga.raw.Flags & flagMask

  if flags != 0 && flags != 1 && flags != 8 {
    err = errors.New("invalid alpha size")
    return
  }

  tga.hasAlpha = ((flags != 0 || tga.raw.BPP == 32) ||
    (tga.raw.ImageType == imageTypeMonoChrome && tga.raw.BPP == 16) ||
    (tga.raw.ImageType == imageTypePaletted && tga.raw.PaletteBPP == 32))

  tga.width = int(tga.raw.Width)
  tga.height = int(tga.raw.Height)
  tga.pixelSize = int(tga.raw.BPP) >> 3

  // default is NOT premultiplied alpha model
  tga.ColorModel = color.NRGBAModel
  tga.outPixelSize = 4
  tga.flip = flip32

  if err == nil {
    err = tga.applyExtensions()
  }

  var formatIsInvalid bool

  switch tga.raw.ImageType {
  case imageTypePaletted:
    formatIsInvalid = (tga.raw.PaletteType == 0 ||
      tga.raw.BPP != 8 ||
      tga.raw.PaletteFirst >= tga.raw.PaletteLength ||
      (tga.raw.PaletteBPP != 15 && tga.raw.PaletteBPP != 16 && tga.raw.PaletteBPP != 24 && tga.raw.PaletteBPP != 32))
    tga.isPaletted = true

  case imageTypeTrueColor:
    formatIsInvalid = (tga.raw.BPP != 32 &&
      tga.raw.BPP != 16 &&
      (tga.raw.BPP != 24 || tga.hasAlpha))
    tga.ColorModel = color.NRGBAModel

  case imageTypeMonoChrome:
    formatIsInvalid = ((tga.hasAlpha && tga.raw.BPP != 16) ||
      (!tga.hasAlpha && tga.raw.BPP != 8))

    if !tga.hasAlpha && tga.raw.BPP == 8 {
      tga.ColorModel = color.GrayModel
      tga.outPixelSize = 1
      tga.flip = flip8
    }

  default:
    err = errors.New("invalid or unsupported image type")
  }

  if formatIsInvalid {
    err = errors.New("invalid image format")
  }

  return
}

func (tga *tga) getPixel(dst []byte) (err error) {
  var R, G, B, A uint8 = 0xff, 0xff, 0xff, 0xff
  src := tga.tmp

  if _, err = io.ReadFull(tga.r, src[0:tga.pixelSize]); err != nil {
    return
  }

  switch tga.pixelSize {
  case 4:
    if tga.hasAlpha {
      A = src[3]
    }
    fallthrough

  case 3:
    B, G, R = src[0], src[1], src[2]

  case 2:
    if tga.raw.ImageType == imageTypeMonoChrome {
      B, G, R = src[0], src[0], src[0]

      if tga.hasAlpha {
        A = src[1]
      }
    } else {
      word := uint16(src[0]) | (uint16(src[1]) << 8)
      B, G, R = wordToBGR(word)

      if tga.hasAlpha && (word&(1<<15)) == 0 {
        A = 0
      }
    }

  case 1:
    if tga.isPaletted {
      index := int(src[0])

      if int(index) >= tga.paletteLength {
        return errors.New("palette index out of range")
      }

      var m int

      if tga.raw.PaletteBPP == 24 {
        m = index * 3
        B, G, R = tga.palette[m+0], tga.palette[m+1], tga.palette[m+2]
      } else if tga.raw.PaletteBPP == 32 {
        m = index * 4
        B, G, R = tga.palette[m+0], tga.palette[m+1], tga.palette[m+2]

        if tga.hasAlpha {
          A = tga.palette[m+3]
        }
      } else if tga.raw.PaletteBPP == 16 {
        m = index * 2
        word := uint16(tga.palette[m+0]) | (uint16(tga.palette[m+1]) << 8)
        B, G, R = wordToBGR(word)
      }
    } else {
      dst[0] = src[0]
      return
    }
  }

  dst[0], dst[1], dst[2], dst[3] = R, G, B, A

  return nil
}

// wordToBGR converts 15-bit color to BGR
func wordToBGR(word uint16) (B, G, R uint8) {
  B = uint8((word >> 0) & 31)
  B = uint8((B << 3) + (B >> 2))
  G = uint8((word >> 5) & 31)
  G = uint8((G << 3) + (G >> 2))
  R = uint8((word >> 10) & 31)
  R = uint8((R << 3) + (R >> 2))
  return
}

// Local Variables:
// indent-tabs-mode: nil
// tab-width: 2
// fill-column: 70
// End:
// ex: set tabstop=2 shiftwidth=2 expandtab:
<commit_msg>Revert "use Gray color model for 8-bit monochrome images"<commit_after>package tga

import (
  "bytes"
  "encoding/binary"
  "errors"
  "image"
  "image/color"
  "io"
)

type rawHeader struct {
  IdLength      uint8
  PaletteType   uint8
  ImageType     uint8
  PaletteFirst  uint16
  PaletteLength uint16
  PaletteBPP    uint8
  OriginX       uint16
  PriginY       uint16
  Width         uint16
  Height        uint16
  BPP           uint8
  Flags         uint8
}

type rawFooter struct {
  ExtAreaOffset uint32
  DevDirOffset  uint32
  Signature     [18]byte // tgaSignature
}

type tga struct {
  r             *bytes.Reader
  raw           rawHeader
  isPaletted    bool
  hasAlpha      bool
  width         int
  height        int
  pixelSize     int
  palette       []byte
  paletteLength int
  ColorModel    color.Model
  tmp           [4]byte
  decode        func(tga *tga, out []byte) (err error)
}

const (
  flagOriginRight = uint8(1 << 4)
  flagOriginTop   = 1 << 5
  flagMask        = 7
)

const (
  imageTypePaletted   = uint8(1)
  imageTypeTrueColor  = 2
  imageTypeMonoChrome = 3
  imageTypeRLE        = 1 << 3
  imageTypeMask       = 3
)

const (
  tgaRawHeaderSize = 18
  tgaRawFooterSize = 26
  tgaSignature     = "TRUEVISION-XFILE.\x00"
)

func decode(r io.Reader) (outImage image.Image, err error) {
  var tga tga
  var data bytes.Buffer

  if _, err = data.ReadFrom(r); err != nil {
    return
  }

  tga.r = bytes.NewReader(data.Bytes())

  if err = tga.getHeader(); err != nil {
    return
  }

  // skip header
  if _, err = tga.r.Seek(int64(tgaRawHeaderSize+tga.raw.IdLength), 0); err != nil {
    return
  }

  if tga.isPaletted {
    // read palette
    entrySize := int((tga.raw.PaletteBPP + 1) >> 3)
    tga.paletteLength = int(tga.raw.PaletteLength - tga.raw.PaletteFirst)
    tga.palette = make([]byte, entrySize*tga.paletteLength)

    // skip to colormap
    if _, err = tga.r.Seek(int64(entrySize)*int64(tga.raw.PaletteFirst), 1); err != nil {
      return
    }

    if _, err = io.ReadFull(tga.r, tga.palette); err != nil {
      return
    }
  }

  var pixels []byte

  // choose a right color model
  if tga.ColorModel == color.NRGBAModel {
    im := image.NewNRGBA(image.Rect(0, 0, tga.width, tga.height))
    outImage = im
    pixels = im.Pix
  } else {
    im := image.NewRGBA(image.Rect(0, 0, tga.width, tga.height))
    outImage = im
    pixels = im.Pix
  }

  if err = tga.decode(&tga, pixels); err == nil {
    tga.flip(pixels)
  }

  return
}

func decodeConfig(r io.Reader) (cfg image.Config, err error) {
  var tga tga
  var data bytes.Buffer

  if _, err = data.ReadFrom(r); err != nil {
    return
  }

  tga.r = bytes.NewReader(data.Bytes())

  if err = tga.getHeader(); err == nil {
    cfg = image.Config{
      ColorModel: tga.ColorModel,
      Width:      tga.width,
      Height:     tga.height,
    }
  }

  return
}

func init() {
  image.RegisterFormat("tga", "", decode, decodeConfig)
}

// applyExtensions reads extensions section (if it exists) and parses attribute type.
func (tga *tga) applyExtensions() (err error) {
  var rawFooter rawFooter

  if _, err = tga.r.Seek(int64(-tgaRawFooterSize), 2); err != nil {
    return
  } else if err = binary.Read(tga.r, binary.LittleEndian, &rawFooter); err != nil {
    return
  } else if bytes.Equal(rawFooter.Signature[:], []byte(tgaSignature)) && rawFooter.ExtAreaOffset != 0 {
    offset := int64(rawFooter.ExtAreaOffset + 0x1ee)

    var n int64
    var t byte

    if n, err = tga.r.Seek(offset, 0); err != nil || n != offset {
      return
    } else if t, err = tga.r.ReadByte(); err != nil {
      return
    } else if t == 3 {
      // alpha
      tga.hasAlpha = true
      tga.ColorModel = color.NRGBAModel
    } else if t == 4 {
      // premultiplied alpha
      tga.hasAlpha = true
      tga.ColorModel = color.RGBAModel
    } else {
      // attribute is not an alpha channel value, ignore it
      tga.hasAlpha = false
    }
  }

  return
}

// decodeRaw decodes a raw (uncompressed) data.
func decodeRaw(tga *tga, out []byte) (err error) {
  for i := 0; i < len(out) && err == nil; i += 4 {
    err = tga.getPixel(out[i:])
  }

  return
}

// decodeRLE decodes run-length encoded data.
func decodeRLE(tga *tga, out []byte) (err error) {
  size := tga.width * tga.height * 4

  for i := 0; i < size && err == nil; {
    var b byte

    if b, err = tga.r.ReadByte(); err != nil {
      break
    }

    count := uint(b)

    if count&(1<<7) != 0 {
      // encoded packet
      count &= ^uint(1 << 7)

      if err = tga.getPixel(tga.tmp[:]); err == nil {
        for count++; count > 0 && i < size; count-- {
          copy(out[i:], tga.tmp[:])
          i += 4
        }
      }
    } else {
      // raw packet
      for count++; count > 0 && i < size && err == nil; count-- {
        err = tga.getPixel(out[i:])
        i += 4
      }
    }
  }

  return
}

// flip flips pixels of image based on its origin.
func (tga *tga) flip(out []byte) {
  flipH := tga.raw.Flags&flagOriginRight != 0
  flipV := tga.raw.Flags&flagOriginTop == 0
  rowSize := tga.width * 4

  if flipH {
    for y := 0; y < tga.height; y++ {
      for x, offset := 0, y*rowSize; x < tga.width/2; x++ {
        a := out[offset+x*4:]
        b := out[offset+(tga.width-x-1)*4:]

        a[0], a[1], a[2], a[3], b[0], b[1], b[2], b[3] = b[0], b[1], b[2], b[3], a[0], a[1], a[2], a[3]
      }
    }
  }

  if flipV {
    for y := 0; y < tga.height/2; y++ {
      for x := 0; x < tga.width; x++ {
        a := out[y*rowSize+x*4:]
        b := out[(tga.height-y-1)*rowSize+x*4:]

        a[0], a[1], a[2], a[3], b[0], b[1], b[2], b[3] = b[0], b[1], b[2], b[3], a[0], a[1], a[2], a[3]
      }
    }
  }
}

// getHeader reads and validates TGA header.
func (tga *tga) getHeader() (err error) {
  if err = binary.Read(tga.r, binary.LittleEndian, &tga.raw); err != nil {
    return
  }

  if tga.raw.ImageType&imageTypeRLE != 0 {
    tga.decode = decodeRLE
  } else {
    tga.decode = decodeRaw
  }

  tga.raw.ImageType &= imageTypeMask
  flags := tga.raw.Flags & flagMask

  if flags != 0 && flags != 1 && flags != 8 {
    err = errors.New("invalid alpha size")
    return
  }

  tga.hasAlpha = ((flags != 0 || tga.raw.BPP == 32) ||
    (tga.raw.ImageType == imageTypeMonoChrome && tga.raw.BPP == 16) ||
    (tga.raw.ImageType == imageTypePaletted && tga.raw.PaletteBPP == 32))

  tga.width = int(tga.raw.Width)
  tga.height = int(tga.raw.Height)
  tga.pixelSize = int(tga.raw.BPP) >> 3

  var formatIsInvalid bool

  switch tga.raw.ImageType {
  case imageTypePaletted:
    formatIsInvalid = (tga.raw.PaletteType == 0 ||
      tga.raw.BPP != 8 ||
      tga.raw.PaletteFirst >= tga.raw.PaletteLength ||
      (tga.raw.PaletteBPP != 15 && tga.raw.PaletteBPP != 16 && tga.raw.PaletteBPP != 24 && tga.raw.PaletteBPP != 32))
    tga.isPaletted = true

  case imageTypeTrueColor:
    formatIsInvalid = (tga.raw.BPP != 32 &&
      tga.raw.BPP != 16 &&
      (tga.raw.BPP != 24 || tga.hasAlpha))

  case imageTypeMonoChrome:
    formatIsInvalid = ((tga.hasAlpha && tga.raw.BPP != 16) ||
      (!tga.hasAlpha && tga.raw.BPP != 8))

  default:
    err = errors.New("invalid or unsupported image type")
  }

  // default is NOT premultiplied alpha model
  tga.ColorModel = color.NRGBAModel

  if formatIsInvalid {
    err = errors.New("invalid image format")
  } else if err == nil {
    err = tga.applyExtensions()
  }

  return
}

func (tga *tga) getPixel(dst []byte) (err error) {
  var R, G, B, A uint8 = 0xff, 0xff, 0xff, 0xff
  src := tga.tmp

  if _, err = io.ReadFull(tga.r, src[0:tga.pixelSize]); err != nil {
    return
  }

  switch tga.pixelSize {
  case 4:
    if tga.hasAlpha {
      A = src[3]
    }
    fallthrough

  case 3:
    B, G, R = src[0], src[1], src[2]

  case 2:
    if tga.raw.ImageType == imageTypeMonoChrome {
      B, G, R = src[0], src[0], src[0]

      if tga.hasAlpha {
        A = src[1]
      }
    } else {
      word := uint16(src[0]) | (uint16(src[1]) << 8)
      B, G, R = wordToBGR(word)

      if tga.hasAlpha && (word&(1<<15)) == 0 {
        A = 0
      }
    }

  case 1:
    if tga.isPaletted {
      index := int(src[0])

      if int(index) >= tga.paletteLength {
        return errors.New("palette index out of range")
      }

      var m int

      if tga.raw.PaletteBPP == 24 {
        m = index * 3
        B, G, R = tga.palette[m+0], tga.palette[m+1], tga.palette[m+2]
      } else if tga.raw.PaletteBPP == 32 {
        m = index * 4
        B, G, R = tga.palette[m+0], tga.palette[m+1], tga.palette[m+2]

        if tga.hasAlpha {
          A = tga.palette[m+3]
        }
      } else if tga.raw.PaletteBPP == 16 {
        m = index * 2
        word := uint16(tga.palette[m+0]) | (uint16(tga.palette[m+1]) << 8)
        B, G, R = wordToBGR(word)
      }
    } else {
      B, G, R = src[0], src[0], src[0]
    }
  }

  dst[0], dst[1], dst[2], dst[3] = R, G, B, A

  return nil
}

// wordToBGR converts 15-bit color to BGR
func wordToBGR(word uint16) (B, G, R uint8) {
  B = uint8((word >> 0) & 31)
  B = uint8((B << 3) + (B >> 2))
  G = uint8((word >> 5) & 31)
  G = uint8((G << 3) + (G >> 2))
  R = uint8((word >> 10) & 31)
  R = uint8((R << 3) + (R >> 2))
  return
}

// Local Variables:
// indent-tabs-mode: nil
// tab-width: 2
// fill-column: 70
// End:
// ex: set tabstop=2 shiftwidth=2 expandtab:
<|endoftext|>


<commit_before>package tnetstring

import (
	"bytes"
	"os"
	"reflect"
	"strconv"
)

func Unmarshal(data []byte, v interface{}) os.Error {
	val := reflect.ValueOf(v)
	val = reflect.Indirect(val)
	if !val.CanSet() {
		return os.NewError("tnetstring: Unmarshal requires a settable value")
	}
	_, err := unmarshal(data, val)
	return err
}

func indirect(v reflect.Value, create bool) reflect.Value {
	for {
		switch v.Kind() {
		case reflect.Ptr:
			if create && v.IsNil() {
				v.Set(reflect.New(v.Type().Elem()))
			}
			v = v.Elem()
		case reflect.Interface:
			if create && v.IsNil() {
				return v
			}
			v = v.Elem()
		default:
			return v
		}
	}
	panic("unreachable")
}

func unmarshal(data []byte, v reflect.Value) (int, os.Error) {
	typ, content, n := readElement(data)
	if n == 0 {
		return 0, os.NewError("tnetstring: invalid data")
	}
	v = indirect(v, true)
	kind := v.Kind()
	// ~ and interface types are special cases
	if typ != '~' && kind != reflect.Interface && typeLookup[kind] != typ {
		return 0, os.NewError("tnetstring: invalid value to unmarshal into")
	}
	switch typ {
	case '!':
		v.Set(reflect.ValueOf(string(content) == "true"))
	case '#':
		switch kind {
		case reflect.Int, reflect.Int8, reflect.Int16, reflect.Int32, reflect.Int64:
			i, err := strconv.Atoi64(string(content))
			if err != nil {
				return 0, err
			}
			v.SetInt(i)
		case reflect.Uint, reflect.Uint8, reflect.Uint16, reflect.Uint32,
			reflect.Uint64, reflect.Uintptr:
			ui, err := strconv.Atoui64(string(content))
			if err != nil {
				return 0, err
			}
			v.SetUint(ui)
		case reflect.Interface:
			i, err := strconv.Atoi64(string(content))
			if err != nil {
				return 0, err
			}
			v.Set(reflect.ValueOf(i))
		}
	case ',':
		v.Set(reflect.ValueOf(string(content)))
	case ']':
		unmarshalArray(content, v)
	case '}':
		var err os.Error
		if kind == reflect.Map {
			err = unmarshalMap(content, v)
		} else {
			err = unmarshalStruct(content, v)
		}
		if err != nil {
			return 0, err
		}
	case '~':
		switch kind {
		case reflect.Interface, reflect.Map, reflect.Ptr, reflect.Slice:
			v.Set(reflect.Zero(v.Type()))
		default:
			return 0, os.NewError("tnetstring: invalid value to unmarshal into")
		}
	default:
		return 0, os.NewError("tnetstring: unknown type")
	}
	return n, nil
}

func unmarshalArray(data []byte, v reflect.Value) os.Error {
	kind := v.Kind()
	n := 0
	i := 0
	elType := v.Type().Elem()
	for len(data)-n > 0 {
		if i >= v.Len() {
			if kind == reflect.Array {
				break
			} else {
				// TODO: could cut down on allocations by calling MakeSlice instead
				v.Set(reflect.Append(v, reflect.New(elType).Elem()))
			}
		}
		el := v.Index(i)
		i++
		nn, err := unmarshal(data[n:], el)
		if err != nil {
			return err
		}
		n += nn
	}
	return nil
}

func unmarshalMap(data []byte, v reflect.Value) os.Error {
	if v.Type().Key().Kind() != reflect.String {
		return os.NewError("tnetstring: only maps with string keys can be unmarshaled")
	}
	if v.IsNil() {
		v.Set(reflect.MakeMap(v.Type()))
	}
	n := 0
	vtype := v.Type().Elem()
	var s string
	key := reflect.ValueOf(&s).Elem()
	val := reflect.New(vtype).Elem()
	for len(data)-n > 0 {
		typ, content, nn := readElement(data[n:])
		if typ != ',' {
			return os.NewError("tnetstring: non-string key in dictionary")
		}
		s = string(content)
		n += nn
		nn, err := unmarshal(data[n:], val)
		if err != nil {
			return err
		}
		n += nn
		v.SetMapIndex(key, val)
	}
	return nil
}

func unmarshalStruct(data []byte, v reflect.Value) os.Error {
	n := 0
	structType := v.Type()
	var s string
	name := reflect.ValueOf(&s).Elem()
	for len(data)-n > 0 {
		nn, err := unmarshal(data[n:], name)
		if err != nil {
			return err
		}
		n += nn
		field := v.FieldByName(s)
		if field.Internal == nil {
			for i := 0; i < structType.NumField(); i++ {
				f := structType.Field(i)
				if f.Tag == s {
					field = v.Field(i)
					break
				}
			}
			if field.Internal == nil {
				var i interface{}
				field = reflect.ValueOf(&i).Elem()
			}
		}
		nn, err = unmarshal(data[n:], field)
		if err != nil {
			return err
		}
		n += nn
	}
	return nil
}

func readElement(data []byte) (typ byte, content []byte, n int) {
	col := bytes.IndexByte(data, ':')
	if col < 1 {
		return
	}
	n, err := strconv.Atoi(string(data[:col]))
	if err != nil || n > len(data[col+1:]) {
		return
	}
	// +1 for colon
	n += col + 1
	content = data[col+1 : n]
	typ = data[n]
	n++
	return
}
<commit_msg>Reduce allocations when unmarshaling to slice<commit_after>package tnetstring

import (
	"bytes"
	"os"
	"reflect"
	"strconv"
)

func Unmarshal(data []byte, v interface{}) os.Error {
	val := reflect.ValueOf(v)
	val = reflect.Indirect(val)
	if !val.CanSet() {
		return os.NewError("tnetstring: Unmarshal requires a settable value")
	}
	_, err := unmarshal(data, val)
	return err
}

func indirect(v reflect.Value, create bool) reflect.Value {
	for {
		switch v.Kind() {
		case reflect.Ptr:
			if create && v.IsNil() {
				v.Set(reflect.New(v.Type().Elem()))
			}
			v = v.Elem()
		case reflect.Interface:
			if create && v.IsNil() {
				return v
			}
			v = v.Elem()
		default:
			return v
		}
	}
	panic("unreachable")
}

func unmarshal(data []byte, v reflect.Value) (int, os.Error) {
	typ, content, n := readElement(data)
	if n == 0 {
		return 0, os.NewError("tnetstring: invalid data")
	}
	v = indirect(v, true)
	kind := v.Kind()
	// ~ and interface types are special cases
	if typ != '~' && kind != reflect.Interface && typeLookup[kind] != typ {
		return 0, os.NewError("tnetstring: invalid value to unmarshal into")
	}
	switch typ {
	case '!':
		v.Set(reflect.ValueOf(string(content) == "true"))
	case '#':
		switch kind {
		case reflect.Int, reflect.Int8, reflect.Int16, reflect.Int32, reflect.Int64:
			i, err := strconv.Atoi64(string(content))
			if err != nil {
				return 0, err
			}
			v.SetInt(i)
		case reflect.Uint, reflect.Uint8, reflect.Uint16, reflect.Uint32,
			reflect.Uint64, reflect.Uintptr:
			ui, err := strconv.Atoui64(string(content))
			if err != nil {
				return 0, err
			}
			v.SetUint(ui)
		case reflect.Interface:
			i, err := strconv.Atoi64(string(content))
			if err != nil {
				return 0, err
			}
			v.Set(reflect.ValueOf(i))
		}
	case ',':
		v.Set(reflect.ValueOf(string(content)))
	case ']':
		unmarshalArray(content, v)
	case '}':
		var err os.Error
		if kind == reflect.Map {
			err = unmarshalMap(content, v)
		} else {
			err = unmarshalStruct(content, v)
		}
		if err != nil {
			return 0, err
		}
	case '~':
		switch kind {
		case reflect.Interface, reflect.Map, reflect.Ptr, reflect.Slice:
			v.Set(reflect.Zero(v.Type()))
		default:
			return 0, os.NewError("tnetstring: invalid value to unmarshal into")
		}
	default:
		return 0, os.NewError("tnetstring: unknown type")
	}
	return n, nil
}

func unmarshalArray(data []byte, v reflect.Value) os.Error {
	kind := v.Kind()
	n := 0
	i := 0
	elType := v.Type().Elem()
	elVal := reflect.Zero(elType)
	for len(data)-n > 0 {
		if i >= v.Len() {
			if kind == reflect.Array {
				break
			} else {
				v.Set(reflect.Append(v, elVal))
			}
		}
		el := v.Index(i)
		i++
		nn, err := unmarshal(data[n:], el)
		if err != nil {
			return err
		}
		n += nn
	}
	return nil
}

func unmarshalMap(data []byte, v reflect.Value) os.Error {
	if v.Type().Key().Kind() != reflect.String {
		return os.NewError("tnetstring: only maps with string keys can be unmarshaled")
	}
	if v.IsNil() {
		v.Set(reflect.MakeMap(v.Type()))
	}
	n := 0
	vtype := v.Type().Elem()
	var s string
	key := reflect.ValueOf(&s).Elem()
	val := reflect.New(vtype).Elem()
	for len(data)-n > 0 {
		typ, content, nn := readElement(data[n:])
		if typ != ',' {
			return os.NewError("tnetstring: non-string key in dictionary")
		}
		s = string(content)
		n += nn
		nn, err := unmarshal(data[n:], val)
		if err != nil {
			return err
		}
		n += nn
		v.SetMapIndex(key, val)
	}
	return nil
}

func unmarshalStruct(data []byte, v reflect.Value) os.Error {
	n := 0
	structType := v.Type()
	var s string
	name := reflect.ValueOf(&s).Elem()
	for len(data)-n > 0 {
		nn, err := unmarshal(data[n:], name)
		if err != nil {
			return err
		}
		n += nn
		field := v.FieldByName(s)
		if field.Internal == nil {
			for i := 0; i < structType.NumField(); i++ {
				f := structType.Field(i)
				if f.Tag == s {
					field = v.Field(i)
					break
				}
			}
			if field.Internal == nil {
				var i interface{}
				field = reflect.ValueOf(&i).Elem()
			}
		}
		nn, err = unmarshal(data[n:], field)
		if err != nil {
			return err
		}
		n += nn
	}
	return nil
}

func readElement(data []byte) (typ byte, content []byte, n int) {
	col := bytes.IndexByte(data, ':')
	if col < 1 {
		return
	}
	n, err := strconv.Atoi(string(data[:col]))
	if err != nil || n > len(data[col+1:]) {
		return
	}
	// +1 for colon
	n += col + 1
	content = data[col+1 : n]
	typ = data[n]
	n++
	return
}
<|endoftext|>


<commit_before>package main

import (
	"flag"
	"log"
	"net/http"
	"time"

	drive "google.golang.org/api/drive/v2"
	"google.golang.org/api/googleapi"
)

const fields = googleapi.Field("items(id,title,md5Checksum),nextPageToken")

var tok = flag.String("tok", "", "OAuth token")

func main() {
	flag.Parse()

	svc, err := drive.New(&http.Client{Transport: authTransport{*tok}})
	if err != nil {
		log.Fatal(err)
	}
	pageToken := ""

	type file struct {
		id, title string
	}

	md5s := map[string][]file{}
	for i := 0; i < 10; i++ {
		log.Printf("querying page token %q\n", pageToken)
		fs, err := svc.Files.List().
			MaxResults(1000).
			PageToken(pageToken).
			Fields(fields).
			Do()
		if err != nil {
			log.Fatal(err)
		}
		for _, f := range fs.Items {
			if f.Md5Checksum != "" {
				md5s[f.Md5Checksum] = append(md5s[f.Md5Checksum], file{f.Id, f.Title})
			}
		}
		pageToken = fs.NextPageToken
		if pageToken == "" {
			break
		}
		time.Sleep(time.Second)
	}

	for k, v := range md5s {
		if len(v) > 1 {
			log.Println("===", k, "===")
			for _, v := range v[1:] {
				log.Println(v.id, v.title)
			}
		}
	}
}

type authTransport struct {
	tok string
}

func (t authTransport) RoundTrip(r *http.Request) (*http.Response, error) {
	r.Header.Set("Authorization", "Bearer "+t.tok)
	return http.DefaultTransport.RoundTrip(r)
}
<commit_msg>better reporting, almost appifyable<commit_after>package main

import (
	"flag"
	"fmt"
	"net/http"
	"os"
	"strings"
	"time"

	drive "google.golang.org/api/drive/v2"
	"google.golang.org/api/googleapi"
)

const fields = googleapi.Field("items(id,title,fileSize,md5Checksum),nextPageToken")

var tok = flag.String("tok", "", "OAuth token")

func main() {
	flag.Parse()

	svc, err := drive.New(&http.Client{Transport: authTransport{*tok}})
	if err != nil {
		fmt.Println(err)
		os.Exit(1)
	}
	pageToken := ""

	type file struct {
		md5, title string
		size       int64
	}

	scannedFiles := 0
	md5s := map[file][]string{}
	fmt.Print("scanning files")
	for {
		fmt.Print(".")
		fs, err := svc.Files.List().
			MaxResults(1000).
			PageToken(pageToken).
			Fields(fields).
			Do()
		if err != nil {
			fmt.Println(err)
			os.Exit(1)
		}
		scannedFiles += len(fs.Items)
		for _, f := range fs.Items {
			if f.Md5Checksum != "" {
				k := file{f.Md5Checksum, f.Title, f.FileSize}
				md5s[k] = append(md5s[k], f.Id)
			}
		}
		pageToken = fs.NextPageToken
		if pageToken == "" {
			break
		}
		time.Sleep(time.Second)
	}

	var totalFiles int
	var totalSize int64
	for k, v := range md5s {
		if len(v) > 1 {
			fmt.Println("===", k.title, "(md5:", k.md5, ") ===")
			fmt.Println("- reapable IDs:", strings.Join(v[1:], ", "))
			totalFiles += len(v) - 1
			totalSize += k.size * int64(len(v)-1)
		}
	}
	fmt.Println("scanned", scannedFiles, "files")
	fmt.Println("can reap", totalFiles, "files")
	fmt.Println("can reclaim", float64(totalSize)/1024/1024/1024, "GiB of space")
}

type authTransport struct {
	tok string
}

func (t authTransport) RoundTrip(r *http.Request) (*http.Response, error) {
	r.Header.Set("Authorization", "Bearer "+t.tok)
	return http.DefaultTransport.RoundTrip(r)
}
<|endoftext|>


<commit_before>package restart

import (
	"fmt"
	"github.com/conductant/gohm/pkg/registry"
	"github.com/conductant/gohm/pkg/template"
	"github.com/conductant/gohm/pkg/zk"
	"github.com/golang/glog"
	. "github.com/infradash/dash/pkg/dash"
	"github.com/infradash/dash/pkg/executor"
	"golang.org/x/net/context"
	"path"
	"strconv"
	"strings"
	"time"
)

type Restart struct {
	RegistryReleaseEntry
	ZkSettings
	RestartConfig

	Initializer *ConfigLoader `json:"-"`

	ExecuteForReal bool
	reg            registry.Registry
}

func mustNot(err error) {
	if err != nil {
		panic(err)
	}
}

func (this *Restart) Run() error {

	if this.Initializer == nil {
		return ErrNoConfig
	}

	url := "zk://" + this.Hosts
	ctx := zk.ContextPutTimeout(context.Background(), this.Timeout)
	reg, err := registry.Dial(ctx, url)
	mustNot(err)
	defer reg.Close()

	this.reg = reg

	// We don't want the application of template to wipe out Domain, Service, etc. variables
	// So escape them.
	this.Initializer.Context = EscapeVars(ConfigVariables...)

	if this.Initializer.ConfigUrl != "" {
		loaded := false
		var err error
		for {
			loaded, err = this.Initializer.Load(this, "", nil)

			if !loaded || err != nil {
				glog.Infoln("Wait then retry:", err)
				time.Sleep(2 * time.Second)

			} else {
				break
			}
		}
	}

	glog.Infoln("Restarting",
		"Domain=", this.Domain, "Service=", this.Service, "Version=", this.Version, "Path=", this.Path,
		"Image=", this.Image, "Build=", this.Build)

	if this.Controller == "" {
		this.Controller = this.Service + "-controller"
	}

	glog.Infoln("ProxyUrl=", this.ProxyUrl)
	glog.Infoln("ControllerPath=", this.GetControllerPath())
	glog.Infoln("MemberWatchPath=", this.GetMemberWatchPath())
	glog.Infoln("ProxyWatchPath=", this.GetProxyWatchPath())
	glog.Infoln("RestartWaitDuration=", this.RestartConfig.RestartWaitDuration.Duration)

	// Load the list of controllers
	controllerPath := this.GetControllerPath()
	controllerPaths, err := reg.List(controllerPath)
	mustNot(err)

	rollingCount := len(controllerPaths)
	glog.Infoln("Checking on containers at", controllerPath, "count=", rollingCount)

	memberWatchPath := this.GetMemberWatchPath()

	// get a count too
	memberPaths, err := reg.List(memberWatchPath)
	memberCount := len(memberPaths)

	glog.Infoln("Watching members in", memberWatchPath, "count=", memberCount)

	if rollingCount != memberCount {
		// It's possibly a misconfiguration.  We expect the member count to be 1:1 with controller count
		panic(fmt.Errorf("Count mismatch: controllers=", rollingCount, "containers=", memberCount))
	}

	clients := []*executor.Client{}
	for _, controllerPath := range controllerPaths {
		glog.Infoln("Found controller path", controllerPath)
		if read, _, err := reg.Get(controllerPath); err == nil {
			host, ps := ParseHostPort(string(read))
			port, err := strconv.Atoi(ps)
			mustNot(err)
			client := executor.NewClient().SetHost(host).SetPort(port)
			clients = append(clients, client)
		}
	}

	if len(clients) != len(controllerPaths) {
		panic("Cannot access all controllers")
	}

	// Get the current value of the watch
	proxyWatchPath := this.GetProxyWatchPath()
	watchValueString, watchValueVersion, err := reg.Get(proxyWatchPath)
	if err != nil {
		panic(fmt.Errorf("Cannot get live watch value:%v", err))
	}
	watchValue, err := strconv.Atoi(string(watchValueString))
	if err != nil {
		panic(fmt.Errorf("Cannot get integer value for watch: %s, err=%v", proxyWatchPath, err))
	}

	beginRollingRestart := make(chan int)
	incrementedWatch := make(chan int)
	restartComplete := make(chan int)
	processComplete := make(chan int)
	go func() {
		<-beginRollingRestart
		glog.Infoln("Received signal -- begin rolling restart....")
		for _, c := range clients {
			if this.ProxyUrl != "" {
				c.SetProxyUrl(this.ProxyUrl)
			}
			info, err := c.GetInfo()
			glog.Infoln("Info=", info, "err=", err)

			// Send a kill
			err = c.RemoteKill()
			glog.Infoln("Kill=", err)

			<-incrementedWatch
		}

		// here we poll until the count is restored
		for {
			c, err := reg.List(memberWatchPath)
			if err != nil {
				panic(err)
			}
			if len(c) == len(clients) {
				break
			}
			glog.Infoln("Members count don't match original client count.  Waiting....")
			time.Sleep(this.RestartWaitDuration.Duration)
		}
		restartComplete <- 0
	}()

	// Now we set up the watch
	memberChanges, memberWatchStop, err := reg.Trigger(registry.Members{Path: memberWatchPath})
	mustNot(err)
	go func() {

		version := watchValueVersion

		for itr := 1; ; itr++ {

			select {
			case <-memberChanges:
				glog.Infoln("Received membership change. Incrementing proxy watch:", proxyWatchPath)

				var err error

				newVal := []byte(fmt.Sprintf("%d", watchValue+itr))
				version, err = reg.PutVersion(proxyWatchPath, newVal, version)
				if err != nil {
					panic(fmt.Errorf("Failed to update watch: %s, err=%v", proxyWatchPath, err))
				}

				glog.Infoln("Wait", this.RestartWaitDuration.Duration, "before shutting down next node.")
				time.Sleep(this.RestartWaitDuration.Duration)
				incrementedWatch <- itr

			case <-restartComplete:

				time.Sleep(this.RestartWaitDuration.Duration)

				// Send one more...
				glog.Infoln("Received restartComplete. Incrementing proxy watch:", proxyWatchPath)

				newVal := []byte(fmt.Sprintf("%d", watchValue+itr))
				version, err = reg.PutVersion(proxyWatchPath, newVal, version)
				if err != nil {
					panic(fmt.Errorf("Failed to update watch: %s, err=%v", proxyWatchPath, err))
				}

				glog.Infoln("Restart process completed.")
				processComplete <- 0
				return
			}
		}
	}()

	if this.ExecuteForReal {
		glog.Infoln("Begin rolling restart.")
		beginRollingRestart <- 0
	} else {
		glog.Infoln("Not going to commit.  Use --restart.commit flag to actually begin.")
		return nil
	}

	<-processComplete
	memberWatchStop <- 0

	return nil
}

func (this *Restart) GetControllerPath() registry.Path {
	applied, err := template.Apply([]byte(this.ControllerPathFormat), this)
	mustNot(err)
	return registry.NewPath(string(applied))
}

func (this *Restart) GetMemberWatchPath() registry.Path {
	applied, err := template.Apply([]byte(this.MemberWatchPathFormat), this)
	mustNot(err)
	return registry.NewPath(string(applied))
}

func (this *Restart) GetProxyWatchPath() registry.Path {
	applied, err := template.Apply([]byte(this.ProxyWatchPathFormat), this)
	mustNot(err)
	return registry.NewPath(string(applied))
}

func (this *Restart) RunningImage() string {
	if this.Image == "" {
		// Determine the current image based on zk values
		applied, err := template.Apply([]byte(this.CurrentImagePathFormat), this)
		mustNot(err)
		p := registry.NewPath(string(applied))
		buff, _, err := this.reg.Get(p)
		mustNot(err)
		fullPath := strings.Split(string(buff), ",")[0]
		d, f := path.Split(fullPath)
		this.Image = path.Join(path.Base(d), f)
		glog.Infoln("Getting current image from", p, "value=", this.Image)

	}
	return this.Image
}
<commit_msg>Handling the case where zk nodes aren't deleted quickly enough by the agents; don't hang<commit_after>package restart

import (
	"fmt"
	"github.com/conductant/gohm/pkg/registry"
	"github.com/conductant/gohm/pkg/template"
	"github.com/conductant/gohm/pkg/zk"
	"github.com/golang/glog"
	. "github.com/infradash/dash/pkg/dash"
	"github.com/infradash/dash/pkg/executor"
	"golang.org/x/net/context"
	"path"
	"strconv"
	"strings"
	"time"
)

type Restart struct {
	RegistryReleaseEntry
	ZkSettings
	RestartConfig

	Initializer *ConfigLoader `json:"-"`

	ExecuteForReal bool
	reg            registry.Registry
}

func mustNot(err error) {
	if err != nil {
		panic(err)
	}
}

func (this *Restart) Run() error {

	if this.Initializer == nil {
		return ErrNoConfig
	}

	url := "zk://" + this.Hosts
	ctx := zk.ContextPutTimeout(context.Background(), this.Timeout)
	reg, err := registry.Dial(ctx, url)
	mustNot(err)
	defer reg.Close()

	this.reg = reg

	// We don't want the application of template to wipe out Domain, Service, etc. variables
	// So escape them.
	this.Initializer.Context = EscapeVars(ConfigVariables...)

	if this.Initializer.ConfigUrl != "" {
		loaded := false
		var err error
		for {
			loaded, err = this.Initializer.Load(this, "", nil)

			if !loaded || err != nil {
				glog.Infoln("Wait then retry:", err)
				time.Sleep(2 * time.Second)

			} else {
				break
			}
		}
	}

	glog.Infoln("Restarting",
		"Domain=", this.Domain, "Service=", this.Service, "Version=", this.Version, "Path=", this.Path,
		"Image=", this.Image, "Build=", this.Build)

	if this.Controller == "" {
		this.Controller = this.Service + "-controller"
	}

	glog.Infoln("ProxyUrl=", this.ProxyUrl)
	glog.Infoln("ControllerPath=", this.GetControllerPath())
	glog.Infoln("MemberWatchPath=", this.GetMemberWatchPath())
	glog.Infoln("ProxyWatchPath=", this.GetProxyWatchPath())
	glog.Infoln("RestartWaitDuration=", this.RestartConfig.RestartWaitDuration.Duration)

	// Load the list of controllers
	controllerPath := this.GetControllerPath()
	controllerPaths, err := reg.List(controllerPath)
	mustNot(err)

	rollingCount := len(controllerPaths)
	glog.Infoln("Checking on containers at", controllerPath, "count=", rollingCount)

	memberWatchPath := this.GetMemberWatchPath()

	// get a count too
	memberPaths, err := reg.List(memberWatchPath)
	memberCount := len(memberPaths)

	glog.Infoln("Watching members in", memberWatchPath, "count=", memberCount)

	if rollingCount != memberCount {
		// It's possibly a misconfiguration.  We expect the member count to be 1:1 with controller count
		panic(fmt.Errorf("Count mismatch: controllers=", rollingCount, "containers=", memberCount))
	}

	clients := []*executor.Client{}
	for _, controllerPath := range controllerPaths {
		glog.Infoln("Found controller path", controllerPath)
		if read, _, err := reg.Get(controllerPath); err == nil {
			host, ps := ParseHostPort(string(read))
			port, err := strconv.Atoi(ps)
			mustNot(err)
			client := executor.NewClient().SetHost(host).SetPort(port)
			clients = append(clients, client)
		}
	}

	if len(clients) != len(controllerPaths) {
		panic("Cannot access all controllers")
	}

	// Get the current value of the watch
	proxyWatchPath := this.GetProxyWatchPath()
	watchValueString, watchValueVersion, err := reg.Get(proxyWatchPath)
	if err != nil {
		panic(fmt.Errorf("Cannot get live watch value:%v", err))
	}
	watchValue, err := strconv.Atoi(string(watchValueString))
	if err != nil {
		panic(fmt.Errorf("Cannot get integer value for watch: %s, err=%v", proxyWatchPath, err))
	}

	beginRollingRestart := make(chan int)
	incrementedWatch := make(chan int)
	restartComplete := make(chan int)
	processComplete := make(chan int)
	go func() {
		<-beginRollingRestart
		glog.Infoln("Received signal -- begin rolling restart....")
		for _, c := range clients {
			if this.ProxyUrl != "" {
				c.SetProxyUrl(this.ProxyUrl)
			}
			info, err := c.GetInfo()
			glog.Infoln("Info=", info, "err=", err)

			// Send a kill
			err = c.RemoteKill()
			glog.Infoln("Kill=", err)

			<-incrementedWatch
		}

		// here we poll until the count is restored
		for {
			c, err := reg.List(memberWatchPath)
			if err != nil {
				panic(err)
			}
			// In some cases the node member doesn't get deleted soon enough...
			if len(c) >= len(clients) {
				break
			}
			glog.Infoln("Members count don't match original client count.  Waiting....")
			time.Sleep(this.RestartWaitDuration.Duration)
		}
		restartComplete <- 0
	}()

	// Now we set up the watch
	memberChanges, memberWatchStop, err := reg.Trigger(registry.Members{Path: memberWatchPath})
	mustNot(err)
	go func() {

		version := watchValueVersion

		for itr := 1; ; itr++ {

			select {
			case <-memberChanges:
				glog.Infoln("Received membership change. Incrementing proxy watch:", proxyWatchPath)

				var err error

				newVal := []byte(fmt.Sprintf("%d", watchValue+itr))
				version, err = reg.PutVersion(proxyWatchPath, newVal, version)
				if err != nil {
					panic(fmt.Errorf("Failed to update watch: %s, err=%v", proxyWatchPath, err))
				}

				glog.Infoln("Wait", this.RestartWaitDuration.Duration, "before shutting down next node.")
				time.Sleep(this.RestartWaitDuration.Duration)
				incrementedWatch <- itr

			case <-restartComplete:

				time.Sleep(this.RestartWaitDuration.Duration)

				// Send one more...
				glog.Infoln("Received restartComplete. Incrementing proxy watch:", proxyWatchPath)

				newVal := []byte(fmt.Sprintf("%d", watchValue+itr))
				version, err = reg.PutVersion(proxyWatchPath, newVal, version)
				if err != nil {
					panic(fmt.Errorf("Failed to update watch: %s, err=%v", proxyWatchPath, err))
				}

				glog.Infoln("Restart process completed.")
				processComplete <- 0
				return
			}
		}
	}()

	if this.ExecuteForReal {
		glog.Infoln("Begin rolling restart.")
		beginRollingRestart <- 0
	} else {
		glog.Infoln("Not going to commit.  Use --restart.commit flag to actually begin.")
		return nil
	}

	<-processComplete
	memberWatchStop <- 0

	return nil
}

func (this *Restart) GetControllerPath() registry.Path {
	applied, err := template.Apply([]byte(this.ControllerPathFormat), this)
	mustNot(err)
	return registry.NewPath(string(applied))
}

func (this *Restart) GetMemberWatchPath() registry.Path {
	applied, err := template.Apply([]byte(this.MemberWatchPathFormat), this)
	mustNot(err)
	return registry.NewPath(string(applied))
}

func (this *Restart) GetProxyWatchPath() registry.Path {
	applied, err := template.Apply([]byte(this.ProxyWatchPathFormat), this)
	mustNot(err)
	return registry.NewPath(string(applied))
}

func (this *Restart) RunningImage() string {
	if this.Image == "" {
		// Determine the current image based on zk values
		applied, err := template.Apply([]byte(this.CurrentImagePathFormat), this)
		mustNot(err)
		p := registry.NewPath(string(applied))
		buff, _, err := this.reg.Get(p)
		mustNot(err)
		fullPath := strings.Split(string(buff), ",")[0]
		d, f := path.Split(fullPath)
		this.Image = path.Join(path.Base(d), f)
		glog.Infoln("Getting current image from", p, "value=", this.Image)

	}
	return this.Image
}
<|endoftext|>


<commit_before>package deploy

import (
	"bufio"
	"errors"
	"fmt"
	"io"
	"io/ioutil"
	"os"
	"regexp"
	"strings"
	"time"

	"github.com/remind101/deploy/Godeps/_workspace/src/github.com/codegangsta/cli"
	"github.com/remind101/deploy/Godeps/_workspace/src/github.com/github/hub/git"
	hub "github.com/remind101/deploy/Godeps/_workspace/src/github.com/github/hub/github"
	"github.com/remind101/deploy/Godeps/_workspace/src/github.com/google/go-github/github"
)

const (
	Name  = "deploy"
	Usage = "A command for creating GitHub deployments"
)

const (
	DefaultRef     = "master"
	DefaultTimeout = 20 * time.Second
)

var errTimeout = errors.New("Timed out waiting for build to start. Did you add a webhook to handle deployment events?")

func init() {
	cli.AppHelpTemplate = `USAGE:
   # Deploy the master branch of remind101/acme-inc to staging
   {{.Name}} --env=staging --ref=master remind101/acme-inc

   # Deploy HEAD of the current branch to staging
   {{.Name}} --env=staging remind101/acme-inc

   # Deploy the current GitHub repo to staging
   {{.Name}} --env=staging
{{if .Flags}}
OPTIONS:
   {{range .Flags}}{{.}}
   {{end}}{{end}}
`
}

var env string
var ProtectedEnvironments = map[string]bool{
	"production": true,
}

var flags = []cli.Flag{
	cli.StringFlag{
		Name:  "ref, branch, commit, tag",
		Value: "",
		Usage: "The git ref to deploy. Can be a git commit, branch or tag.",
	},
	cli.StringFlag{
		Name:        "env, e",
		Value:       "",
		Usage:       "The environment to deploy to.",
		Destination: &env,
	},
	cli.BoolFlag{
		Name:  "force, f",
		Usage: "Ignore commit status checks.",
	},
	cli.BoolFlag{
		Name:  "detached, d",
		Usage: "Don't wait for the deployment to complete.",
	},
	cli.BoolFlag{
		Name:  "quiet, q",
		Usage: "Silence any output to STDOUT.",
	},
	cli.BoolFlag{
		Name:  "update, u",
		Usage: "Update the binary",
	},
}

// NewApp returns a new cli.App for the deploy command.
func NewApp() *cli.App {
	app := cli.NewApp()
	app.Version = Version
	app.Name = Name
	app.Usage = Usage
	app.Flags = flags
	app.Action = func(c *cli.Context) {
		if c.Bool("update") {
			updater := NewUpdater()
			if err := updater.Update(); err != nil {
				fmt.Println(err)
				os.Exit(-1)
			} else {
				os.Exit(0)
			}
		}

		if err := RunDeploy(c); err != nil {
			msg := err.Error()
			if err, ok := err.(*github.ErrorResponse); ok {
				if strings.HasPrefix(err.Message, "Conflict: Commit status checks failed for") {
					msg = "Commit status checks failed. You can bypass commit status checks with the --force flag."
				} else if strings.HasPrefix(err.Message, "No ref found for") {
					msg = fmt.Sprintf("%s. Did you push it to GitHub?", err.Message)
				} else {
					msg = err.Message
				}
			}

			fmt.Println(msg)
			os.Exit(-1)
		}
	}

	return app
}

// RunDeploy performs a deploy.
func RunDeploy(c *cli.Context) error {
	var w io.Writer
	if c.Bool("quiet") {
		w = ioutil.Discard
	} else {
		w = c.App.Writer
	}

	h, err := hub.CurrentConfig().PromptForHost("github.com")
	if err != nil {
		return err
	}

	client, err := newGitHubClient(h)
	if err != nil {
		return err
	}

	nwo, err := Repo(c.Args())
	if err != nil {
		return err
	}

	owner, repo, err := SplitRepo(nwo, os.Getenv("GITHUB_ORGANIZATION"))
	if err != nil {
		return fmt.Errorf("Invalid GitHub repo: %s", nwo)
	}

	if env == "" {
		return fmt.Errorf("--env flag is required")
	}
	env = aliasEnvironment(env)

	err = displayNewCommits(owner, repo, c, client)
	if err != nil {
		return err
	}

	r, err := newDeploymentRequest(c)
	if err != nil {
		return err
	}

	fmt.Fprintf(w, "Deploying %s/%s@%s to %s...\n", owner, repo, *r.Ref, *r.Environment)

	if c.Bool("detached") {
		return nil
	}

	d, _, err := client.Repositories.CreateDeployment(owner, repo, r)
	if err != nil {
		return err
	}

	started := make(chan *github.DeploymentStatus)
	completed := make(chan *github.DeploymentStatus)

	go func() {
		started <- waitState(pendingStates, owner, repo, *d.ID, client)
	}()

	go func() {
		completed <- waitState(completedStates, owner, repo, *d.ID, client)
	}()

	select {
	case <-time.After(DefaultTimeout):
		return errTimeout
	case status := <-started:
		var url string
		if status.TargetURL != nil {
			url = *status.TargetURL
		}
		fmt.Fprintf(w, "%s\n", url)
	}

	status := <-completed

	if isFailed(*status.State) {
		return errors.New("Failed to deploy")
	}

	return nil
}

func displayNewCommits(owner string, repo string, c *cli.Context, client *github.Client) error {
	ref := Ref(c.String("ref"), git.Head)

	opt := &github.DeploymentsListOptions{
		Environment: env,
	}

	deployments, _, err := client.Repositories.ListDeployments(owner, repo, opt)
	if err != nil || len(deployments) == 0 {
		return err
	}

	sha := *deployments[0].SHA
	compare, _, err := client.Repositories.CompareCommits(owner, repo, sha, ref)
	if err != nil || len(compare.Commits) == 0 {
		return err
	}

	fmt.Println("Deploying the following commits:\n")
	for _, commit := range compare.Commits {
		message := *commit.Commit.Message
		fmt.Printf("%-20s\t%s\n", *commit.Commit.Author.Name, strings.Split(message, "\n")[0])
	}
	fmt.Printf("\nSee entire diff here: https://github.com/%s/%s/compare/%s...%s\n\n", owner, repo, sha, ref)
	return nil
}

var EnvironmentAliases = map[string]string{
	"prod":  "production",
	"stage": "staging",
}

func aliasEnvironment(env string) string {
	if a, ok := EnvironmentAliases[env]; ok {
		return a
	}

	return env
}

func newDeploymentRequest(c *cli.Context) (*github.DeploymentRequest, error) {
	ref := Ref(c.String("ref"), git.Head)

	if ProtectedEnvironments[env] {
		yes := askYN(fmt.Sprintf("Are you sure you want to deploy %s to %s?", ref, env))
		if !yes {
			return nil, fmt.Errorf("Deployment aborted.")
		}
	}

	var contexts *[]string
	if c.Bool("force") {
		s := []string{}
		contexts = &s
	}

	return &github.DeploymentRequest{
		Ref:              github.String(ref),
		Task:             github.String("deploy"),
		AutoMerge:        github.Bool(false),
		Environment:      github.String(env),
		RequiredContexts: contexts,
		Payload: map[string]interface{}{
			"force": c.Bool("force"),
		},
		// TODO Description:
	}, nil
}

var (
	pendingStates   = []string{"pending"}
	completedStates = []string{"success", "error", "failure"}
)

func isFailed(state string) bool {
	return state == "error" || state == "failure"
}

// waitState waits for a deployment status that matches the given states, then
// sends on the returned channel.
func waitState(states []string, owner, repo string, deploymentID int, c *github.Client) *github.DeploymentStatus {
	for {
		<-time.After(1 * time.Second)

		statuses, _, err := c.Repositories.ListDeploymentStatuses(owner, repo, deploymentID, nil)
		if err != nil {
			continue
		}

		status := firstStatus(states, statuses)
		if status != nil {
			return status
		}
	}
}

// firstStatus takes a slice of github.DeploymentStatus and returns the
// first status that matches the provided slice of states.
func firstStatus(states []string, statuses []github.DeploymentStatus) *github.DeploymentStatus {
	for _, ds := range statuses {
		for _, s := range states {
			if ds.State != nil && *ds.State == s {
				return &ds
			}
		}
	}

	return nil
}

// refRegex is a regular expression that matches a full git HEAD ref.
var refRegex = regexp.MustCompile(`^refs/heads/(.*)$`)

// Ref attempts to return the proper git ref to deploy. If a ref is provided,
// that will be returned. If not, it will fallback to calling headFunc. If an
// error is returned (not in a git repo), then it will fallback to DefaultRef.
func Ref(ref string, headFunc func() (string, error)) string {
	if ref != "" {
		return ref
	}

	ref, err := headFunc()
	if err != nil {
		// An error means that we're either not in a GitRepo or we're
		// not on a branch. In this case, we just fallback to the
		// DefaultRef.
		return DefaultRef
	}

	// Convert `refs/heads/test-deploy` => `test-deploy`
	return refRegex.ReplaceAllString(ref, "$1")
}

// Repo will determine the correct GitHub repo to deploy to, based on a set of
// arguments.
func Repo(arguments []string) (string, error) {
	if len(arguments) != 0 {
		return arguments[0], nil
	}

	remotes, err := hub.Remotes()
	if err != nil {
		return "", err
	}

	repo := GitHubRepo(remotes)
	if repo == "" {
		return repo, errors.New("no GitHub repo found in .git/config")
	}

	return repo, nil
}

// A regular expression that can convert a URL.Path into a GitHub repo name.
var remoteRegex = regexp.MustCompile(`^/(.*)\.git$`)

// GitHubRepo, given a list of git remotes, will determine what the GitHub repo
// is.
func GitHubRepo(remotes []hub.Remote) string {
	// We only want to look at the `origin` remote.
	remote := findRemote("origin", remotes)
	if remote == nil {
		return ""
	}

	// Remotes that are not pointed at a GitHub repo are not valid.
	if remote.URL.Host != "github.com" {
		return ""
	}

	// Convert `/remind101/acme-inc.git` => `remind101/acme-inc`.
	return remoteRegex.ReplaceAllString(remote.URL.Path, "$1")
}

func findRemote(name string, remotes []hub.Remote) *hub.Remote {
	for _, r := range remotes {
		if r.Name == name {
			return &r
		}
	}

	return nil
}

var errInvalidRepo = errors.New("invalid repo")

// SplitRepo splits a repo string in the form remind101/acme-inc into it's owner
// and repo components.
func SplitRepo(nwo, defaultOrg string) (owner string, repo string, err error) {
	parts := strings.Split(nwo, "/")

	// If we were only given a repo name, and a default organization is set,
	// we'll use the defaultOrg as the owner.
	if len(parts) == 1 && defaultOrg != "" && parts[0] != "" {
		owner = defaultOrg
		repo = parts[0]
		return
	}

	if len(parts) != 2 {
		err = errInvalidRepo
		return
	}

	owner = parts[0]
	repo = parts[1]

	return
}

func askYN(prompt string) bool {
	r := bufio.NewReader(os.Stdin)
	fmt.Printf("%s (y/N)\n", prompt)
	a, _ := r.ReadString('\n')
	return strings.ToUpper(a) == "Y\n"
}
<commit_msg>Small refactor.<commit_after>package deploy

import (
	"bufio"
	"errors"
	"fmt"
	"io"
	"io/ioutil"
	"os"
	"regexp"
	"strings"
	"time"

	"github.com/github/hub/git"
	"github.com/remind101/deploy/Godeps/_workspace/src/github.com/codegangsta/cli"
	hub "github.com/remind101/deploy/Godeps/_workspace/src/github.com/github/hub/github"
	"github.com/remind101/deploy/Godeps/_workspace/src/github.com/google/go-github/github"
)

const (
	Name  = "deploy"
	Usage = "A command for creating GitHub deployments"
)

const (
	DefaultRef     = "master"
	DefaultTimeout = 20 * time.Second
)

var errTimeout = errors.New("Timed out waiting for build to start. Did you add a webhook to handle deployment events?")

func init() {
	cli.AppHelpTemplate = `USAGE:
   # Deploy the master branch of remind101/acme-inc to staging
   {{.Name}} --env=staging --ref=master remind101/acme-inc

   # Deploy HEAD of the current branch to staging
   {{.Name}} --env=staging remind101/acme-inc

   # Deploy the current GitHub repo to staging
   {{.Name}} --env=staging
{{if .Flags}}
OPTIONS:
   {{range .Flags}}{{.}}
   {{end}}{{end}}
`
}

var env string
var ProtectedEnvironments = map[string]bool{
	"production": true,
}

var flags = []cli.Flag{
	cli.StringFlag{
		Name:  "ref, branch, commit, tag",
		Value: "",
		Usage: "The git ref to deploy. Can be a git commit, branch or tag.",
	},
	cli.StringFlag{
		Name:        "env, e",
		Value:       "",
		Usage:       "The environment to deploy to.",
		Destination: &env,
	},
	cli.BoolFlag{
		Name:  "force, f",
		Usage: "Ignore commit status checks.",
	},
	cli.BoolFlag{
		Name:  "detached, d",
		Usage: "Don't wait for the deployment to complete.",
	},
	cli.BoolFlag{
		Name:  "quiet, q",
		Usage: "Silence any output to STDOUT.",
	},
	cli.BoolFlag{
		Name:  "update, u",
		Usage: "Update the binary",
	},
}

// NewApp returns a new cli.App for the deploy command.
func NewApp() *cli.App {
	app := cli.NewApp()
	app.Version = Version
	app.Name = Name
	app.Usage = Usage
	app.Flags = flags
	app.Action = func(c *cli.Context) {
		if c.Bool("update") {
			updater := NewUpdater()
			if err := updater.Update(); err != nil {
				fmt.Println(err)
				os.Exit(-1)
			} else {
				os.Exit(0)
			}
		}

		if err := RunDeploy(c); err != nil {
			msg := err.Error()
			if err, ok := err.(*github.ErrorResponse); ok {
				if strings.HasPrefix(err.Message, "Conflict: Commit status checks failed for") {
					msg = "Commit status checks failed. You can bypass commit status checks with the --force flag."
				} else if strings.HasPrefix(err.Message, "No ref found for") {
					msg = fmt.Sprintf("%s. Did you push it to GitHub?", err.Message)
				} else {
					msg = err.Message
				}
			}

			fmt.Println(msg)
			os.Exit(-1)
		}
	}

	return app
}

// RunDeploy performs a deploy.
func RunDeploy(c *cli.Context) error {
	var w io.Writer
	if c.Bool("quiet") {
		w = ioutil.Discard
	} else {
		w = c.App.Writer
	}

	h, err := hub.CurrentConfig().PromptForHost("github.com")
	if err != nil {
		return err
	}

	client, err := newGitHubClient(h)
	if err != nil {
		return err
	}

	nwo, err := Repo(c.Args())
	if err != nil {
		return err
	}

	owner, repo, err := SplitRepo(nwo, os.Getenv("GITHUB_ORGANIZATION"))
	if err != nil {
		return fmt.Errorf("Invalid GitHub repo: %s", nwo)
	}

	if env == "" {
		return fmt.Errorf("--env flag is required")
	}
	env = AliasEnvironment(env)

	ref := Ref(c.String("ref"), git.Head)
	err = displayNewCommits(owner, repo, ref, env, client)
	if err != nil {
		return err
	}

	r, err := newDeploymentRequest(c, ref, env)
	if err != nil {
		return err
	}

	fmt.Fprintf(w, "Deploying %s/%s@%s to %s...\n", owner, repo, *r.Ref, *r.Environment)

	if c.Bool("detached") {
		return nil
	}

	d, _, err := client.Repositories.CreateDeployment(owner, repo, r)
	if err != nil {
		return err
	}

	started := make(chan *github.DeploymentStatus)
	completed := make(chan *github.DeploymentStatus)

	go func() {
		started <- waitState(pendingStates, owner, repo, *d.ID, client)
	}()

	go func() {
		completed <- waitState(completedStates, owner, repo, *d.ID, client)
	}()

	select {
	case <-time.After(DefaultTimeout):
		return errTimeout
	case status := <-started:
		var url string
		if status.TargetURL != nil {
			url = *status.TargetURL
		}
		fmt.Fprintf(w, "%s\n", url)
	}

	status := <-completed

	if isFailed(*status.State) {
		return errors.New("Failed to deploy")
	}

	return nil
}

func displayNewCommits(owner string, repo string, ref string, env string, client *github.Client) error {
	opt := &github.DeploymentsListOptions{
		Environment: env,
	}

	deployments, _, err := client.Repositories.ListDeployments(owner, repo, opt)
	if err != nil || len(deployments) == 0 {
		return err
	}

	sha := *deployments[0].SHA
	compare, _, err := client.Repositories.CompareCommits(owner, repo, sha, ref)
	if err != nil || len(compare.Commits) == 0 {
		return err
	}

	fmt.Println("Deploying the following commits:\n")
	for _, commit := range compare.Commits {
		message := *commit.Commit.Message
		fmt.Printf("%-20s\t%s\n", *commit.Commit.Author.Name, strings.Split(message, "\n")[0])
	}
	fmt.Printf("\nSee entire diff here: https://github.com/%s/%s/compare/%s...%s\n\n", owner, repo, sha, ref)
	return nil
}

var EnvironmentAliases = map[string]string{
	"prod":  "production",
	"stage": "staging",
}

func AliasEnvironment(env string) string {
	if a, ok := EnvironmentAliases[env]; ok {
		return a
	}

	return env
}

func newDeploymentRequest(c *cli.Context, ref string, env string) (*github.DeploymentRequest, error) {
	if ProtectedEnvironments[env] {
		yes := askYN(fmt.Sprintf("Are you sure you want to deploy %s to %s?", ref, env))
		if !yes {
			return nil, fmt.Errorf("Deployment aborted.")
		}
	}

	var contexts *[]string
	if c.Bool("force") {
		s := []string{}
		contexts = &s
	}

	return &github.DeploymentRequest{
		Ref:              github.String(ref),
		Task:             github.String("deploy"),
		AutoMerge:        github.Bool(false),
		Environment:      github.String(env),
		RequiredContexts: contexts,
		Payload: map[string]interface{}{
			"force": c.Bool("force"),
		},
		// TODO Description:
	}, nil
}

var (
	pendingStates   = []string{"pending"}
	completedStates = []string{"success", "error", "failure"}
)

func isFailed(state string) bool {
	return state == "error" || state == "failure"
}

// waitState waits for a deployment status that matches the given states, then
// sends on the returned channel.
func waitState(states []string, owner, repo string, deploymentID int, c *github.Client) *github.DeploymentStatus {
	for {
		<-time.After(1 * time.Second)

		statuses, _, err := c.Repositories.ListDeploymentStatuses(owner, repo, deploymentID, nil)
		if err != nil {
			continue
		}

		status := firstStatus(states, statuses)
		if status != nil {
			return status
		}
	}
}

// firstStatus takes a slice of github.DeploymentStatus and returns the
// first status that matches the provided slice of states.
func firstStatus(states []string, statuses []github.DeploymentStatus) *github.DeploymentStatus {
	for _, ds := range statuses {
		for _, s := range states {
			if ds.State != nil && *ds.State == s {
				return &ds
			}
		}
	}

	return nil
}

// refRegex is a regular expression that matches a full git HEAD ref.
var refRegex = regexp.MustCompile(`^refs/heads/(.*)$`)

// Ref attempts to return the proper git ref to deploy. If a ref is provided,
// that will be returned. If not, it will fallback to calling headFunc. If an
// error is returned (not in a git repo), then it will fallback to DefaultRef.
func Ref(ref string, headFunc func() (string, error)) string {
	if ref != "" {
		return ref
	}

	ref, err := headFunc()
	if err != nil {
		// An error means that we're either not in a GitRepo or we're
		// not on a branch. In this case, we just fallback to the
		// DefaultRef.
		return DefaultRef
	}

	// Convert `refs/heads/test-deploy` => `test-deploy`
	return refRegex.ReplaceAllString(ref, "$1")
}

// Repo will determine the correct GitHub repo to deploy to, based on a set of
// arguments.
func Repo(arguments []string) (string, error) {
	if len(arguments) != 0 {
		return arguments[0], nil
	}

	remotes, err := hub.Remotes()
	if err != nil {
		return "", err
	}

	repo := GitHubRepo(remotes)
	if repo == "" {
		return repo, errors.New("no GitHub repo found in .git/config")
	}

	return repo, nil
}

// A regular expression that can convert a URL.Path into a GitHub repo name.
var remoteRegex = regexp.MustCompile(`^/(.*)\.git$`)

// GitHubRepo, given a list of git remotes, will determine what the GitHub repo
// is.
func GitHubRepo(remotes []hub.Remote) string {
	// We only want to look at the `origin` remote.
	remote := findRemote("origin", remotes)
	if remote == nil {
		return ""
	}

	// Remotes that are not pointed at a GitHub repo are not valid.
	if remote.URL.Host != "github.com" {
		return ""
	}

	// Convert `/remind101/acme-inc.git` => `remind101/acme-inc`.
	return remoteRegex.ReplaceAllString(remote.URL.Path, "$1")
}

func findRemote(name string, remotes []hub.Remote) *hub.Remote {
	for _, r := range remotes {
		if r.Name == name {
			return &r
		}
	}

	return nil
}

var errInvalidRepo = errors.New("invalid repo")

// SplitRepo splits a repo string in the form remind101/acme-inc into it's owner
// and repo components.
func SplitRepo(nwo, defaultOrg string) (owner string, repo string, err error) {
	parts := strings.Split(nwo, "/")

	// If we were only given a repo name, and a default organization is set,
	// we'll use the defaultOrg as the owner.
	if len(parts) == 1 && defaultOrg != "" && parts[0] != "" {
		owner = defaultOrg
		repo = parts[0]
		return
	}

	if len(parts) != 2 {
		err = errInvalidRepo
		return
	}

	owner = parts[0]
	repo = parts[1]

	return
}

func askYN(prompt string) bool {
	r := bufio.NewReader(os.Stdin)
	fmt.Printf("%s (y/N)\n", prompt)
	a, _ := r.ReadString('\n')
	return strings.ToUpper(a) == "Y\n"
}
<|endoftext|>


<commit_before>package main

import (
	"errors"
	"log"
	"os"
	"regexp"
	"strconv"
	"time"

	"github.com/coreos/go-etcd/etcd"
	"github.com/fsouza/go-dockerclient"
)

const (
	timeout time.Duration = 10 * time.Second
	etcdTTL time.Duration = timeout * 2
)

func getopt(name, dfault string) string {
	value := os.Getenv(name)
	if value == "" {
		value = dfault
	}
	return value
}

func main() {
	endpoint := getopt("DOCKER_HOST", "unix:///var/run/docker.sock")
	etcdHost := getopt("ETCD_HOST", "127.0.0.1")

	client, err := docker.NewClient(endpoint)
	if err != nil {
		log.Fatal(err)
	}

	etcdClient := etcd.NewClient([]string{"http://" + etcdHost + ":4001"})

	go listenContainers(client, etcdClient, etcdTTL)

	for {
		go pollContainers(client, etcdClient, etcdTTL)
		time.Sleep(timeout)
	}
}

func listenContainers(client *docker.Client, etcdClient *etcd.Client, ttl time.Duration) {

	listener := make(chan *docker.APIEvents)
	// TODO: figure out why we need to sleep for 10 milliseconds
	// https://github.com/fsouza/go-dockerclient/blob/0236a64c6c4bd563ec277ba00e370cc753e1677c/event_test.go#L43
	defer func() { time.Sleep(10 * time.Millisecond); client.RemoveEventListener(listener) }()
	err := client.AddEventListener(listener)
	if err != nil {
		log.Fatal(err)
	}
	for {
		select {
		case event := <-listener:
			if event.Status == "start" {
				container, err := getContainer(client, event.ID)
				if err != nil {
					log.Fatal(err)
				}
				publishContainer(etcdClient, container, ttl)
			}
		}
	}
}

func getContainer(client *docker.Client, id string) (*docker.APIContainers, error) {
	containers, err := client.ListContainers(docker.ListContainersOptions{})
	if err != nil {
		return nil, err
	}

	for _, container := range containers {
		// send container to channel for processing
		if container.ID == id {
			return &container, nil
		}
	}
	return nil, errors.New("could not find container")
}

func pollContainers(client *docker.Client, etcdClient *etcd.Client, ttl time.Duration) {
	containers, err := client.ListContainers(docker.ListContainersOptions{})
	if err != nil {
		log.Fatal(err)
	}

	for _, container := range containers {
		// send container to channel for processing
		publishContainer(etcdClient, &container, ttl)
	}
}

func publishContainer(client *etcd.Client, container *docker.APIContainers, ttl time.Duration) {

	var publishableContainerName = regexp.MustCompile(`[a-z0-9-]+_v[1-9][0-9]*.(cmd|web).[1-9][0-9]*`)
	var publishableContainerBaseName = regexp.MustCompile(`^[a-z0-9-]+`)

	// this is where we publish to etcd
	for _, name := range container.Names {
		// HACK: remove slash from container name
		// see https://github.com/docker/docker/issues/7519
		containerName := name[1:]
		if !publishableContainerName.MatchString(containerName) {
			continue
		}
		containerBaseName := publishableContainerBaseName.FindString(containerName)
		keyPath := "/deis/services/" + containerBaseName + "/" + containerName
		for _, p := range container.Ports {
			host := os.Getenv("HOST")
			port := strconv.Itoa(int(p.PublicPort))
			setEtcd(client, keyPath, host+":"+port, uint64(ttl.Seconds()))
			// TODO: support multiple exposed ports
			break
		}
	}
}

func setEtcd(client *etcd.Client, key, value string, ttl uint64) {
	_, err := client.Set(key, value, ttl)
	if err != nil {
		log.Println(err)
	}
	log.Println("set", key, "->", value)
}

func unsetEtcd(client *etcd.Client, key string) {
	_, err := client.Delete(key, true)
	if err != nil {
		log.Println(err)
	}
	log.Println("unset", key)
}
<commit_msg>fix(publisher): skip when failed to retrieve container<commit_after>package main

import (
	"errors"
	"log"
	"os"
	"regexp"
	"strconv"
	"time"

	"github.com/coreos/go-etcd/etcd"
	"github.com/fsouza/go-dockerclient"
)

const (
	timeout time.Duration = 10 * time.Second
	etcdTTL time.Duration = timeout * 2
)

func getopt(name, dfault string) string {
	value := os.Getenv(name)
	if value == "" {
		value = dfault
	}
	return value
}

func main() {
	endpoint := getopt("DOCKER_HOST", "unix:///var/run/docker.sock")
	etcdHost := getopt("ETCD_HOST", "127.0.0.1")

	client, err := docker.NewClient(endpoint)
	if err != nil {
		log.Fatal(err)
	}

	etcdClient := etcd.NewClient([]string{"http://" + etcdHost + ":4001"})

	go listenContainers(client, etcdClient, etcdTTL)

	for {
		go pollContainers(client, etcdClient, etcdTTL)
		time.Sleep(timeout)
	}
}

func listenContainers(client *docker.Client, etcdClient *etcd.Client, ttl time.Duration) {

	listener := make(chan *docker.APIEvents)
	// TODO: figure out why we need to sleep for 10 milliseconds
	// https://github.com/fsouza/go-dockerclient/blob/0236a64c6c4bd563ec277ba00e370cc753e1677c/event_test.go#L43
	defer func() { time.Sleep(10 * time.Millisecond); client.RemoveEventListener(listener) }()
	err := client.AddEventListener(listener)
	if err != nil {
		log.Fatal(err)
	}
	for {
		select {
		case event := <-listener:
			if event.Status == "start" {
				container, err := getContainer(client, event.ID)
				if err != nil {
					log.Println(err)
					continue
				}
				publishContainer(etcdClient, container, ttl)
			}
		}
	}
}

func getContainer(client *docker.Client, id string) (*docker.APIContainers, error) {
	containers, err := client.ListContainers(docker.ListContainersOptions{})
	if err != nil {
		return nil, err
	}

	for _, container := range containers {
		// send container to channel for processing
		if container.ID == id {
			return &container, nil
		}
	}
	return nil, errors.New("could not find container")
}

func pollContainers(client *docker.Client, etcdClient *etcd.Client, ttl time.Duration) {
	containers, err := client.ListContainers(docker.ListContainersOptions{})
	if err != nil {
		log.Fatal(err)
	}

	for _, container := range containers {
		// send container to channel for processing
		publishContainer(etcdClient, &container, ttl)
	}
}

func publishContainer(client *etcd.Client, container *docker.APIContainers, ttl time.Duration) {

	var publishableContainerName = regexp.MustCompile(`[a-z0-9-]+_v[1-9][0-9]*.(cmd|web).[1-9][0-9]*`)
	var publishableContainerBaseName = regexp.MustCompile(`^[a-z0-9-]+`)

	// this is where we publish to etcd
	for _, name := range container.Names {
		// HACK: remove slash from container name
		// see https://github.com/docker/docker/issues/7519
		containerName := name[1:]
		if !publishableContainerName.MatchString(containerName) {
			continue
		}
		containerBaseName := publishableContainerBaseName.FindString(containerName)
		keyPath := "/deis/services/" + containerBaseName + "/" + containerName
		for _, p := range container.Ports {
			host := os.Getenv("HOST")
			port := strconv.Itoa(int(p.PublicPort))
			setEtcd(client, keyPath, host+":"+port, uint64(ttl.Seconds()))
			// TODO: support multiple exposed ports
			break
		}
	}
}

func setEtcd(client *etcd.Client, key, value string, ttl uint64) {
	_, err := client.Set(key, value, ttl)
	if err != nil {
		log.Println(err)
	}
	log.Println("set", key, "->", value)
}

func unsetEtcd(client *etcd.Client, key string) {
	_, err := client.Delete(key, true)
	if err != nil {
		log.Println(err)
	}
	log.Println("unset", key)
}
<|endoftext|>


<commit_before>package hashtree

import (
	"testing"
)

type treeLevels struct {
	len    Bytes
	levels Level
}

var expectedTreeLevels = []treeLevels{
	{0, 1}, {1, 1}, {treeNodeSize, 1},
	{treeNodeSize + 1, 2}, {2 * treeNodeSize, 2},
	{2*treeNodeSize + 1, 3}, {4 * treeNodeSize, 3},
	{4*treeNodeSize + 1, 4}, {8 * treeNodeSize, 4},
	{8*treeNodeSize + 1, 5}, {16 * treeNodeSize, 5},
}

func TestTreeLevels(t *testing.T) {
	h := NewTree()
	for i := 0; i < len(expectedTreeLevels); i++ {
		e := expectedTreeLevels[i]
		if h.Levels(e.len) != e.levels {
			t.Fatalf("Levels(%d) = %d want %d", e.len, h.Levels(e.len), e.levels)
		}
	}
}

type levelWidth struct {
	len   Bytes
	level Level
	width Nodes
}

var expectedLevelWidth = []levelWidth{
	{0, 0, 1}, {1, 0, 1}, {treeNodeSize, 0, 1},

	{treeNodeSize + 1, 0, 2}, {2 * treeNodeSize, 0, 2},
	{treeNodeSize + 1, 1, 1}, {2 * treeNodeSize, 1, 1},

	{2*treeNodeSize + 1, 0, 3}, {4 * treeNodeSize, 0, 4},
	{2*treeNodeSize + 1, 1, 2}, {4 * treeNodeSize, 1, 2},
	{2*treeNodeSize + 1, 2, 1}, {4 * treeNodeSize, 2, 1},

	{5 * treeNodeSize, 0, 5}, {6 * treeNodeSize, 0, 6}, {7 * treeNodeSize, 0, 7}, {8 * treeNodeSize, 0, 8},
	{5 * treeNodeSize, 1, 3}, {6 * treeNodeSize, 1, 3}, {7 * treeNodeSize, 1, 4}, {8 * treeNodeSize, 1, 4},
	{5 * treeNodeSize, 2, 2}, {6 * treeNodeSize, 2, 2}, {7 * treeNodeSize, 2, 2}, {8 * treeNodeSize, 2, 2},
	{5 * treeNodeSize, 3, 1}, {6 * treeNodeSize, 3, 1}, {7 * treeNodeSize, 3, 1}, {8 * treeNodeSize, 3, 1},
}

func TestLevelWidth(t *testing.T) {
	h := NewTree()
	for i := 0; i < len(expectedLevelWidth); i++ {
		e := expectedLevelWidth[i]
		if h.LevelWidth(e.len, e.level) != e.width {
			t.Fatalf("LevelWidth(%d,%d) = %d want %d", e.len, e.level, h.LevelWidth(e.len, e.level), e.width)
		}
	}
}

var expectedInner = [][]int32{
	{1, 1, 2, 3, 5, 8},
	{0, -1, -3},
	{1, -3},
	{4},
}

func TestInnerHashListener(t *testing.T) {
	listener := func(level Level, index Nodes, hash *H256) {
		//TODO: check hash
	}
	c := NewTree2(NoPad32bytes, minus).(*treeDigest)
	c.SetInnerHashListener(listener)
	for _, n := range expectedInner[0] {
		data := H256{uint32(n)}
		c.Write(data.toBytes())
	}
}
<commit_msg>test calls on inner hash listener<commit_after>package hashtree

import (
	"testing"
)

type treeLevels struct {
	len    Bytes
	levels Level
}

var expectedTreeLevels = []treeLevels{
	{0, 1}, {1, 1}, {treeNodeSize, 1},
	{treeNodeSize + 1, 2}, {2 * treeNodeSize, 2},
	{2*treeNodeSize + 1, 3}, {4 * treeNodeSize, 3},
	{4*treeNodeSize + 1, 4}, {8 * treeNodeSize, 4},
	{8*treeNodeSize + 1, 5}, {16 * treeNodeSize, 5},
}

func TestTreeLevels(t *testing.T) {
	h := NewTree()
	for i := 0; i < len(expectedTreeLevels); i++ {
		e := expectedTreeLevels[i]
		if h.Levels(e.len) != e.levels {
			t.Fatalf("Levels(%d) = %d want %d", e.len, h.Levels(e.len), e.levels)
		}
	}
}

type levelWidth struct {
	len   Bytes
	level Level
	width Nodes
}

var expectedLevelWidth = []levelWidth{
	{0, 0, 1}, {1, 0, 1}, {treeNodeSize, 0, 1},

	{treeNodeSize + 1, 0, 2}, {2 * treeNodeSize, 0, 2},
	{treeNodeSize + 1, 1, 1}, {2 * treeNodeSize, 1, 1},

	{2*treeNodeSize + 1, 0, 3}, {4 * treeNodeSize, 0, 4},
	{2*treeNodeSize + 1, 1, 2}, {4 * treeNodeSize, 1, 2},
	{2*treeNodeSize + 1, 2, 1}, {4 * treeNodeSize, 2, 1},

	{5 * treeNodeSize, 0, 5}, {6 * treeNodeSize, 0, 6}, {7 * treeNodeSize, 0, 7}, {8 * treeNodeSize, 0, 8},
	{5 * treeNodeSize, 1, 3}, {6 * treeNodeSize, 1, 3}, {7 * treeNodeSize, 1, 4}, {8 * treeNodeSize, 1, 4},
	{5 * treeNodeSize, 2, 2}, {6 * treeNodeSize, 2, 2}, {7 * treeNodeSize, 2, 2}, {8 * treeNodeSize, 2, 2},
	{5 * treeNodeSize, 3, 1}, {6 * treeNodeSize, 3, 1}, {7 * treeNodeSize, 3, 1}, {8 * treeNodeSize, 3, 1},
}

func TestLevelWidth(t *testing.T) {
	h := NewTree()
	for i := 0; i < len(expectedLevelWidth); i++ {
		e := expectedLevelWidth[i]
		if h.LevelWidth(e.len, e.level) != e.width {
			t.Fatalf("LevelWidth(%d,%d) = %d want %d", e.len, e.level, h.LevelWidth(e.len, e.level), e.width)
		}
	}
}

func TestInnerHashListener(t *testing.T) {
	inner := [][]int32{
		{1, 1, 2, 3, 5, 8},
		{0, -1, -3},
		{1, -3},
		{4},
	}
	listener := func(l Level, i Nodes, hash *H256) {
		h := int32(hash[0])
		if inner[l][i] != h {
			if inner[l][i] == h+2000 {
				t.Fatalf("Level:%d, Node:%d was repeated", l, i)
			}
			t.Fatalf("Level:%d, Node:%d, hash:%d, should be %d", l, i, h, inner[l][i])
		} else {
			inner[l][i] += 2000 //mark heard
		}
	}
	c := NewTree2(NoPad32bytes, minus).(*treeDigest)
	c.SetInnerHashListener(listener)
	for _, n := range inner[0] {
		data := H256{uint32(n)}
		c.Write(data.toBytes())
	}
	c.Sum(nil)

	for l, array := range inner {
		for i, n := range array {
			if n < 1000 {
				t.Fatalf("Level:%d, Node:%d was not heard", l, i)
			}
		}
	}
}
<|endoftext|>


<commit_before>/*
   Hockeypuck - OpenPGP key server
   Copyright (C) 2012  Casey Marshall

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

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

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

package sks

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

	"gopkg.in/errgo.v1"
	"hockeypuck/hkp/storage"
)

type LoadStat struct {
	Inserted int
	Updated  int
}

type LoadStatMap map[time.Time]*LoadStat

func (m LoadStatMap) MarshalJSON() ([]byte, error) {
	doc := map[string]*LoadStat{}
	for k, v := range m {
		doc[k.Format(time.RFC3339)] = v
	}
	return json.Marshal(&doc)
}

func (m LoadStatMap) UnmarshalJSON(b []byte) error {
	doc := map[string]*LoadStat{}
	err := json.Unmarshal(b, &doc)
	if err != nil {
		return err
	}
	for k, v := range doc {
		t, err := time.Parse(time.RFC3339, k)
		if err != nil {
			return err
		}
		m[t] = v
	}
	return nil
}

func (m LoadStatMap) update(t time.Time, kc storage.KeyChange) {
	ls, ok := m[t]
	if !ok {
		ls = &LoadStat{}
		m[t] = ls
	}
	switch kc.(type) {
	case storage.KeyAdded:
		ls.Inserted++
	case storage.KeyReplaced:
		ls.Updated++
	}
}

type Stats struct {
	Total int

	mu     sync.Mutex
	Hourly LoadStatMap
	Daily  LoadStatMap
}

func NewStats() *Stats {
	return &Stats{
		Hourly: LoadStatMap{},
		Daily:  LoadStatMap{},
	}
}

func (s *Stats) prune() {
	yesterday := time.Now().UTC().Add(-24 * time.Hour)
	lastWeek := time.Now().UTC().Add(-24 * 7 * time.Hour)
	s.mu.Lock()
	for k := range s.Hourly {
		if k.Before(yesterday) {
			delete(s.Hourly, k)
		}
	}
	for k := range s.Daily {
		if k.Before(lastWeek) {
			delete(s.Daily, k)
		}
	}
	s.mu.Unlock()
}

func (s *Stats) Update(kc storage.KeyChange) {
	s.mu.Lock()
	s.Hourly.update(time.Now().UTC().Truncate(time.Hour), kc)
	s.Daily.update(time.Now().UTC().Truncate(24*time.Hour), kc)
	switch kc.(type) {
	case storage.KeyAdded:
		s.Total++
	}
	s.mu.Unlock()
}

func (s *Stats) clone() *Stats {
	s.mu.Lock()
	result := &Stats{
		Total:  s.Total,
		Hourly: LoadStatMap{},
		Daily:  LoadStatMap{},
	}
	for k, v := range s.Hourly {
		result.Hourly[k] = v
	}
	for k, v := range s.Daily {
		result.Daily[k] = v
	}
	s.mu.Unlock()
	return result
}

func (s *Stats) ReadFile(path string) error {
	f, err := os.Open(path)
	if err != nil {
		if os.IsNotExist(err) {
			empty := NewStats()
			*s = *empty
			return nil
		} else {
			return errgo.Notef(err, "cannot open stats %q", path)
		}
	} else {
		defer f.Close()
		err = json.NewDecoder(f).Decode(s)
		if err != nil {
			return errgo.Notef(err, "cannot decode stats")
		}
	}
	return nil
}

func (s *Stats) WriteFile(path string) error {
	f, err := os.Create(path)
	if err != nil {
		return errgo.Notef(err, "cannot open stats %q", path)
	}
	defer f.Close()
	err = json.NewEncoder(f).Encode(s)
	if err != nil {
		return errgo.Notef(err, "cannot encode stats")
	}
	return nil
}
<commit_msg>Fix a bug in Stats.ReadFile<commit_after>/*
   Hockeypuck - OpenPGP key server
   Copyright (C) 2012  Casey Marshall

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

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

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

package sks

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

	"gopkg.in/errgo.v1"
	"hockeypuck/hkp/storage"
)

type LoadStat struct {
	Inserted int
	Updated  int
}

type LoadStatMap map[time.Time]*LoadStat

func (m LoadStatMap) MarshalJSON() ([]byte, error) {
	doc := map[string]*LoadStat{}
	for k, v := range m {
		doc[k.Format(time.RFC3339)] = v
	}
	return json.Marshal(&doc)
}

func (m LoadStatMap) UnmarshalJSON(b []byte) error {
	doc := map[string]*LoadStat{}
	err := json.Unmarshal(b, &doc)
	if err != nil {
		return err
	}
	for k, v := range doc {
		t, err := time.Parse(time.RFC3339, k)
		if err != nil {
			return err
		}
		m[t] = v
	}
	return nil
}

func (m LoadStatMap) update(t time.Time, kc storage.KeyChange) {
	ls, ok := m[t]
	if !ok {
		ls = &LoadStat{}
		m[t] = ls
	}
	switch kc.(type) {
	case storage.KeyAdded:
		ls.Inserted++
	case storage.KeyReplaced:
		ls.Updated++
	}
}

type Stats struct {
	Total int

	mu     sync.Mutex
	Hourly LoadStatMap
	Daily  LoadStatMap
}

func NewStats() *Stats {
	return &Stats{
		Hourly: LoadStatMap{},
		Daily:  LoadStatMap{},
	}
}

func (s *Stats) reset() {
	s.mu.Lock()
	defer s.mu.Unlock()
	s.Total = 0
	s.Hourly = LoadStatMap{}
	s.Daily = LoadStatMap{}
}

func (s *Stats) prune() {
	yesterday := time.Now().UTC().Add(-24 * time.Hour)
	lastWeek := time.Now().UTC().Add(-24 * 7 * time.Hour)
	s.mu.Lock()
	for k := range s.Hourly {
		if k.Before(yesterday) {
			delete(s.Hourly, k)
		}
	}
	for k := range s.Daily {
		if k.Before(lastWeek) {
			delete(s.Daily, k)
		}
	}
	s.mu.Unlock()
}

func (s *Stats) Update(kc storage.KeyChange) {
	s.mu.Lock()
	s.Hourly.update(time.Now().UTC().Truncate(time.Hour), kc)
	s.Daily.update(time.Now().UTC().Truncate(24*time.Hour), kc)
	switch kc.(type) {
	case storage.KeyAdded:
		s.Total++
	}
	s.mu.Unlock()
}

func (s *Stats) clone() *Stats {
	s.mu.Lock()
	result := &Stats{
		Total:  s.Total,
		Hourly: LoadStatMap{},
		Daily:  LoadStatMap{},
	}
	for k, v := range s.Hourly {
		result.Hourly[k] = v
	}
	for k, v := range s.Daily {
		result.Daily[k] = v
	}
	s.mu.Unlock()
	return result
}

func (s *Stats) ReadFile(path string) error {
	f, err := os.Open(path)
	if err != nil {
		if os.IsNotExist(err) {
			s.reset()
			return nil
		} else {
			return errgo.Notef(err, "cannot open stats %q", path)
		}
	} else {
		defer f.Close()
		// TODO(jsing): This is modifying the maps without holding mu.
		err = json.NewDecoder(f).Decode(s)
		if err != nil {
			return errgo.Notef(err, "cannot decode stats")
		}
	}
	return nil
}

func (s *Stats) WriteFile(path string) error {
	f, err := os.Create(path)
	if err != nil {
		return errgo.Notef(err, "cannot open stats %q", path)
	}
	defer f.Close()
	// TODO(jsing): This is reading the maps without holding mu.
	err = json.NewEncoder(f).Encode(s)
	if err != nil {
		return errgo.Notef(err, "cannot encode stats")
	}
	return nil
}
<|endoftext|>


<commit_before>// Host detection.

package embd

import (
	"fmt"
	"io/ioutil"
	"os/exec"
	"strconv"
	"strings"
)

// The Host type represents all the supported host types.
type Host string

const (
	// HostNull reprents a null host.
	HostNull Host = ""

	// HostRPi represents the RaspberryPi.
	HostRPi = "Raspberry Pi"

	// HostBBB represents the BeagleBone Black.
	HostBBB = "BeagleBone Black"

	// HostGalileo represents the Intel Galileo board.
	HostGalileo = "Intel Galileo"

	// HostCubieTruck represents the Cubie Truck.
	HostCubieTruck = "CubieTruck"

	// HostRadxa represents the Radxa board.
	HostRadxa = "Radxa"

	// HostCHIP represents the NextThing C.H.I.P.
	HostCHIP = "CHIP"
)

func execOutput(name string, arg ...string) (output string, err error) {
	var out []byte
	if out, err = exec.Command(name, arg...).Output(); err != nil {
		return
	}
	output = strings.TrimSpace(string(out))
	return
}

func parseVersion(str string) (major, minor, patch int, err error) {
	versionNumber := strings.Split(str, "-")
	parts := strings.Split(versionNumber[0], ".")
	len := len(parts)

	if major, err = strconv.Atoi(parts[0]); err != nil {
		return 0, 0, 0, err
	}
	if minor, err = strconv.Atoi(parts[1]); err != nil {
		return 0, 0, 0, err
	}
	if len > 2 {
		part := parts[2]
		part = strings.TrimSuffix(part, "+")
		if patch, err = strconv.Atoi(part); err != nil {
			return 0, 0, 0, err
		}
	}

	return major, minor, patch, err
}

func kernelVersion() (major, minor, patch int, err error) {
	output, err := execOutput("uname", "-r")
	if err != nil {
		return 0, 0, 0, err
	}

	return parseVersion(output)
}

func cpuInfo() (model, hardware string, revision int, err error) {
	output, err := ioutil.ReadFile("/proc/cpuinfo")
	if err != nil {
		return "", "", 0, err
	}
	for _, line := range strings.Split(string(output), "\n") {
		fields := strings.Split(line, ":")
		if len(fields) < 1 {
			continue
		}
		switch {
		case strings.HasPrefix(fields[0], "Revision"):
			revs := strings.TrimSpace(fields[1])
			rev, err := strconv.ParseInt(revs, 16, 32)
			if err != nil {
				continue
			}
			revision = int(rev)
		case strings.HasPrefix(fields[0], "Hardware"):
			hardware = strings.TrimSpace(fields[1])
		case strings.HasPrefix(fields[0], "model name"):
			model = fields[1]
		}
	}
	return model, hardware, revision, nil
}

// DetectHost returns the detected host and its revision number.
func DetectHost() (host Host, rev int, err error) {
	major, minor, patch, err := kernelVersion()
	if err != nil {
		return HostNull, 0, err
	}

	if major < 3 || (major == 3 && minor < 8) {
		return HostNull, 0, fmt.Errorf(
			"embd: linux kernel versions lower than 3.8 are not supported, "+
				"you have %v.%v.%v", major, minor, patch)
	}

	model, hardware, rev, err := cpuInfo()
	if err != nil {
		return HostNull, 0, err
	}

	switch {
	case strings.Contains(model, "ARMv7") && (strings.Contains(hardware, "AM33XX") || strings.Contains(hardware, "AM335X")):
		return HostBBB, rev, nil
	case strings.Contains(hardware, "BCM2708") || strings.Contains(hardware, "BCM2709"):
		return HostRPi, rev, nil
	case hardware == "Allwinner sun4i/sun5i Families":
		if major < 4 || (major == 4 && minor < 4) {
			return HostNull, 0, fmt.Errorf(
				"embd: linux kernel version 4.4+ required, you have %v.%v",
				major, minor)
		}
		return HostCHIP, rev, nil
	default:
		return HostNull, 0, fmt.Errorf(`embd: your host "%v:%v" is not supported at this moment. request support at https://github.com/kidoman/embd/issues`, host, model)
	}
}
<commit_msg>Correctly detect Raspberry Pi 3 on Kernel 4.9<commit_after>// Host detection.

package embd

import (
	"fmt"
	"io/ioutil"
	"os/exec"
	"strconv"
	"strings"
)

// The Host type represents all the supported host types.
type Host string

const (
	// HostNull reprents a null host.
	HostNull Host = ""

	// HostRPi represents the RaspberryPi.
	HostRPi = "Raspberry Pi"

	// HostBBB represents the BeagleBone Black.
	HostBBB = "BeagleBone Black"

	// HostGalileo represents the Intel Galileo board.
	HostGalileo = "Intel Galileo"

	// HostCubieTruck represents the Cubie Truck.
	HostCubieTruck = "CubieTruck"

	// HostRadxa represents the Radxa board.
	HostRadxa = "Radxa"

	// HostCHIP represents the NextThing C.H.I.P.
	HostCHIP = "CHIP"
)

func execOutput(name string, arg ...string) (output string, err error) {
	var out []byte
	if out, err = exec.Command(name, arg...).Output(); err != nil {
		return
	}
	output = strings.TrimSpace(string(out))
	return
}

func parseVersion(str string) (major, minor, patch int, err error) {
	versionNumber := strings.Split(str, "-")
	parts := strings.Split(versionNumber[0], ".")
	len := len(parts)

	if major, err = strconv.Atoi(parts[0]); err != nil {
		return 0, 0, 0, err
	}
	if minor, err = strconv.Atoi(parts[1]); err != nil {
		return 0, 0, 0, err
	}
	if len > 2 {
		part := parts[2]
		part = strings.TrimSuffix(part, "+")
		if patch, err = strconv.Atoi(part); err != nil {
			return 0, 0, 0, err
		}
	}

	return major, minor, patch, err
}

func kernelVersion() (major, minor, patch int, err error) {
	output, err := execOutput("uname", "-r")
	if err != nil {
		return 0, 0, 0, err
	}

	return parseVersion(output)
}

func cpuInfo() (model, hardware string, revision int, err error) {
	output, err := ioutil.ReadFile("/proc/cpuinfo")
	if err != nil {
		return "", "", 0, err
	}
	for _, line := range strings.Split(string(output), "\n") {
		fields := strings.Split(line, ":")
		if len(fields) < 1 {
			continue
		}
		switch {
		case strings.HasPrefix(fields[0], "Revision"):
			revs := strings.TrimSpace(fields[1])
			rev, err := strconv.ParseInt(revs, 16, 32)
			if err != nil {
				continue
			}
			revision = int(rev)
		case strings.HasPrefix(fields[0], "Hardware"):
			hardware = strings.TrimSpace(fields[1])
		case strings.HasPrefix(fields[0], "model name"):
			model = fields[1]
		}
	}
	return model, hardware, revision, nil
}

// DetectHost returns the detected host and its revision number.
func DetectHost() (host Host, rev int, err error) {
	major, minor, patch, err := kernelVersion()
	if err != nil {
		return HostNull, 0, err
	}

	if major < 3 || (major == 3 && minor < 8) {
		return HostNull, 0, fmt.Errorf(
			"embd: linux kernel versions lower than 3.8 are not supported, "+
				"you have %v.%v.%v", major, minor, patch)
	}

	model, hardware, rev, err := cpuInfo()
	if err != nil {
		return HostNull, 0, err
	}

	switch {
	case strings.Contains(model, "ARMv7") && (strings.Contains(hardware, "AM33XX") || strings.Contains(hardware, "AM335X")):
		return HostBBB, rev, nil
	case strings.Contains(hardware, "BCM2708") || strings.Contains(hardware, "BCM2709") || strings.Contains(hardware, "BCM2835"):
		return HostRPi, rev, nil
	case hardware == "Allwinner sun4i/sun5i Families":
		if major < 4 || (major == 4 && minor < 4) {
			return HostNull, 0, fmt.Errorf(
				"embd: linux kernel version 4.4+ required, you have %v.%v",
				major, minor)
		}
		return HostCHIP, rev, nil
	default:
		return HostNull, 0, fmt.Errorf(`embd: your host "%v:%v" is not supported at this moment. request support at https://github.com/kidoman/embd/issues`, host, model)
	}
}
<|endoftext|>