<commit_before>package main

import (
  "strings"
  "log"
  // "fmt"
//  "regexp"
  "bytes"
  "encoding/json"
  "os"
  "time"
  "strconv"
  "regexp"
  "github.com/lookify/town/cluster"
  dockerapi "github.com/fsouza/go-dockerclient"
)

var (
  scaleNumRegexp, _ = regexp.Compile("\\$\\{SCALE_NUM:(.+)\\}")
//  SCALE_TOTAL_REG, _ = regexp.Compile("\\$\\{SCALE_NUM:(.+)\\}")
  hostsRegexp, _ = regexp.Compile("\\$\\{(.+)_HOSTS\\}")
)

// Town describe cluster and docker clients.
type Town struct {
  cluster *cluster.Cluster
  docker *dockerapi.Client // TODO change to multiple clients
}

// NewTown create new town with default values
func NewTown() *Town {
  return &Town{
    cluster: nil,
    docker: nil,
  }
}

// ReadFile - read town configuration fail in current direcotry or /etc/town (ext .yml)
func (t *Town) ReadFile(name string) {
  var pathLocs = [...]string{
    name + ".yml",
    "/etc/town/" + name + ".yml",
  }

  for _, path := range pathLocs {
    if _, err := os.Stat(path); err == nil {
      t.cluster = cluster.NewCluster(path)
      t.cluster.ReadFile()
      return
    }
  }

  log.Println("ERROR: Could not find file ", name, ".yml")
}

// Connect - connect to docker hosts.
func (t *Town) Connect() {
  endpoint := t.cluster.Application.Docker.Hosts[0] // at the moment use only first
  log.Println("Using Docker API endpont: ", endpoint)
  docker, err := dockerapi.NewClient( endpoint )
  if err != nil {
    log.Println("Can't connect to the docker")
  }
  t.docker = docker
}

// Provision running containers.
func (t *Town) Provision(checkChanged bool) {
  // update containers
  pull := true
  repository := t.cluster.Application.Docker.Repository
  if len(repository) == 1 && repository[0] == "local" {
    pull = false
  }

  if pull {
    for _, node := range t.cluster.Nodes {
      var buf bytes.Buffer
      var image = strings.Split(node.Container.Image, ":")

      opts := dockerapi.PullImageOptions{
          Repository: image[0],
          // Registry:     "docker.tsuru.io",
          // Tag:          "latest",
          OutputStream: &buf,
      }

      if len(image) > 1 {
        opts.Tag = image[1]
      }
      err := t.docker.PullImage(opts, dockerapi.AuthConfiguration{});
      if err != nil {
        log.Println("Could not pull image ", image)
      }
    }
  }

  allContainers, err := t.docker.ListContainers(dockerapi.ListContainersOptions{
    All: true,
  })
  if err == nil {
    for _, listing := range allContainers {
      container, err := t.docker.InspectContainer(listing.ID)
      if err == nil {
        name := container.Name[1:]
        node, index := t.cluster.FindNodeByName(name)
        if node != nil && index > 0 {
          if node.Container.Exist == nil {
            node.Container.Exist = []cluster.ExistContainer{}
          }
          runningContainer := cluster.NewExistContainer(listing.ID, name, index, container.State.Running)
          runningContainer.Pid = container.State.Pid
          runningContainer.User = container.Config.User
          if checkChanged {
            node.Container.Changed = t.isChangedImage(node, container)
          } else {
            node.Container.Changed = true
          }
          node.Container.Exist = append(node.Container.Exist, runningContainer)
        }
      } else {
        log.Println("[ERROR] Unable to inspect container:", listing.ID[:12], err)
      }
    }

    if checkChanged {
      t.cluster.AddChangeDependant()
    }
  } else {
    log.Println("[ERROR] Can't start provision")
  }
}

// Info - print current cluster information.
func (t *Town) Info() {
  for i := len(t.cluster.Nodes) - 1; i >= 0; i-- {
    node := t.cluster.Nodes[i]
    log.Print("Node ", node.Container.Name, " image ")
	  if node.Container.Changed {
      log.Print("(Changed)")
    }
	  log.Println(": ", node.Container.Image)
    for _, container := range node.Container.Exist {
      log.Print("      ", container.Name, "\t")
  	  if container.Running {
  	    log.Println("Running")
  	  } else {
  	    log.Print("Stoped")
  	  }
    }
  }
}

/**
 * Check node and running container for changes.
 * TODO: add cache to docker call.
 **/
func (t *Town) isChangedImage(node *cluster.Node, container *dockerapi.Container) bool {
  var imageName = container.Image
  image , error := t.docker.InspectImage(imageName)
  if error == nil {
    secondImage , secondError := t.docker.InspectImage(node.Container.Image)
    if secondError == nil {
      return secondImage.Created.After(image.Created)
    }
  }
  log.Println("[ERROR] Could not inspect image ", node.Container.Name)
  return false
}

// StopContainers - stop all containers or only containers with changed images.
func (t *Town) StopContainers(checkChanged bool) {
  log.Println("Stop...")
  //for node := range t.cluster.nodes {
  for i := len(t.cluster.Nodes) - 1; i >= 0; i-- {
    node := t.cluster.Nodes[i]
    if (!checkChanged || node.Container.Changed) && len(node.Container.Exist) > 0 {
      for _, container := range node.Container.Exist {
        if container.Running {
          err := t.docker.StopContainer(container.ID, 10)
          if err == nil {
            log.Println("   -  ", container.Name)
          } else {
            log.Println("   -  ", container.Name, " failed ", err)
          }
        }
      }
    }
  }
  log.Println("=============================")
}

// RemoveContainers - remove container form local repository.
func (t *Town) RemoveContainers(checkChanged bool) {
  log.Println("Remove...")
  //for node := range t.cluster.nodes {
  for i := len(t.cluster.Nodes) - 1; i >= 0; i-- {
    node := t.cluster.Nodes[i]
    if (!checkChanged || node.Container.Changed) && len(node.Container.Exist) > 0 {
      for _, container := range node.Container.Exist {
        err := t.docker.RemoveContainer(dockerapi.RemoveContainerOptions{
          ID: container.ID,
          RemoveVolumes: false,
        })
        if err == nil {
          log.Println("   -  ", container.Name)
        } else {
          log.Println("   -  ", container.Name, " failed ", err)
        }
      }
    }
  }
  log.Println("=============================")
}

// CreateContainer - create container.
func (t *Town) CreateContainer(node *cluster.Node, index int) (string, string, string) {
  containerName := node.Container.Name + "-" + strconv.Itoa(index)

  log.Println("   -  ", containerName)

  node.Container.Hostname = containerName // ?? Help !!!!

  env :=  make([]string, 0, cap(node.Container.Environment))
  for _, e := range node.Container.Environment {
    env = append(env, t.exec(e, index))
  }

  volumes := make(map[string]struct{})
  binds := make([]string, 0, cap(node.Container.Volumes))
  if len(node.Container.Volumes) > 0 {
    for _, volume := range node.Container.Volumes {
      volume = t.exec(volume, index)
      vol := strings.Split(volume, ":")
      if len(vol) > 1 {
        volumes[vol[1]] = struct{}{}
      } else {
        volumes[vol[0]] = struct{}{}
      }
      binds = append(binds, volume)
    }
  }

  dockerConfig := dockerapi.Config{
      Image: node.Container.Image,
      Hostname: node.Container.Hostname,
      PortSpecs: node.Container.Ports,
      Env: env,
      Volumes: volumes,

      AttachStdout: false,
      AttachStdin: false,
      AttachStderr: false,

      Tty: false,

      //Cmd: []
  }

  if len(node.Container.Command) > 0 {
    cmd := t.exec(node.Container.Command, index)
    dockerConfig.Cmd = []string{ cmd }
  }

  // just info
  //for _, l := range node.status.links {
  //  log.Println("     * ", l)
  //}

  // create links
  links := t.cluster.GetLinks(node)

  portBindings := map[dockerapi.Port][]dockerapi.PortBinding{}
  // create ports
  for _, ports := range node.Container.Ports {

    port := strings.Split(ports, ":")
    var p dockerapi.Port

    if len(port) > 1 {
      p = dockerapi.Port(port[1] + "/tcp")
    } else {
      p = dockerapi.Port(port[0] + "/tcp")
    }

    if portBindings[p] == nil {
      portBindings[p] = [] dockerapi.PortBinding {}
    }

    portBindings[p] = append(portBindings[p], dockerapi.PortBinding{
      HostIP: "",
      HostPort: port[0],
    })
  }

  hostConfig := dockerapi.HostConfig{
    Binds: binds,
    Links: links, //, [],
    PortBindings: portBindings,
    NetworkMode: "bridge",
    PublishAllPorts: false,
    Privileged: node.Container.Privileged,
  }

  opts := dockerapi.CreateContainerOptions{Name: containerName, Config: &dockerConfig, HostConfig: &hostConfig}
  container, err := t.docker.CreateContainer(opts)
  if err == nil {
    runningContainer := cluster.NewExistContainer(container.ID, containerName, index, true)
    node.Container.Exist = append(node.Container.Exist, runningContainer)
    // runningContainer.Pid = container.State.Pid
    // runningContainer.User = container.Config.User
    // if checkChanged {
    //   node.Container.Changed = t.isChangedImage(node, container)
    // } else {
    //   node.Container.Changed = true
    // }

    retry := 5
    for retry > 0 {
      error := t.docker.StartContainer(container.ID, &hostConfig)
      if error != nil {
        // log.Println("start error: ", error);

        out, err := json.Marshal(container)
        if err != nil {
            panic (err)
        }
        // fmt.Println(string(out))

        retry--;
        if retry == 0 {
          log.Println(" Start failed after 5 retries: ", string(out))
        }
        // log.Println("retry: ", retry);
      } else {
        inspect, inspectError := t.docker.InspectContainer(container.ID)
        if inspectError == nil {
          //links = append(links, inspect.NetworkSettings.IPAddress + "  " + containerName)
          //ids = append(ids, container.ID)
          return container.ID, inspect.NetworkSettings.IPAddress + "  " + containerName, containerName
        }

        log.Println("Inpect ", container.ID, " error ", inspectError)

        //retry = 0
        break;
      }
    }
  } else {
    log.Println("Create container ", containerName, " error: ", err);
  }

  return "", "", ""
}

// CreateContainers - create list of containers.
func (t *Town) CreateContainers(checkChanged bool) {
  log.Println("Create...")
  for _, node := range t.cluster.Nodes {

    if !checkChanged || node.Container.Changed {
      ids := make([]string, 0, node.Container.Scale )

      hosts := make([]string, 0, node.Container.Scale)

      log.Println(node.Container.Name, "  image: ", node.Container.Image)
      for i := 1; i <= node.Container.Scale; i++ {

        _, host, containerName := t.CreateContainer(node, i) //id

        if len(node.Container.Exec.Post) > 0 {
          t.bashCommand(containerName, node.Container.Exec.Post)
        }

        ids = append(ids, containerName)
        hosts = append(hosts, host)
      }

      if len(ids) > 1 {
        for index, id := range ids {
          var buffer bytes.Buffer

          buffer.WriteString("echo -e '")
          for i := 0; i < len(hosts); i++ {
            if i != index {
              buffer.WriteString("\n")
              buffer.WriteString(hosts[i])
            }
          }
          buffer.WriteString("' >> /etc/hosts; touch /tmp/host-generated")
          t.bashCommand(id, buffer.String() )
        }
      }

      time.Sleep(1000 * time.Millisecond)
    } else if len(node.Container.Exist) < node.Container.Scale {
      log.Println(node.Container.Name, "  image: ", node.Container.Image, "  ", node.Container.Scale)
      var create =  make([]bool, node.Container.Scale)
      for i := 0; i < node.Container.Scale; i++ {
        create[i] = true
      }
      for _, container := range node.Container.Exist {
        if container.Running {
          create[container.Index - 1] = false;
        }
      }

      for i := 0; i < node.Container.Scale; i++ {
        if create[i] {
          // TODO create or start
          _, _, containerName := t.CreateContainer(node, i + 1)
          // TODO add hosts
          if len(node.Container.Exec.Post) > 0 {
            t.bashCommand(containerName, node.Container.Exec.Post)
          }
        }
      }
    }
  }
}

// bashCommand - execute bash command inside container.
func (t *Town) bashCommand(id string, command string)  {
  config := dockerapi.CreateExecOptions{
    Container:    id,
    AttachStdin:  false,
    AttachStdout: false,
    AttachStderr: false,
    Tty:          false,
    Cmd:          []string{"bash", "-c", command},
  }
  execObj, err := t.docker.CreateExec(config)
  if err == nil {
    startConfig := dockerapi.StartExecOptions{
      Detach: false,
    }
    err = t.docker.StartExec(execObj.ID, startConfig)
    if err != nil {
      log.Println("Container ", id, " command failed with error: ", err, "\n", command)
    }
  } else {
    log.Println("Container ", id, " command failed with error: ", err, "\n", command)
  }
}

func (t *Town) exec(text string, scale int) string {
  replace := strings.Replace(text, "${SCALE_NUM}", strconv.Itoa(scale), -1)
  match := scaleNumRegexp.FindAllStringSubmatch(replace, -1)
  hostMatch := hostsRegexp.FindAllStringSubmatch(replace, -1)
  if len(match) > 0 {
    if len(match[0]) > 1 {
      nums := strings.Split(match[0][1], ",")
      if len(nums) > (scale - 1) {
        replace = strings.Replace(replace, match[0][0], nums[scale - 1], -1)
      }
    }
  }
  if len(hostMatch) > 0 {
    if len(hostMatch[0]) > 1 {
      //nums := strings.Split(, ",")
      name := strings.ToLower(hostMatch[0][1])
      node := t.cluster.FindNodeByID(name)

      var buffer bytes.Buffer
      for i := 1; i <= node.Container.Scale; i++ {
        buffer.WriteString(name)
        buffer.WriteString("-")
        buffer.WriteString(strconv.Itoa( i ))
        if i != node.Container.Scale {
          buffer.WriteString(",")
        }
      }
      replace = strings.Replace(replace, hostMatch[0][0], buffer.String(), -1)
    }
  }
  return replace
}
<commit_msg>add info<commit_after>package main

import (
  "strings"
  "log"
  // "fmt"
//  "regexp"
  "bytes"
  "encoding/json"
  "os"
  "time"
  "strconv"
  "regexp"
  "github.com/lookify/town/cluster"
  dockerapi "github.com/fsouza/go-dockerclient"
)

var (
  scaleNumRegexp, _ = regexp.Compile("\\$\\{SCALE_NUM:(.+)\\}")
//  SCALE_TOTAL_REG, _ = regexp.Compile("\\$\\{SCALE_NUM:(.+)\\}")
  hostsRegexp, _ = regexp.Compile("\\$\\{(.+)_HOSTS\\}")
)

// Town describe cluster and docker clients.
type Town struct {
  cluster *cluster.Cluster
  docker *dockerapi.Client // TODO change to multiple clients
}

// NewTown create new town with default values
func NewTown() *Town {
  return &Town{
    cluster: nil,
    docker: nil,
  }
}

// ReadFile - read town configuration fail in current direcotry or /etc/town (ext .yml)
func (t *Town) ReadFile(name string) {
  var pathLocs = [...]string{
    name + ".yml",
    "/etc/town/" + name + ".yml",
  }

  for _, path := range pathLocs {
    if _, err := os.Stat(path); err == nil {
      t.cluster = cluster.NewCluster(path)
      t.cluster.ReadFile()
      return
    }
  }

  log.Println("ERROR: Could not find file ", name, ".yml")
}

// Connect - connect to docker hosts.
func (t *Town) Connect() {
  endpoint := t.cluster.Application.Docker.Hosts[0] // at the moment use only first
  log.Println("Using Docker API endpont: ", endpoint)
  docker, err := dockerapi.NewClient( endpoint )
  if err != nil {
    log.Println("Can't connect to the docker")
  }
  t.docker = docker
}

// Provision running containers.
func (t *Town) Provision(checkChanged bool) {
  // update containers
  pull := true
  repository := t.cluster.Application.Docker.Repository
  log.Println("repository ", len(repository))
  if len(repository) == 1 && repository[0] == "local" {
    pull = false
  }

  if pull {
    for _, node := range t.cluster.Nodes {
      var buf bytes.Buffer
      var image = strings.Split(node.Container.Image, ":")

      opts := dockerapi.PullImageOptions{
          Repository: image[0],
          // Registry:     "docker.tsuru.io",
          // Tag:          "latest",
          OutputStream: &buf,
      }

      if len(image) > 1 {
        opts.Tag = image[1]
      }
      err := t.docker.PullImage(opts, dockerapi.AuthConfiguration{});
      if err != nil {
        log.Println("Could not pull image ", image)
      }
    }
  }

  allContainers, err := t.docker.ListContainers(dockerapi.ListContainersOptions{
    All: true,
  })
  if err == nil {
    for _, listing := range allContainers {
      container, err := t.docker.InspectContainer(listing.ID)
      if err == nil {
        name := container.Name[1:]
        node, index := t.cluster.FindNodeByName(name)
        if node != nil && index > 0 {
          if node.Container.Exist == nil {
            node.Container.Exist = []cluster.ExistContainer{}
          }
          runningContainer := cluster.NewExistContainer(listing.ID, name, index, container.State.Running)
          runningContainer.Pid = container.State.Pid
          runningContainer.User = container.Config.User
          if checkChanged {
            node.Container.Changed = t.isChangedImage(node, container)
          } else {
            node.Container.Changed = true
          }
          node.Container.Exist = append(node.Container.Exist, runningContainer)
        }
      } else {
        log.Println("[ERROR] Unable to inspect container:", listing.ID[:12], err)
      }
    }

    if checkChanged {
      t.cluster.AddChangeDependant()
    }
  } else {
    log.Println("[ERROR] Can't start provision")
  }
}

// Info - print current cluster information.
func (t *Town) Info() {
  for i := len(t.cluster.Nodes) - 1; i >= 0; i-- {
    node := t.cluster.Nodes[i]
    log.Print("Node ", node.Container.Name, " image ")
	  if node.Container.Changed {
      log.Print("(Changed)")
    }
	  log.Println(": ", node.Container.Image)
    for _, container := range node.Container.Exist {
      log.Print("      ", container.Name, "\t")
  	  if container.Running {
  	    log.Println("Running")
  	  } else {
  	    log.Print("Stoped")
  	  }
    }
  }
}

/**
 * Check node and running container for changes.
 * TODO: add cache to docker call.
 **/
func (t *Town) isChangedImage(node *cluster.Node, container *dockerapi.Container) bool {
  var imageName = container.Image
  image , error := t.docker.InspectImage(imageName)
  if error == nil {
    secondImage , secondError := t.docker.InspectImage(node.Container.Image)
    if secondError == nil {
      return secondImage.Created.After(image.Created)
    }
  }
  log.Println("[ERROR] Could not inspect image ", node.Container.Name)
  return false
}

// StopContainers - stop all containers or only containers with changed images.
func (t *Town) StopContainers(checkChanged bool) {
  log.Println("Stop...")
  //for node := range t.cluster.nodes {
  for i := len(t.cluster.Nodes) - 1; i >= 0; i-- {
    node := t.cluster.Nodes[i]
    if (!checkChanged || node.Container.Changed) && len(node.Container.Exist) > 0 {
      for _, container := range node.Container.Exist {
        if container.Running {
          err := t.docker.StopContainer(container.ID, 10)
          if err == nil {
            log.Println("   -  ", container.Name)
          } else {
            log.Println("   -  ", container.Name, " failed ", err)
          }
        }
      }
    }
  }
  log.Println("=============================")
}

// RemoveContainers - remove container form local repository.
func (t *Town) RemoveContainers(checkChanged bool) {
  log.Println("Remove...")
  //for node := range t.cluster.nodes {
  for i := len(t.cluster.Nodes) - 1; i >= 0; i-- {
    node := t.cluster.Nodes[i]
    if (!checkChanged || node.Container.Changed) && len(node.Container.Exist) > 0 {
      for _, container := range node.Container.Exist {
        err := t.docker.RemoveContainer(dockerapi.RemoveContainerOptions{
          ID: container.ID,
          RemoveVolumes: false,
        })
        if err == nil {
          log.Println("   -  ", container.Name)
        } else {
          log.Println("   -  ", container.Name, " failed ", err)
        }
      }
    }
  }
  log.Println("=============================")
}

// CreateContainer - create container.
func (t *Town) CreateContainer(node *cluster.Node, index int) (string, string, string) {
  containerName := node.Container.Name + "-" + strconv.Itoa(index)

  log.Println("   -  ", containerName)

  node.Container.Hostname = containerName // ?? Help !!!!

  env :=  make([]string, 0, cap(node.Container.Environment))
  for _, e := range node.Container.Environment {
    env = append(env, t.exec(e, index))
  }

  volumes := make(map[string]struct{})
  binds := make([]string, 0, cap(node.Container.Volumes))
  if len(node.Container.Volumes) > 0 {
    for _, volume := range node.Container.Volumes {
      volume = t.exec(volume, index)
      vol := strings.Split(volume, ":")
      if len(vol) > 1 {
        volumes[vol[1]] = struct{}{}
      } else {
        volumes[vol[0]] = struct{}{}
      }
      binds = append(binds, volume)
    }
  }

  dockerConfig := dockerapi.Config{
      Image: node.Container.Image,
      Hostname: node.Container.Hostname,
      PortSpecs: node.Container.Ports,
      Env: env,
      Volumes: volumes,

      AttachStdout: false,
      AttachStdin: false,
      AttachStderr: false,

      Tty: false,

      //Cmd: []
  }

  if len(node.Container.Command) > 0 {
    cmd := t.exec(node.Container.Command, index)
    dockerConfig.Cmd = []string{ cmd }
  }

  // just info
  //for _, l := range node.status.links {
  //  log.Println("     * ", l)
  //}

  // create links
  links := t.cluster.GetLinks(node)

  portBindings := map[dockerapi.Port][]dockerapi.PortBinding{}
  // create ports
  for _, ports := range node.Container.Ports {

    port := strings.Split(ports, ":")
    var p dockerapi.Port

    if len(port) > 1 {
      p = dockerapi.Port(port[1] + "/tcp")
    } else {
      p = dockerapi.Port(port[0] + "/tcp")
    }

    if portBindings[p] == nil {
      portBindings[p] = [] dockerapi.PortBinding {}
    }

    portBindings[p] = append(portBindings[p], dockerapi.PortBinding{
      HostIP: "",
      HostPort: port[0],
    })
  }

  hostConfig := dockerapi.HostConfig{
    Binds: binds,
    Links: links, //, [],
    PortBindings: portBindings,
    NetworkMode: "bridge",
    PublishAllPorts: false,
    Privileged: node.Container.Privileged,
  }

  opts := dockerapi.CreateContainerOptions{Name: containerName, Config: &dockerConfig, HostConfig: &hostConfig}
  container, err := t.docker.CreateContainer(opts)
  if err == nil {
    runningContainer := cluster.NewExistContainer(container.ID, containerName, index, true)
    node.Container.Exist = append(node.Container.Exist, runningContainer)
    // runningContainer.Pid = container.State.Pid
    // runningContainer.User = container.Config.User
    // if checkChanged {
    //   node.Container.Changed = t.isChangedImage(node, container)
    // } else {
    //   node.Container.Changed = true
    // }

    retry := 5
    for retry > 0 {
      error := t.docker.StartContainer(container.ID, &hostConfig)
      if error != nil {
        // log.Println("start error: ", error);

        out, err := json.Marshal(container)
        if err != nil {
            panic (err)
        }
        // fmt.Println(string(out))

        retry--;
        if retry == 0 {
          log.Println(" Start failed after 5 retries: ", string(out))
        }
        // log.Println("retry: ", retry);
      } else {
        inspect, inspectError := t.docker.InspectContainer(container.ID)
        if inspectError == nil {
          //links = append(links, inspect.NetworkSettings.IPAddress + "  " + containerName)
          //ids = append(ids, container.ID)
          return container.ID, inspect.NetworkSettings.IPAddress + "  " + containerName, containerName
        }

        log.Println("Inpect ", container.ID, " error ", inspectError)

        //retry = 0
        break;
      }
    }
  } else {
    log.Println("Create container ", containerName, " error: ", err);
  }

  return "", "", ""
}

// CreateContainers - create list of containers.
func (t *Town) CreateContainers(checkChanged bool) {
  log.Println("Create...")
  for _, node := range t.cluster.Nodes {

    if !checkChanged || node.Container.Changed {
      ids := make([]string, 0, node.Container.Scale )

      hosts := make([]string, 0, node.Container.Scale)

      log.Println(node.Container.Name, "  image: ", node.Container.Image)
      for i := 1; i <= node.Container.Scale; i++ {

        _, host, containerName := t.CreateContainer(node, i) //id

        if len(node.Container.Exec.Post) > 0 {
          t.bashCommand(containerName, node.Container.Exec.Post)
        }

        ids = append(ids, containerName)
        hosts = append(hosts, host)
      }

      if len(ids) > 1 {
        for index, id := range ids {
          var buffer bytes.Buffer

          buffer.WriteString("echo -e '")
          for i := 0; i < len(hosts); i++ {
            if i != index {
              buffer.WriteString("\n")
              buffer.WriteString(hosts[i])
            }
          }
          buffer.WriteString("' >> /etc/hosts; touch /tmp/host-generated")
          t.bashCommand(id, buffer.String() )
        }
      }

      time.Sleep(1000 * time.Millisecond)
    } else if len(node.Container.Exist) < node.Container.Scale {
      log.Println(node.Container.Name, "  image: ", node.Container.Image, "  ", node.Container.Scale)
      var create =  make([]bool, node.Container.Scale)
      for i := 0; i < node.Container.Scale; i++ {
        create[i] = true
      }
      for _, container := range node.Container.Exist {
        if container.Running {
          create[container.Index - 1] = false;
        }
      }

      for i := 0; i < node.Container.Scale; i++ {
        if create[i] {
          // TODO create or start
          _, _, containerName := t.CreateContainer(node, i + 1)
          // TODO add hosts
          if len(node.Container.Exec.Post) > 0 {
            t.bashCommand(containerName, node.Container.Exec.Post)
          }
        }
      }
    }
  }
}

// bashCommand - execute bash command inside container.
func (t *Town) bashCommand(id string, command string)  {
  config := dockerapi.CreateExecOptions{
    Container:    id,
    AttachStdin:  false,
    AttachStdout: false,
    AttachStderr: false,
    Tty:          false,
    Cmd:          []string{"bash", "-c", command},
  }
  execObj, err := t.docker.CreateExec(config)
  if err == nil {
    startConfig := dockerapi.StartExecOptions{
      Detach: false,
    }
    err = t.docker.StartExec(execObj.ID, startConfig)
    if err != nil {
      log.Println("Container ", id, " command failed with error: ", err, "\n", command)
    }
  } else {
    log.Println("Container ", id, " command failed with error: ", err, "\n", command)
  }
}

func (t *Town) exec(text string, scale int) string {
  replace := strings.Replace(text, "${SCALE_NUM}", strconv.Itoa(scale), -1)
  match := scaleNumRegexp.FindAllStringSubmatch(replace, -1)
  hostMatch := hostsRegexp.FindAllStringSubmatch(replace, -1)
  if len(match) > 0 {
    if len(match[0]) > 1 {
      nums := strings.Split(match[0][1], ",")
      if len(nums) > (scale - 1) {
        replace = strings.Replace(replace, match[0][0], nums[scale - 1], -1)
      }
    }
  }
  if len(hostMatch) > 0 {
    if len(hostMatch[0]) > 1 {
      //nums := strings.Split(, ",")
      name := strings.ToLower(hostMatch[0][1])
      node := t.cluster.FindNodeByID(name)

      var buffer bytes.Buffer
      for i := 1; i <= node.Container.Scale; i++ {
        buffer.WriteString(name)
        buffer.WriteString("-")
        buffer.WriteString(strconv.Itoa( i ))
        if i != node.Container.Scale {
          buffer.WriteString(",")
        }
      }
      replace = strings.Replace(replace, hostMatch[0][0], buffer.String(), -1)
    }
  }
  return replace
}
<|endoftext|>


<commit_before>// Implementation of an R-Way Trie data structure.
//
// A Trie has a root Node which is the base of the tree.
// Each subsequent Node has a letter and children, which are
// nodes that have letter values associated with them.

package trie

type Node struct {
	val      rune
	term     bool
	meta     interface{}
	mask     uint64
	parent   *Node
	children map[rune]*Node
}

type Trie struct {
	root *Node
	size int
}

const nul = 0x0

func newNode(parent *Node, val rune, m uint64, term bool) *Node {
	return &Node{
		val:      val,
		mask:     m,
		term:     term,
		parent:   parent,
		children: make(map[rune]*Node),
	}
}

// Creates and returns a pointer to a new child for the node.
func (n *Node) NewChild(r rune, bitmask uint64, val rune, term bool) *Node {
	node := newNode(n, val, bitmask, term)
	n.children[r] = node
	return node
}

func (n *Node) RemoveChild(r rune) {
	delete(n.children, r)

	n.recalculateMask()
	for parent := n.Parent(); parent != nil; parent = parent.Parent() {
		parent.recalculateMask()
	}
}

func (n *Node) recalculateMask() {
	n.mask = maskrune(n.Val())
	for k, c := range n.Children() {
		n.mask |= (maskrune(k) | c.Mask())
	}
}

// Returns the parent of this node.
func (n Node) Parent() *Node {
	return n.parent
}

// Returns the meta information of this node.
func (n Node) Meta() interface{} {
	return n.meta
}

// Returns the children of this node.
func (n Node) Children() map[rune]*Node {
	return n.children
}

func (n Node) Val() rune {
	return n.val
}

// Returns a uint64 representing the current
// mask of this node.
func (n Node) Mask() uint64 {
	return n.mask
}

// Creates a new Trie with an initialized root Node.
func CreateTrie() *Trie {
	node := newNode(nil, 0, 0, false)
	return &Trie{
		root: node,
		size: 0,
	}
}

// Returns the root node for the Trie.
func (t *Trie) Root() *Node {
	return t.root
}

// Adds the key to the Trie.
func (t *Trie) Add(key string, meta interface{}) *Node {
	t.size++
	runes := []rune(key)
	node := t.addrune(t.Root(), runes, 0)
	node.meta = meta
	return node
}

// Removes a key from the trie, ensuring that
// all bitmasks up to root are appropriately recalculated.
func (t *Trie) Remove(key string) {
	var (
		i    int
		rs   = []rune(key)
		node = t.nodeAtPath(key)
	)

	t.size--
	for n := node.Parent(); n != nil; n = n.Parent() {
		i++
		if len(n.Children()) > 1 {
			r := rs[len(rs)-i]
			n.RemoveChild(r)
			break
		}
	}
}

// Returns all the keys currently stored in the trie.
func (t *Trie) Keys() []string {
	return t.PrefixSearch("")
}

// Performs a fuzzy search against the keys in the trie.
func (t Trie) FuzzySearch(pre string) []string {
	var (
		keys []string
		pm   []rune
	)

	fuzzycollect(t.Root(), pm, []rune(pre), &keys)
	return keys
}

// Performs a prefix search against the keys in the trie.
func (t Trie) PrefixSearch(pre string) []string {
	var keys []string

	node := t.nodeAtPath(pre)
	if node == nil {
		return keys
	}

	collect(node, []rune(pre), &keys)
	return keys
}

func (t Trie) nodeAtPath(pre string) *Node {
	runes := []rune(pre)
	return findNode(t.Root(), runes, 0)
}

func findNode(node *Node, runes []rune, d int) *Node {
	if node == nil {
		return nil
	}

	if len(runes) == 0 {
		return node
	}

	upper := len(runes)
	if d == upper {
		return node
	}

	n, ok := node.Children()[runes[d]]
	if !ok {
		return nil
	}

	d++
	return findNode(n, runes, d)
}

func (t Trie) addrune(node *Node, runes []rune, i int) *Node {
	if len(runes) == 0 {
		return node.NewChild(0, 0, nul, true)
	}

	r := runes[0]
	c := node.Children()

	n, ok := c[r]
	bitmask := maskruneslice(runes)
	if !ok {
		n = node.NewChild(r, bitmask, r, false)
	}
	n.mask |= bitmask

	i++
	return t.addrune(n, runes[1:], i)
}

func maskruneslice(rs []rune) uint64 {
	var m uint64
	for _, r := range rs {
		m |= maskrune(r)
	}

	return m
}

func maskrune(r rune) uint64 {
	i := uint64(1)
	return i << (uint64(r) - 97)
}

func collect(node *Node, pre []rune, keys *[]string) {
	children := node.Children()
	for r, n := range children {
		if n.term {
			*keys = append(*keys, string(pre))
			continue
		}

		npre := append(pre, r)
		collect(n, npre, keys)
	}
}

func fuzzycollect(node *Node, partialmatch, partial []rune, keys *[]string) {
	partiallen := len(partial)

	if partiallen == 0 {
		collect(node, partialmatch, keys)
		return
	}

	m := maskruneslice(partial)
	children := node.Children()
	for v, n := range children {
		xor := n.Mask() ^ m
		if (xor & m) != 0 {
			continue
		}

		npartial := partial
		if v == partial[0] {
			if partiallen > 1 {
				npartial = partial[1:]
			} else {
				npartial = partial[0:0]
			}
		}

		fuzzycollect(n, append(partialmatch, v), npartial, keys)
	}
}
<commit_msg>Add documentation for Add()<commit_after>// Implementation of an R-Way Trie data structure.
//
// A Trie has a root Node which is the base of the tree.
// Each subsequent Node has a letter and children, which are
// nodes that have letter values associated with them.

package trie

type Node struct {
	val      rune
	term     bool
	meta     interface{}
	mask     uint64
	parent   *Node
	children map[rune]*Node
}

type Trie struct {
	root *Node
	size int
}

const nul = 0x0

func newNode(parent *Node, val rune, m uint64, term bool) *Node {
	return &Node{
		val:      val,
		mask:     m,
		term:     term,
		parent:   parent,
		children: make(map[rune]*Node),
	}
}

// Creates and returns a pointer to a new child for the node.
func (n *Node) NewChild(r rune, bitmask uint64, val rune, term bool) *Node {
	node := newNode(n, val, bitmask, term)
	n.children[r] = node
	return node
}

func (n *Node) RemoveChild(r rune) {
	delete(n.children, r)

	n.recalculateMask()
	for parent := n.Parent(); parent != nil; parent = parent.Parent() {
		parent.recalculateMask()
	}
}

func (n *Node) recalculateMask() {
	n.mask = maskrune(n.Val())
	for k, c := range n.Children() {
		n.mask |= (maskrune(k) | c.Mask())
	}
}

// Returns the parent of this node.
func (n Node) Parent() *Node {
	return n.parent
}

// Returns the meta information of this node.
func (n Node) Meta() interface{} {
	return n.meta
}

// Returns the children of this node.
func (n Node) Children() map[rune]*Node {
	return n.children
}

func (n Node) Val() rune {
	return n.val
}

// Returns a uint64 representing the current
// mask of this node.
func (n Node) Mask() uint64 {
	return n.mask
}

// Creates a new Trie with an initialized root Node.
func CreateTrie() *Trie {
	node := newNode(nil, 0, 0, false)
	return &Trie{
		root: node,
		size: 0,
	}
}

// Returns the root node for the Trie.
func (t *Trie) Root() *Node {
	return t.root
}

// Adds the key to the Trie, including meta data. Meta data
// is stored as `interface{}` and must be type cast by
// the caller.
func (t *Trie) Add(key string, meta interface{}) *Node {
	t.size++
	runes := []rune(key)
	node := t.addrune(t.Root(), runes, 0)
	node.meta = meta
	return node
}

// Removes a key from the trie, ensuring that
// all bitmasks up to root are appropriately recalculated.
func (t *Trie) Remove(key string) {
	var (
		i    int
		rs   = []rune(key)
		node = t.nodeAtPath(key)
	)

	t.size--
	for n := node.Parent(); n != nil; n = n.Parent() {
		i++
		if len(n.Children()) > 1 {
			r := rs[len(rs)-i]
			n.RemoveChild(r)
			break
		}
	}
}

// Returns all the keys currently stored in the trie.
func (t *Trie) Keys() []string {
	return t.PrefixSearch("")
}

// Performs a fuzzy search against the keys in the trie.
func (t Trie) FuzzySearch(pre string) []string {
	var (
		keys []string
		pm   []rune
	)

	fuzzycollect(t.Root(), pm, []rune(pre), &keys)
	return keys
}

// Performs a prefix search against the keys in the trie.
func (t Trie) PrefixSearch(pre string) []string {
	var keys []string

	node := t.nodeAtPath(pre)
	if node == nil {
		return keys
	}

	collect(node, []rune(pre), &keys)
	return keys
}

func (t Trie) nodeAtPath(pre string) *Node {
	runes := []rune(pre)
	return findNode(t.Root(), runes, 0)
}

func findNode(node *Node, runes []rune, d int) *Node {
	if node == nil {
		return nil
	}

	if len(runes) == 0 {
		return node
	}

	upper := len(runes)
	if d == upper {
		return node
	}

	n, ok := node.Children()[runes[d]]
	if !ok {
		return nil
	}

	d++
	return findNode(n, runes, d)
}

func (t Trie) addrune(node *Node, runes []rune, i int) *Node {
	if len(runes) == 0 {
		return node.NewChild(0, 0, nul, true)
	}

	r := runes[0]
	c := node.Children()

	n, ok := c[r]
	bitmask := maskruneslice(runes)
	if !ok {
		n = node.NewChild(r, bitmask, r, false)
	}
	n.mask |= bitmask

	i++
	return t.addrune(n, runes[1:], i)
}

func maskruneslice(rs []rune) uint64 {
	var m uint64
	for _, r := range rs {
		m |= maskrune(r)
	}

	return m
}

func maskrune(r rune) uint64 {
	i := uint64(1)
	return i << (uint64(r) - 97)
}

func collect(node *Node, pre []rune, keys *[]string) {
	children := node.Children()
	for r, n := range children {
		if n.term {
			*keys = append(*keys, string(pre))
			continue
		}

		npre := append(pre, r)
		collect(n, npre, keys)
	}
}

func fuzzycollect(node *Node, partialmatch, partial []rune, keys *[]string) {
	partiallen := len(partial)

	if partiallen == 0 {
		collect(node, partialmatch, keys)
		return
	}

	m := maskruneslice(partial)
	children := node.Children()
	for v, n := range children {
		xor := n.Mask() ^ m
		if (xor & m) != 0 {
			continue
		}

		npartial := partial
		if v == partial[0] {
			if partiallen > 1 {
				npartial = partial[1:]
			} else {
				npartial = partial[0:0]
			}
		}

		fuzzycollect(n, append(partialmatch, v), npartial, keys)
	}
}
<|endoftext|>


<commit_before>package arc

func min(x, y int) int {
	if x < y {
		return x
	} else {
		return y
	}
}

func max(x, y int) int {
	if x < y {
		return x
	} else {
		return y
	}
}
<commit_msg>utils.go<commit_after>package arc

func min(x, y int) int {
	if x < y {
		return x
	}
	return y
}

func max(x, y int) int {
	if x > y {
		return x
	}
	return y
}
<|endoftext|>


<commit_before>package pixel

import "image/color"

// colorToRGBA converts a color from image/color to RGBA components in interval [0, 1].
func colorToRGBA(c color.Color) (r, g, b, a float64) {
	ri, gi, bi, ai := c.RGBA()
	r = float64(ri) / 0xffff
	g = float64(gi) / 0xffff
	b = float64(bi) / 0xffff
	a = float64(ai) / 0xffff
	return
}
<commit_msg>change float64 to float32 in colorToRGBA<commit_after>package pixel

import "image/color"

// colorToRGBA converts a color from image/color to RGBA components in interval [0, 1].
func colorToRGBA(c color.Color) (r, g, b, a float32) {
	ri, gi, bi, ai := c.RGBA()
	r = float32(ri) / 0xffff
	g = float32(gi) / 0xffff
	b = float32(bi) / 0xffff
	a = float32(ai) / 0xffff
	return
}
<|endoftext|>


<commit_before>package CloudForest

import ()

//Tree represents a single decision tree.
type Tree struct {
	//Tree int
	Root *Node
}

//AddNode adds a node a the specified path with the specivied pred value and/or
//Splitter. Paths are specified in the same format as in rf-aces sf files, as a
//string of 'L' and 'R'. Nodes must be added from the root up as the case where
//the path specifies a node whose parent does not allready exist in the tree is
//not handled well.
func (t *Tree) AddNode(path string, pred string, splitter *Splitter) {
	n := new(Node)
	n.Pred = pred
	n.Splitter = splitter
	if t.Root == nil {
		t.Root = n
	} else {
		loc := t.Root
		for i := 0; i < len(path); i++ {
			switch path[i : i+1] {
			case "L":
				if loc.Left == nil {
					loc.Left = n
				}
				loc = loc.Left
			case "R":
				if loc.Right == nil {
					loc.Right = n
				}
				loc = loc.Right
			}

		}
	}

}

/*
tree.Grow grows the reciever tree through recursion. It uses impurity decrease to selct spliters at
each node as in Brieman's Random Forest. It should be called on a tree with only a root node defined.

fm is a feature matrix of training data.

target is the feature to predict via regression or classification as determined by feature type.

cases specifies the cases to calculate impurity decrease over and can contain repeated values
to allow for sampeling of cases with replacment as in RF.

mTry specifies the number of canidate features to evaluate for each split.

leafSize specifies the minimum number of cases at a leafNode.
*/
func (t *Tree) Grow(fm *FeatureMatrix,
	target *Feature,
	cases []int,
	canidates []int,
	mTry int,
	leafSize int,
	itter bool,
	l *[]int,
	r *[]int) {

	t.Root.Recurse(func(n *Node, innercases []int) {

		if 2*leafSize <= len(innercases) {
			SampleFirstN(&canidates, mTry)
			best, impDec := target.BestSplitter(fm, innercases, canidates[:mTry], itter, l, r)
			if best != nil && impDec > minImp {
				//not a leaf node so define the spliter and left and right nodes
				//so recursion will continue
				n.Splitter = best
				n.Pred = ""
				n.Left = &Node{nil, nil, "", nil}
				n.Right = &Node{nil, nil, "", nil}
				return
			}
		}

		//Leaf node so find the predictive value and set it in n.Pred
		n.Splitter = nil
		n.Pred = target.FindPredicted(innercases)

	}, fm, cases)
}

//GetSplits returns the arrays of all Numeric spliters of a tree.
func (t *Tree) GetSplits(fm *FeatureMatrix, fbycase *SparseCounter, relativeSplitCount *SparseCounter) []Splitter {
	splitters := make([]Splitter, 0)
	ncases := len(fm.Data[0].Missing) // grab the number of samples for the first feature
	cases := make([]int, ncases)      //make an array that large
	for i, _ := range cases {
		cases[i] = i
	}

	t.Root.Recurse(func(n *Node, cases []int) {
		//if we're on a splitting node
		if fbycase != nil && n.Splitter != nil && n.Splitter.Numerical == true {
			//add this splitter to the list
			splitters = append(splitters, Splitter{n.Splitter.Feature,
				n.Splitter.Numerical,
				n.Splitter.Value,
				n.Splitter.Left})
			f_id := n.Splitter.Feature               //get the feature at this splitter
			f := fm.Data[fm.Map[n.Splitter.Feature]] //get the feature at this splitter
			for _, c := range cases {                //for each case

				if f.Missing[c] == false { //if there isa value for this case
					fbycase.Add(c, fm.Map[f_id], 1) //count the number of times each case is present for a split by a feature

					switch f.NumData[c] <= n.Splitter.Value {
					case true: //if the value was split to the left
						relativeSplitCount.Add(c, fm.Map[f_id], -1) //subtract one
					case false:
						relativeSplitCount.Add(c, fm.Map[f_id], 1) //add one
					}
				}
			}
		}
	}, fm, cases)
	return splitters //return the array of all splitters from the tree

}

//GetLeaves is called by the leaf count utility to
//gather statistics about the nodes of a tree including the sets of cases at
//"leaf" nodes that aren't split further and the number of times each feature
//is used to split away each case.
func (t *Tree) GetLeaves(fm *FeatureMatrix, fbycase *SparseCounter) []Leaf {
	leaves := make([]Leaf, 0)
	ncases := len(fm.Data[0].Missing)
	cases := make([]int, 0, ncases)
	for i := 0; i < ncases; i++ {
		cases = append(cases, i)
	}

	t.Root.Recurse(func(n *Node, cases []int) {
		if n.Left == nil && n.Right == nil { // I'm in a leaf node
			leaves = append(leaves, Leaf{cases, n.Pred})
		}
		if fbycase != nil && n.Splitter != nil { //I'm not in a leaf node?
			for _, c := range cases {
				fbycase.Add(c, fm.Map[n.Splitter.Feature], 1)
			}
		}

	}, fm, cases)
	return leaves

}

//Leaf is a struct for storing the index of the cases at a terminal "Leaf" node
//along with the Numeric predicted value.
type Leaf struct {
	Cases []int
	Pred  string
}

//Tree.Vote casts a vote for the predicted value of each case in fm *FeatureMatrix.
//into bb *BallotBox. Since BallotBox is not thread safe trees should not vote
//into the same BallotBox in parralel.
func (t *Tree) Vote(fm *FeatureMatrix, bb VoteTallyer) {
	ncases := len(fm.Data[0].Missing)
	cases := make([]int, 0, ncases)
	for i := 0; i < ncases; i++ {
		cases = append(cases, i)
	}

	t.Root.Recurse(func(n *Node, cases []int) {
		if n.Left == nil && n.Right == nil {
			// I'm in a leaf node
			for i := 0; i < len(cases); i++ {
				bb.Vote(cases[i], n.Pred)
			}
		}
	}, fm, cases)
}
<commit_msg>Don't split small leafs<commit_after>package CloudForest

import ()

//Tree represents a single decision tree.
type Tree struct {
	//Tree int
	Root *Node
}

//AddNode adds a node a the specified path with the specivied pred value and/or
//Splitter. Paths are specified in the same format as in rf-aces sf files, as a
//string of 'L' and 'R'. Nodes must be added from the root up as the case where
//the path specifies a node whose parent does not allready exist in the tree is
//not handled well.
func (t *Tree) AddNode(path string, pred string, splitter *Splitter) {
	n := new(Node)
	n.Pred = pred
	n.Splitter = splitter
	if t.Root == nil {
		t.Root = n
	} else {
		loc := t.Root
		for i := 0; i < len(path); i++ {
			switch path[i : i+1] {
			case "L":
				if loc.Left == nil {
					loc.Left = n
				}
				loc = loc.Left
			case "R":
				if loc.Right == nil {
					loc.Right = n
				}
				loc = loc.Right
			}

		}
	}

}

/*
tree.Grow grows the reciever tree through recursion. It uses impurity decrease to selct spliters at
each node as in Brieman's Random Forest. It should be called on a tree with only a root node defined.

fm is a feature matrix of training data.

target is the feature to predict via regression or classification as determined by feature type.

cases specifies the cases to calculate impurity decrease over and can contain repeated values
to allow for sampeling of cases with replacment as in RF.

mTry specifies the number of canidate features to evaluate for each split.

leafSize specifies the minimum number of cases at a leafNode.
*/
func (t *Tree) Grow(fm *FeatureMatrix,
	target *Feature,
	cases []int,
	canidates []int,
	mTry int,
	leafSize int,
	itter bool,
	l *[]int,
	r *[]int) {

	t.Root.Recurse(func(n *Node, innercases []int) {

		if (2 * leafSize) <= len(innercases) {
			SampleFirstN(&canidates, mTry)
			best, impDec := target.BestSplitter(fm, innercases, canidates[:mTry], itter, l, r)
			if best != nil && impDec > minImp {
				//not a leaf node so define the spliter and left and right nodes
				//so recursion will continue
				n.Splitter = best
				n.Pred = ""
				n.Left = &Node{nil, nil, "", nil}
				n.Right = &Node{nil, nil, "", nil}
				return
			}
		}

		//Leaf node so find the predictive value and set it in n.Pred
		n.Splitter = nil
		n.Pred = target.FindPredicted(innercases)

	}, fm, cases)
}

//GetSplits returns the arrays of all Numeric spliters of a tree.
func (t *Tree) GetSplits(fm *FeatureMatrix, fbycase *SparseCounter, relativeSplitCount *SparseCounter) []Splitter {
	splitters := make([]Splitter, 0)
	ncases := len(fm.Data[0].Missing) // grab the number of samples for the first feature
	cases := make([]int, ncases)      //make an array that large
	for i, _ := range cases {
		cases[i] = i
	}

	t.Root.Recurse(func(n *Node, cases []int) {
		//if we're on a splitting node
		if fbycase != nil && n.Splitter != nil && n.Splitter.Numerical == true {
			//add this splitter to the list
			splitters = append(splitters, Splitter{n.Splitter.Feature,
				n.Splitter.Numerical,
				n.Splitter.Value,
				n.Splitter.Left})
			f_id := n.Splitter.Feature               //get the feature at this splitter
			f := fm.Data[fm.Map[n.Splitter.Feature]] //get the feature at this splitter
			for _, c := range cases {                //for each case

				if f.Missing[c] == false { //if there isa value for this case
					fbycase.Add(c, fm.Map[f_id], 1) //count the number of times each case is present for a split by a feature

					switch f.NumData[c] <= n.Splitter.Value {
					case true: //if the value was split to the left
						relativeSplitCount.Add(c, fm.Map[f_id], -1) //subtract one
					case false:
						relativeSplitCount.Add(c, fm.Map[f_id], 1) //add one
					}
				}
			}
		}
	}, fm, cases)
	return splitters //return the array of all splitters from the tree

}

//GetLeaves is called by the leaf count utility to
//gather statistics about the nodes of a tree including the sets of cases at
//"leaf" nodes that aren't split further and the number of times each feature
//is used to split away each case.
func (t *Tree) GetLeaves(fm *FeatureMatrix, fbycase *SparseCounter) []Leaf {
	leaves := make([]Leaf, 0)
	ncases := len(fm.Data[0].Missing)
	cases := make([]int, 0, ncases)
	for i := 0; i < ncases; i++ {
		cases = append(cases, i)
	}

	t.Root.Recurse(func(n *Node, cases []int) {
		if n.Left == nil && n.Right == nil { // I'm in a leaf node
			leaves = append(leaves, Leaf{cases, n.Pred})
		}
		if fbycase != nil && n.Splitter != nil { //I'm not in a leaf node?
			for _, c := range cases {
				fbycase.Add(c, fm.Map[n.Splitter.Feature], 1)
			}
		}

	}, fm, cases)
	return leaves

}

//Leaf is a struct for storing the index of the cases at a terminal "Leaf" node
//along with the Numeric predicted value.
type Leaf struct {
	Cases []int
	Pred  string
}

//Tree.Vote casts a vote for the predicted value of each case in fm *FeatureMatrix.
//into bb *BallotBox. Since BallotBox is not thread safe trees should not vote
//into the same BallotBox in parralel.
func (t *Tree) Vote(fm *FeatureMatrix, bb VoteTallyer) {
	ncases := len(fm.Data[0].Missing)
	cases := make([]int, 0, ncases)
	for i := 0; i < ncases; i++ {
		cases = append(cases, i)
	}

	t.Root.Recurse(func(n *Node, cases []int) {
		if n.Left == nil && n.Right == nil {
			// I'm in a leaf node
			for i := 0; i < len(cases); i++ {
				bb.Vote(cases[i], n.Pred)
			}
		}
	}, fm, cases)
}
<|endoftext|>


<commit_before>package trie

import (
	"errors"
)

type Trie struct {
	Value    string
	Children []Child
	End      bool
}

type Child struct {
	Value rune
	Link  *Trie
}

func NewTrie() *Trie {
	return &Trie{}
}

func NewChild(value rune) Child {
	child := Child{}
	child.Value = value
	return child
}

func (t *Trie) Add(value rune) *Trie {
	link := NewTrie()
	link.Value = t.Value + string(value)
	child := NewChild(value)
	child.Link = link

	t.Children = append(t.Children, child)

	return link
}

func (t *Trie) Insert(value string) {
	trie := t

	for i, val := range value {
		link, err := trie.Find(val)

		if i < len(value) {
			if err != nil {
				link = trie.Add(val)
			}

			trie = link
		}

	}

	trie.End = true
}

func (t *Trie) Find(value rune) (*Trie, error) {
	for _, child := range t.Children {
		if child.Value == value {
			return child.Link, nil
		}
	}
	return nil, errors.New("Value not found")
}

func (t *Trie) AllPrefixes() []string {
	results := []string{}

	if t.End == true {
		results = append(results, t.Value)
	}

	for _, val := range t.Children {
		prefixes := val.Link.AllPrefixes()
		results = append(results, prefixes...)
	}

	return results
}

func (t *Trie) AutoComplete(prefix string) []string {
	trie := t

	for _, val := range prefix {
		link, _ := trie.Find(val)
		trie = link
	}

	results := trie.AllPrefixes()

	return results
}
<commit_msg>Fixed bug with not available chars<commit_after>package trie

import (
	"errors"
)

type Trie struct {
	Value    string
	Children []Child
	End      bool
}

type Child struct {
	Value rune
	Link  *Trie
}

func NewTrie() *Trie {
	return &Trie{}
}

func NewChild(value rune) Child {
	child := Child{}
	child.Value = value
	return child
}

func (t *Trie) Add(value rune) *Trie {
	link := NewTrie()
	link.Value = t.Value + string(value)
	child := NewChild(value)
	child.Link = link

	t.Children = append(t.Children, child)

	return link
}

func (t *Trie) Insert(value string) {
	trie := t

	for i, val := range value {
		link, err := trie.Find(val)

		if i < len(value) {
			if err != nil {
				link = trie.Add(val)
			}

			trie = link
		}

	}

	trie.End = true
}

func (t *Trie) Find(value rune) (*Trie, error) {
	if t != nil {
		for _, child := range t.Children {
			if child.Value == value {
				return child.Link, nil
			}
		}
	}
	return nil, errors.New("Value not found")
}

func (t *Trie) AllPrefixes() ([]string, error) {
	results := []string{}
	if t != nil {
		if t.End == true {
			results = append(results, t.Value)
		}

		for _, val := range t.Children {
			prefixes, err := val.Link.AllPrefixes()
			if err != nil {
				break
			}
			results = append(results, prefixes...)
		}
	} else {
		return results, errors.New("No end found, so no completions available")
	}
	return results, nil
}

func (t *Trie) AutoComplete(prefix string) ([]string, error) {
	trie := t

	for _, val := range prefix {
		link, _ := trie.Find(val)
		trie = link
	}

	results, err := trie.AllPrefixes()

	return results, err
}
<|endoftext|>


<commit_before>package archive

import (
	"bufio"
	"bytes"
	"fmt"
	"log"
	"os"
	"os/exec"
	"path/filepath"
	"strings"

	version "github.com/hashicorp/go-version"
)

// VCS is a struct that explains how to get the file list for a given
// VCS.
type VCS struct {
	Name string

	// Detect is a list of files/folders that if they exist, signal that
	// this VCS is the VCS in use.
	Detect []string

	// Files returns the files that are under version control for the
	// given path.
	Files VCSFilesFunc

	// Metadata returns arbitrary metadata about the underlying VCS for the
	// given path.
	Metadata VCSMetadataFunc

	// Preflight is a function to run before looking for VCS files.
	Preflight VCSPreflightFunc
}

// VCSList is the list of VCS we recognize.
var VCSList = []*VCS{
	&VCS{
		Name:   "git",
		Detect: []string{".git/"},
		Preflight: func(path string) error {
			var stderr, stdout bytes.Buffer

			cmd := exec.Command("git", "--version")
			cmd.Dir = path
			cmd.Stdout = &stdout
			cmd.Stderr = &stderr
			if err := cmd.Run(); err != nil {
				return fmt.Errorf("error getting git version: %s\nstdout: %s\nstderr: %s",
					err, stdout.String(), stderr.String())
			}

			// Check if the output is valid
			output := strings.Split(strings.TrimSpace(stdout.String()), " ")
			if len(output) < 1 {
				log.Printf("[WARN] could not extract version output from Git")
				return nil
			}

			// Parse the version
			gitv, err := version.NewVersion(output[len(output)-1])
			if err != nil {
				log.Printf("[WARN] could not parse version output from Git")
				return nil
			}

			constraint, err := version.NewConstraint("> 1.8")
			if err != nil {
				log.Printf("[WARN] could not create version constraint to check")
				return nil
			}
			if !constraint.Check(gitv) {
				return fmt.Errorf("git version (%s) is too old, please upgrade", gitv.String())
			}

			return nil
		},
		Files: vcsFilesCmd("git", "ls-files"),
		Metadata: func(path string) (map[string]string, error) {
			// Future-self note: Git is NOT threadsafe, so we cannot run these
			// operations in go routines or else you're going to have a really really
			// bad day and Panda.State == "Sad" :(

			branch, err := gitBranch(path)
			if err != nil {
				return nil, err
			}

			commit, err := gitCommit(path)
			if err != nil {
				return nil, err
			}

			remotes, err := gitRemotes(path)
			if err != nil {
				return nil, err
			}

			// Make the return result (we already know the size)
			result := make(map[string]string, 2+len(remotes))

			result["branch"] = branch
			result["commit"] = commit
			for remote, value := range remotes {
				result[remote] = value
			}

			return result, nil
		},
	},
	&VCS{
		Name:   "hg",
		Detect: []string{".hg/"},
		Files:  vcsTrimCmd(vcsFilesCmd("hg", "locate", "-f", "--include", ".")),
	},
	&VCS{
		Name:   "svn",
		Detect: []string{".svn/"},
		Files:  vcsFilesCmd("svn", "ls"),
	},
}

// VCSFilesFunc is the callback invoked to return the files in the VCS.
//
// The return value should be paths relative to the given path.
type VCSFilesFunc func(string) ([]string, error)

// VCSMetadataFunc is the callback invoked to get arbitrary information about
// the current VCS.
//
// The return value should be a map of key-value pairs.
type VCSMetadataFunc func(string) (map[string]string, error)

// VCSPreflightFunc is a function that runs before VCS detection to be
// configured by the user. It may be used to check if pre-requisites (like the
// actual VCS) are installed or that a program is at the correct version. If an
// error is returned, the VCS will not be processed and the error will be
// returned up the stack.
//
// The given argument is the path where the VCS is running.
type VCSPreflightFunc func(string) error

// vcsDetect detects the VCS that is used for path.
func vcsDetect(path string) (*VCS, error) {
	for _, v := range VCSList {
		dir := path
		for {
			for _, f := range v.Detect {
				check := filepath.Join(dir, f)
				if _, err := os.Stat(check); err == nil {
					return v, nil
				}
			}

			lastDir := dir
			dir = filepath.Dir(dir)
			if dir == lastDir {
				break
			}
		}
	}

	return nil, nil
}

// vcsPreflight returns the metadata for the VCS directory path.
func vcsPreflight(path string) error {
	vcs, err := vcsDetect(path)
	if err != nil {
		return fmt.Errorf("error detecting VCS: %s", err)
	}
	if vcs == nil {
		return fmt.Errorf("no VCS found for path: %s", path)
	}

	if vcs.Preflight != nil {
		return vcs.Preflight(path)
	}

	return nil
}

// vcsFiles returns the files for the VCS directory path.
func vcsFiles(path string) ([]string, error) {
	vcs, err := vcsDetect(path)
	if err != nil {
		return nil, fmt.Errorf("error detecting VCS: %s", err)
	}
	if vcs == nil {
		return nil, fmt.Errorf("no VCS found for path: %s", path)
	}

	if vcs.Files != nil {
		return vcs.Files(path)
	}

	return nil, nil
}

// vcsFilesCmd creates a Files-compatible function that reads the files
// by executing the command in the repository path and returning each
// line in stdout.
func vcsFilesCmd(args ...string) VCSFilesFunc {
	return func(path string) ([]string, error) {
		var stderr, stdout bytes.Buffer

		cmd := exec.Command(args[0], args[1:]...)
		cmd.Dir = path
		cmd.Stdout = &stdout
		cmd.Stderr = &stderr
		if err := cmd.Run(); err != nil {
			return nil, fmt.Errorf(
				"error executing %s: %s",
				strings.Join(args, " "),
				err)
		}

		// Read each line of output as a path
		result := make([]string, 0, 100)
		scanner := bufio.NewScanner(&stdout)
		for scanner.Scan() {
			result = append(result, scanner.Text())
		}

		return result, nil
	}
}

// vcsTrimCmd trims the prefix from the paths returned by another VCSFilesFunc.
// This should be used to wrap another function if the return value is known
// to have full paths rather than relative paths
func vcsTrimCmd(f VCSFilesFunc) VCSFilesFunc {
	return func(path string) ([]string, error) {
		absPath, err := filepath.Abs(path)
		if err != nil {
			return nil, fmt.Errorf(
				"error expanding VCS path: %s", err)
		}

		// Now that we have the root path, get the inner files
		fs, err := f(path)
		if err != nil {
			return nil, err
		}

		// Trim the root path from the files
		result := make([]string, 0, len(fs))
		for _, f := range fs {
			if !strings.HasPrefix(f, absPath) {
				continue
			}

			f, err = filepath.Rel(absPath, f)
			if err != nil {
				return nil, fmt.Errorf(
					"error determining path: %s", err)
			}

			result = append(result, f)
		}

		return result, nil
	}
}

// vcsMetadata returns the metadata for the VCS directory path.
func vcsMetadata(path string) (map[string]string, error) {
	vcs, err := vcsDetect(path)
	if err != nil {
		return nil, fmt.Errorf("error detecting VCS: %s", err)
	}
	if vcs == nil {
		return nil, fmt.Errorf("no VCS found for path: %s", path)
	}

	if vcs.Metadata != nil {
		return vcs.Metadata(path)
	}

	return nil, nil
}

// gitBranch gets and returns the current git branch for the Git repository
// at the given path. It is assumed that the VCS is git.
func gitBranch(path string) (string, error) {
	var stderr, stdout bytes.Buffer

	cmd := exec.Command("git", "symbolic-ref", "--short", "HEAD")
	cmd.Dir = path
	cmd.Stdout = &stdout
	cmd.Stderr = &stderr
	if err := cmd.Run(); err != nil {
		return "", fmt.Errorf("error getting git branch: %s\nstdout: %s\nstderr: %s",
			err, stdout.String(), stderr.String())
	}

	branch := strings.TrimSpace(stdout.String())

	return branch, nil
}

// gitCommit gets the SHA of the latest commit for the Git repository at the
// given path. It is assumed that the VCS is git.
func gitCommit(path string) (string, error) {
	var stderr, stdout bytes.Buffer

	cmd := exec.Command("git", "log", "-n1", "--pretty=format:%H")
	cmd.Dir = path
	cmd.Stdout = &stdout
	cmd.Stderr = &stderr
	if err := cmd.Run(); err != nil {
		return "", fmt.Errorf("error getting git commit: %s\nstdout: %s\nstderr: %s",
			err, stdout.String(), stderr.String())
	}

	commit := strings.TrimSpace(stdout.String())

	return commit, nil
}

// gitRemotes gets and returns a map of all remotes for the Git repository. The
// map key is the name of the remote of the format "remote.NAME" and the value
// is the endpoint for the remote. It is assumed that the VCS is git.
func gitRemotes(path string) (map[string]string, error) {
	var stderr, stdout bytes.Buffer

	cmd := exec.Command("git", "remote", "-v")
	cmd.Dir = path
	cmd.Stdout = &stdout
	cmd.Stderr = &stderr
	if err := cmd.Run(); err != nil {
		return nil, fmt.Errorf("error getting git remotes: %s\nstdout: %s\nstderr: %s",
			err, stdout.String(), stderr.String())
	}

	// Read each line of output as a remote
	result := make(map[string]string)
	scanner := bufio.NewScanner(&stdout)
	for scanner.Scan() {
		line := scanner.Text()
		split := strings.Split(line, "\t")

		if len(split) < 2 {
			return nil, fmt.Errorf("invalid response from git remote: %s", stdout.String())
		}

		remote := fmt.Sprintf("remote.%s", strings.TrimSpace(split[0]))
		if _, ok := result[remote]; !ok {
			// https://github.com/foo/bar.git (fetch) #=> https://github.com/foo/bar.git
			urlSplit := strings.Split(split[1], " ")
			result[remote] = strings.TrimSpace(urlSplit[0])
		}
	}

	return result, nil
}
<commit_msg>Move functions out of struct<commit_after>package archive

import (
	"bufio"
	"bytes"
	"fmt"
	"log"
	"os"
	"os/exec"
	"path/filepath"
	"strings"

	version "github.com/hashicorp/go-version"
)

// VCS is a struct that explains how to get the file list for a given
// VCS.
type VCS struct {
	Name string

	// Detect is a list of files/folders that if they exist, signal that
	// this VCS is the VCS in use.
	Detect []string

	// Files returns the files that are under version control for the
	// given path.
	Files VCSFilesFunc

	// Metadata returns arbitrary metadata about the underlying VCS for the
	// given path.
	Metadata VCSMetadataFunc

	// Preflight is a function to run before looking for VCS files.
	Preflight VCSPreflightFunc
}

// VCSList is the list of VCS we recognize.
var VCSList = []*VCS{
	&VCS{
		Name:      "git",
		Detect:    []string{".git/"},
		Preflight: gitPreflight,
		Files:     vcsFilesCmd("git", "ls-files"),
		Metadata:  gitMetadata,
	},
	&VCS{
		Name:   "hg",
		Detect: []string{".hg/"},
		Files:  vcsTrimCmd(vcsFilesCmd("hg", "locate", "-f", "--include", ".")),
	},
	&VCS{
		Name:   "svn",
		Detect: []string{".svn/"},
		Files:  vcsFilesCmd("svn", "ls"),
	},
}

// VCSFilesFunc is the callback invoked to return the files in the VCS.
//
// The return value should be paths relative to the given path.
type VCSFilesFunc func(string) ([]string, error)

// VCSMetadataFunc is the callback invoked to get arbitrary information about
// the current VCS.
//
// The return value should be a map of key-value pairs.
type VCSMetadataFunc func(string) (map[string]string, error)

// VCSPreflightFunc is a function that runs before VCS detection to be
// configured by the user. It may be used to check if pre-requisites (like the
// actual VCS) are installed or that a program is at the correct version. If an
// error is returned, the VCS will not be processed and the error will be
// returned up the stack.
//
// The given argument is the path where the VCS is running.
type VCSPreflightFunc func(string) error

// vcsDetect detects the VCS that is used for path.
func vcsDetect(path string) (*VCS, error) {
	for _, v := range VCSList {
		dir := path
		for {
			for _, f := range v.Detect {
				check := filepath.Join(dir, f)
				if _, err := os.Stat(check); err == nil {
					return v, nil
				}
			}

			lastDir := dir
			dir = filepath.Dir(dir)
			if dir == lastDir {
				break
			}
		}
	}

	return nil, nil
}

// vcsPreflight returns the metadata for the VCS directory path.
func vcsPreflight(path string) error {
	vcs, err := vcsDetect(path)
	if err != nil {
		return fmt.Errorf("error detecting VCS: %s", err)
	}
	if vcs == nil {
		return fmt.Errorf("no VCS found for path: %s", path)
	}

	if vcs.Preflight != nil {
		return vcs.Preflight(path)
	}

	return nil
}

// vcsFiles returns the files for the VCS directory path.
func vcsFiles(path string) ([]string, error) {
	vcs, err := vcsDetect(path)
	if err != nil {
		return nil, fmt.Errorf("error detecting VCS: %s", err)
	}
	if vcs == nil {
		return nil, fmt.Errorf("no VCS found for path: %s", path)
	}

	if vcs.Files != nil {
		return vcs.Files(path)
	}

	return nil, nil
}

// vcsFilesCmd creates a Files-compatible function that reads the files
// by executing the command in the repository path and returning each
// line in stdout.
func vcsFilesCmd(args ...string) VCSFilesFunc {
	return func(path string) ([]string, error) {
		var stderr, stdout bytes.Buffer

		cmd := exec.Command(args[0], args[1:]...)
		cmd.Dir = path
		cmd.Stdout = &stdout
		cmd.Stderr = &stderr
		if err := cmd.Run(); err != nil {
			return nil, fmt.Errorf(
				"error executing %s: %s",
				strings.Join(args, " "),
				err)
		}

		// Read each line of output as a path
		result := make([]string, 0, 100)
		scanner := bufio.NewScanner(&stdout)
		for scanner.Scan() {
			result = append(result, scanner.Text())
		}

		return result, nil
	}
}

// vcsTrimCmd trims the prefix from the paths returned by another VCSFilesFunc.
// This should be used to wrap another function if the return value is known
// to have full paths rather than relative paths
func vcsTrimCmd(f VCSFilesFunc) VCSFilesFunc {
	return func(path string) ([]string, error) {
		absPath, err := filepath.Abs(path)
		if err != nil {
			return nil, fmt.Errorf(
				"error expanding VCS path: %s", err)
		}

		// Now that we have the root path, get the inner files
		fs, err := f(path)
		if err != nil {
			return nil, err
		}

		// Trim the root path from the files
		result := make([]string, 0, len(fs))
		for _, f := range fs {
			if !strings.HasPrefix(f, absPath) {
				continue
			}

			f, err = filepath.Rel(absPath, f)
			if err != nil {
				return nil, fmt.Errorf(
					"error determining path: %s", err)
			}

			result = append(result, f)
		}

		return result, nil
	}
}

// vcsMetadata returns the metadata for the VCS directory path.
func vcsMetadata(path string) (map[string]string, error) {
	vcs, err := vcsDetect(path)
	if err != nil {
		return nil, fmt.Errorf("error detecting VCS: %s", err)
	}
	if vcs == nil {
		return nil, fmt.Errorf("no VCS found for path: %s", path)
	}

	if vcs.Metadata != nil {
		return vcs.Metadata(path)
	}

	return nil, nil
}

// gitBranch gets and returns the current git branch for the Git repository
// at the given path. It is assumed that the VCS is git.
func gitBranch(path string) (string, error) {
	var stderr, stdout bytes.Buffer

	cmd := exec.Command("git", "symbolic-ref", "--short", "HEAD")
	cmd.Dir = path
	cmd.Stdout = &stdout
	cmd.Stderr = &stderr
	if err := cmd.Run(); err != nil {
		return "", fmt.Errorf("error getting git branch: %s\nstdout: %s\nstderr: %s",
			err, stdout.String(), stderr.String())
	}

	branch := strings.TrimSpace(stdout.String())

	return branch, nil
}

// gitCommit gets the SHA of the latest commit for the Git repository at the
// given path. It is assumed that the VCS is git.
func gitCommit(path string) (string, error) {
	var stderr, stdout bytes.Buffer

	cmd := exec.Command("git", "log", "-n1", "--pretty=format:%H")
	cmd.Dir = path
	cmd.Stdout = &stdout
	cmd.Stderr = &stderr
	if err := cmd.Run(); err != nil {
		return "", fmt.Errorf("error getting git commit: %s\nstdout: %s\nstderr: %s",
			err, stdout.String(), stderr.String())
	}

	commit := strings.TrimSpace(stdout.String())

	return commit, nil
}

// gitRemotes gets and returns a map of all remotes for the Git repository. The
// map key is the name of the remote of the format "remote.NAME" and the value
// is the endpoint for the remote. It is assumed that the VCS is git.
func gitRemotes(path string) (map[string]string, error) {
	var stderr, stdout bytes.Buffer

	cmd := exec.Command("git", "remote", "-v")
	cmd.Dir = path
	cmd.Stdout = &stdout
	cmd.Stderr = &stderr
	if err := cmd.Run(); err != nil {
		return nil, fmt.Errorf("error getting git remotes: %s\nstdout: %s\nstderr: %s",
			err, stdout.String(), stderr.String())
	}

	// Read each line of output as a remote
	result := make(map[string]string)
	scanner := bufio.NewScanner(&stdout)
	for scanner.Scan() {
		line := scanner.Text()
		split := strings.Split(line, "\t")

		if len(split) < 2 {
			return nil, fmt.Errorf("invalid response from git remote: %s", stdout.String())
		}

		remote := fmt.Sprintf("remote.%s", strings.TrimSpace(split[0]))
		if _, ok := result[remote]; !ok {
			// https://github.com/foo/bar.git (fetch) #=> https://github.com/foo/bar.git
			urlSplit := strings.Split(split[1], " ")
			result[remote] = strings.TrimSpace(urlSplit[0])
		}
	}

	return result, nil
}

// gitPreflight is the pre-flight command that runs for Git-based VCSs
func gitPreflight(path string) error {
	var stderr, stdout bytes.Buffer

	cmd := exec.Command("git", "--version")
	cmd.Dir = path
	cmd.Stdout = &stdout
	cmd.Stderr = &stderr
	if err := cmd.Run(); err != nil {
		return fmt.Errorf("error getting git version: %s\nstdout: %s\nstderr: %s",
			err, stdout.String(), stderr.String())
	}

	// Check if the output is valid
	output := strings.Split(strings.TrimSpace(stdout.String()), " ")
	if len(output) < 1 {
		log.Printf("[WARN] could not extract version output from Git")
		return nil
	}

	// Parse the version
	gitv, err := version.NewVersion(output[len(output)-1])
	if err != nil {
		log.Printf("[WARN] could not parse version output from Git")
		return nil
	}

	constraint, err := version.NewConstraint("> 1.8")
	if err != nil {
		log.Printf("[WARN] could not create version constraint to check")
		return nil
	}
	if !constraint.Check(gitv) {
		return fmt.Errorf("git version (%s) is too old, please upgrade", gitv.String())
	}

	return nil
}

// gitMetadata is the function to parse and return Git metadata
func gitMetadata(path string) (map[string]string, error) {
	// Future-self note: Git is NOT threadsafe, so we cannot run these
	// operations in go routines or else you're going to have a really really
	// bad day and Panda.State == "Sad" :(

	branch, err := gitBranch(path)
	if err != nil {
		return nil, err
	}

	commit, err := gitCommit(path)
	if err != nil {
		return nil, err
	}

	remotes, err := gitRemotes(path)
	if err != nil {
		return nil, err
	}

	// Make the return result (we already know the size)
	result := make(map[string]string, 2+len(remotes))

	result["branch"] = branch
	result["commit"] = commit
	for remote, value := range remotes {
		result[remote] = value
	}

	return result, nil
}
<|endoftext|>


<commit_before>package sshvault

import (
	"bufio"
	"bytes"
	"crypto/x509"
	"encoding/base64"
	"encoding/pem"
	"fmt"
	"io/ioutil"
	"os"
	"strings"

	"github.com/ssh-vault/crypto/aead"
	"github.com/ssh-vault/crypto/oaep"
)

// View decrypts data and print it to stdout
func (v *vault) View() ([]byte, error) {
	file, err := os.Open(v.vault)
	if err != nil {
		return nil, err
	}
	defer file.Close()

	var (
		// ssh-vault;AES256;fingerprint
		header     []string
		rawPayload bytes.Buffer
	)

	scanner := bufio.NewScanner(file)
	scanner.Split(bufio.ScanLines)
	l := 1
	for scanner.Scan() {
		line := scanner.Text()
		if l == 1 {
			header = strings.Split(line, ";")
		} else {
			rawPayload.WriteString(line)
		}
		l++
	}

	// password, body
	payload := strings.Split(rawPayload.String(), ";")

	// use private key only
	if strings.HasSuffix(v.key, ".pub") {
		v.key = strings.Trim(v.key, ".pub")
	}

	keyFile, err := ioutil.ReadFile(v.key)
	if err != nil {
		return nil, fmt.Errorf("Error reading private key: %s", err)
	}

	block, _ := pem.Decode(keyFile)
	if block == nil || block.Type != "RSA PRIVATE KEY" {
		return nil, fmt.Errorf("No valid PEM (private key) data found")
	}

	if x509.IsEncryptedPEMBlock(block) {
		keyPassword, err := v.GetPassword()
		if err != nil {
			return nil, fmt.Errorf("Unable to get private key password, Decryption failed.")
		}

		block.Bytes, err = x509.DecryptPEMBlock(block, keyPassword)
		if err != nil {
			return nil, fmt.Errorf("Password incorrect, Decryption failed.")
		}
	}

	privateKey, err := x509.ParsePKCS1PrivateKey(block.Bytes)
	if err != nil {
		return nil, err
	}

	ciphertext, err := base64.StdEncoding.DecodeString(payload[0])
	if err != nil {
		return nil, err
	}

	v.Password, err = oaep.Decrypt(privateKey, ciphertext, []byte(""))
	if err != nil {
		return nil, fmt.Errorf("Decryption failed, use private key with fingerprint: %s", v.Fingerprint)
	}

	ciphertext, err = base64.StdEncoding.DecodeString(payload[1])
	if err != nil {
		return nil, err
	}

	// decrypt ciphertext using fingerprint as additionalData
	data, err := aead.Decrypt(v.Password, ciphertext, []byte(header[2]))
	if err != nil {
		return nil, err
	}
	return data, nil
}
<commit_msg>fixed message when trying to open a vault using incorrect key so that it can print the proper key<commit_after>package sshvault

import (
	"bufio"
	"bytes"
	"crypto/x509"
	"encoding/base64"
	"encoding/pem"
	"fmt"
	"io/ioutil"
	"os"
	"strings"

	"github.com/ssh-vault/crypto/aead"
	"github.com/ssh-vault/crypto/oaep"
)

// View decrypts data and print it to stdout
func (v *vault) View() ([]byte, error) {
	file, err := os.Open(v.vault)
	if err != nil {
		return nil, err
	}
	defer file.Close()

	var (
		// ssh-vault;AES256;fingerprint
		header     []string
		rawPayload bytes.Buffer
	)

	scanner := bufio.NewScanner(file)
	scanner.Split(bufio.ScanLines)
	l := 1
	for scanner.Scan() {
		line := scanner.Text()
		if l == 1 {
			header = strings.Split(line, ";")
		} else {
			rawPayload.WriteString(line)
		}
		l++
	}

	// password, body
	payload := strings.Split(rawPayload.String(), ";")

	// use private key only
	if strings.HasSuffix(v.key, ".pub") {
		v.key = strings.Trim(v.key, ".pub")
	}

	keyFile, err := ioutil.ReadFile(v.key)
	if err != nil {
		return nil, fmt.Errorf("Error reading private key: %s", err)
	}

	block, _ := pem.Decode(keyFile)
	if block == nil || block.Type != "RSA PRIVATE KEY" {
		return nil, fmt.Errorf("No valid PEM (private key) data found")
	}

	if x509.IsEncryptedPEMBlock(block) {
		keyPassword, err := v.GetPassword()
		if err != nil {
			return nil, fmt.Errorf("Unable to get private key password, Decryption failed.")
		}

		block.Bytes, err = x509.DecryptPEMBlock(block, keyPassword)
		if err != nil {
			return nil, fmt.Errorf("Password incorrect, Decryption failed.")
		}
	}

	privateKey, err := x509.ParsePKCS1PrivateKey(block.Bytes)
	if err != nil {
		return nil, err
	}

	ciphertext, err := base64.StdEncoding.DecodeString(payload[0])
	if err != nil {
		return nil, err
	}

	v.Password, err = oaep.Decrypt(privateKey, ciphertext, []byte(""))
	if err != nil {
		return nil, fmt.Errorf("Decryption failed, use private key with fingerprint: %s", header[2])
	}

	ciphertext, err = base64.StdEncoding.DecodeString(payload[1])
	if err != nil {
		return nil, err
	}

	// decrypt ciphertext using fingerprint as additionalData
	data, err := aead.Decrypt(v.Password, ciphertext, []byte(header[2]))
	if err != nil {
		return nil, err
	}
	return data, nil
}
<|endoftext|>


<commit_before>package scripts_test

import (
	. "github.com/onsi/ginkgo"
	. "github.com/onsi/gomega"

	"bytes"
	"encoding/json"
	"io/ioutil"
	"net"
	"os"
	"os/exec"
	"path/filepath"
)

var _ = Describe("VPC peering", func() {
	var examplePeerConfig = map[string]string{
		"peer_name":   "cheese",
		"account_id":  "123cheese",
		"vpc_id":      "vpc123cheese",
		"subnet_cidr": "0.0.0.0/32",
	}

	var files []string

	BeforeEach(func() {
		files, _ = filepath.Glob("../*.vpc_peering.json")
	})

	Context("config files", func() {
		It("have all required values", func() {
			for _, filename := range files {
				file, _ := ioutil.ReadFile(filename)

				var peers []map[string]string
				err := json.Unmarshal([]byte(file), &peers)
				Expect(err).To(BeNil(), "Config file: %s", filename)

				for _, peer := range peers {
					for k, _ := range examplePeerConfig {
						Expect(peer).To(HaveKeyWithValue(k, Not(BeEmpty())), "Config %s missing %s", filename, k)
					}
				}
			}
		})

		It("do not have intersecting CIDRs", func() {
			for _, filename := range files {
				file, _ := ioutil.ReadFile(filename)

				var peers []map[string]string
				err := json.Unmarshal([]byte(file), &peers)
				Expect(err).To(BeNil(), "Config file: %s", filename)

				var nets []*net.IPNet
				for _, peer := range peers {
					_, net, err := net.ParseCIDR(peer["subnet_cidr"])
					Expect(err).To(BeNil(), "Config file: %s", filename)

					for _, othernet := range nets {
						intersection := (net.Contains(othernet.IP) || othernet.Contains(net.IP))
						Expect(intersection).To(Equal(false), "%s has intersecting CIDRs: %s, %s", filename, net, othernet)
					}
					nets = append(nets, net)
				}
			}
		})
	})

	Context("tfvar output", func() {
		It("can be generated", func() {
			for _, filename := range files {
				command := exec.Command("./generate_vpc_peering_tfvars.rb", filename)
				var out bytes.Buffer
				command.Stdout = &out
				err := command.Run()
				Expect(err).To(BeNil(), "Config file: %s", filename)
				Expect(out.String()).ToNot(BeEmpty(), "Config file: %s", filename)
			}
		})

		It("is empty with no environment config", func() {
			command := exec.Command("./generate_vpc_peering_tfvars.rb", "cheese")
			var out bytes.Buffer
			command.Stdout = &out
			err := command.Run()
			Expect(err).To(BeNil())
			Expect(out.String()).To(BeEmpty())
		})

		It("crashes when not given a config", func() {
			command := exec.Command("./generate_vpc_peering_tfvars.rb")
			err := command.Run()
			Expect(err).ToNot(BeNil())
		})

		It("crashes when a field is missing", func() {
			for field, _ := range examplePeerConfig {
				var config = map[string]string{}
				for k, v := range examplePeerConfig {
					if k == field {
						continue
					}
					config[k] = v
				}

				file, err := ioutil.TempFile(os.TempDir(), "vpcpeer")
				Expect(err).To(BeNil())

				defer os.Remove(file.Name())

				data, _ := json.Marshal(config)
				err = ioutil.WriteFile(file.Name(), data, 0644)
				Expect(err).To(BeNil())

				command := exec.Command("./generate_vpc_peering_tfvars.rb", file.Name())
				var out bytes.Buffer
				command.Stdout = &out
				err = command.Run()
				Expect(err).ToNot(BeNil(), "Missing field: %s", field)
			}
		})
	})

	Context("ops file", func() {
		It("can be generated", func() {
			for _, filename := range files {
				command := exec.Command("./generate_vpc_peering_opsfile.rb", filename)
				var out bytes.Buffer
				command.Stdout = &out
				err := command.Run()
				Expect(err).To(BeNil(), "Config file: %s", filename)
				Expect(out.String()).ToNot(BeEmpty(), "Config file: %s", filename)
				Expect(out.String()).ToNot(Equal(`--- []
`), "Config file: %s", filename)
			}
		})

		It("has no operations with no environment config", func() {
			command := exec.Command("./generate_vpc_peering_opsfile.rb", "cheese")
			var out bytes.Buffer
			command.Stdout = &out
			err := command.Run()
			Expect(err).To(BeNil())
			Expect(out.String()).To(Equal(`--- []
`))
		})

		It("generates one operation per config", func() {
			file, err := ioutil.TempFile(os.TempDir(), "vpcpeer")
			Expect(err).To(BeNil())

			defer os.Remove(file.Name())

			data, _ := json.Marshal([]interface{}{examplePeerConfig})
			err = ioutil.WriteFile(file.Name(), data, 0644)
			Expect(err).To(BeNil())

			command := exec.Command("./generate_vpc_peering_opsfile.rb", file.Name())
			var out bytes.Buffer
			command.Stdout = &out
			err = command.Run()
			Expect(err).To(BeNil())
			Expect(out.String()).To(Equal(`---
- type: replace
  path: "/instance_groups/name=api/jobs/name=cloud_controller_ng/properties/cc/security_group_definitions?/-"
  value:
    name: vpc_peer_cheese
    rules:
    - protocol: all
      destination: 0.0.0.0/32
`))
		})

		It("crashes when not given a config", func() {
			command := exec.Command("./generate_vpc_peering_opsfile.rb")
			err := command.Run()
			Expect(err).ToNot(BeNil())
		})

		It("crashes when a field is missing", func() {
			for field, _ := range examplePeerConfig {
				var config = map[string]string{}
				for k, v := range examplePeerConfig {
					if k == field {
						continue
					}
					config[k] = v
				}

				file, err := ioutil.TempFile(os.TempDir(), "vpcpeer")
				Expect(err).To(BeNil())

				defer os.Remove(file.Name())

				data, _ := json.Marshal(config)
				err = ioutil.WriteFile(file.Name(), data, 0644)
				Expect(err).To(BeNil())

				command := exec.Command("./generate_vpc_peering_opsfile.rb", file.Name())
				var out bytes.Buffer
				command.Stdout = &out
				err = command.Run()
				Expect(err).ToNot(BeNil(), "Missing field: %s", field)
			}
		})
	})
})
<commit_msg>Refactor tests for VPC peering configuration<commit_after>package scripts_test

import (
	. "github.com/onsi/ginkgo"
	. "github.com/onsi/gomega"

	"bytes"
	"encoding/json"
	"io/ioutil"
	"net"
	"os"
	"os/exec"
	"path/filepath"
)

var _ = Describe("VPC peering", func() {

	var KEY_NAMES = []string{"peer_name", "account_id", "vpc_id", "subnet_cidr"}

	Describe("config files", func() {
		var files []string

		BeforeEach(func() {
			files, _ = filepath.Glob("../*.vpc_peering.json")
		})

		It("have all required values", func() {
			for _, filename := range files {
				file, _ := ioutil.ReadFile(filename)

				var peers []map[string]string
				err := json.Unmarshal([]byte(file), &peers)
				Expect(err).To(BeNil(), "Config file: %s", filename)

				for _, peer := range peers {
					for _, v := range KEY_NAMES {
						Expect(peer).To(HaveKeyWithValue(v, Not(BeEmpty())), "Config %s missing %s", filename, v)
					}
				}
			}
		})

		It("do not have intersecting CIDRs", func() {
			for _, filename := range files {
				file, _ := ioutil.ReadFile(filename)

				var peers []map[string]string
				err := json.Unmarshal([]byte(file), &peers)
				Expect(err).To(BeNil(), "Config file: %s", filename)

				var nets []*net.IPNet
				for _, peer := range peers {
					_, net, err := net.ParseCIDR(peer["subnet_cidr"])
					Expect(err).To(BeNil(), "Config file: %s", filename)

					for _, othernet := range nets {
						intersection := (net.Contains(othernet.IP) || othernet.Contains(net.IP))
						Expect(intersection).To(Equal(false), "%s has intersecting CIDRs: %s, %s", filename, net, othernet)
					}
					nets = append(nets, net)
				}
			}
		})

		Describe("tfvar output", func() {
			It("can be generated", func() {
				for _, filename := range files {
					command := exec.Command("./generate_vpc_peering_tfvars.rb", filename)
					var out bytes.Buffer
					command.Stdout = &out
					err := command.Run()
					Expect(err).To(BeNil(), "Config file: %s", filename)
					Expect(out.String()).ToNot(BeEmpty(), "Config file: %s", filename)
				}
			})
		})

		Describe("ops file", func() {
			It("can be generated", func() {
				for _, filename := range files {
					command := exec.Command("./generate_vpc_peering_opsfile.rb", filename)
					var out bytes.Buffer
					command.Stdout = &out
					err := command.Run()
					Expect(err).To(BeNil(), "Config file: %s", filename)
					Expect(out.String()).ToNot(BeEmpty(), "Config file: %s", filename)
					Expect(out.String()).ToNot(Equal(`--- []
`), "Config file: %s", filename)
				}
			})
		})
	})

	Describe("generate_vpc_peering_tfvars.rb", func() {
		var examplePeerConfig map[string]string

		BeforeEach(func() {
			examplePeerConfig = map[string]string{
				"peer_name":   "cheese",
				"account_id":  "123cheese",
				"vpc_id":      "vpc123cheese",
				"subnet_cidr": "0.0.0.0/32",
			}
		})

		Context("tfvar output", func() {
			It("is empty with no environment config", func() {
				command := exec.Command("./generate_vpc_peering_tfvars.rb", "cheese")
				var out bytes.Buffer
				command.Stdout = &out
				err := command.Run()
				Expect(err).To(BeNil())
				Expect(out.String()).To(BeEmpty())
			})

			It("crashes when not given a config", func() {
				command := exec.Command("./generate_vpc_peering_tfvars.rb")
				err := command.Run()
				Expect(err).ToNot(BeNil())
			})

			It("crashes when a field is missing", func() {
				for field, _ := range examplePeerConfig {
					var config = map[string]string{}
					for k, v := range examplePeerConfig {
						if k == field {
							continue
						}
						config[k] = v
					}

					file, err := ioutil.TempFile(os.TempDir(), "vpcpeer")
					Expect(err).To(BeNil())

					defer os.Remove(file.Name())

					data, _ := json.Marshal(config)
					err = ioutil.WriteFile(file.Name(), data, 0644)
					Expect(err).To(BeNil())

					command := exec.Command("./generate_vpc_peering_tfvars.rb", file.Name())
					var out bytes.Buffer
					command.Stdout = &out
					err = command.Run()
					Expect(err).ToNot(BeNil(), "Missing field: %s", field)
				}
			})
		})

		Context("ops file", func() {
			It("has no operations with no environment config", func() {
				command := exec.Command("./generate_vpc_peering_opsfile.rb", "cheese")
				var out bytes.Buffer
				command.Stdout = &out
				err := command.Run()
				Expect(err).To(BeNil())
				Expect(out.String()).To(Equal(`--- []
`))
			})

			It("generates one operation per config", func() {
				file, err := ioutil.TempFile(os.TempDir(), "vpcpeer")
				Expect(err).To(BeNil())

				defer os.Remove(file.Name())

				data, _ := json.Marshal([]interface{}{examplePeerConfig})
				err = ioutil.WriteFile(file.Name(), data, 0644)
				Expect(err).To(BeNil())

				command := exec.Command("./generate_vpc_peering_opsfile.rb", file.Name())
				var out bytes.Buffer
				command.Stdout = &out
				err = command.Run()
				Expect(err).To(BeNil())
				Expect(out.String()).To(Equal(`---
- type: replace
  path: "/instance_groups/name=api/jobs/name=cloud_controller_ng/properties/cc/security_group_definitions?/-"
  value:
    name: vpc_peer_cheese
    rules:
    - protocol: all
      destination: 0.0.0.0/32
`))
			})

			It("crashes when not given a config", func() {
				command := exec.Command("./generate_vpc_peering_opsfile.rb")
				err := command.Run()
				Expect(err).ToNot(BeNil())
			})

			It("crashes when a field is missing", func() {
				for field, _ := range examplePeerConfig {
					var config = map[string]string{}
					for k, v := range examplePeerConfig {
						if k == field {
							continue
						}
						config[k] = v
					}

					file, err := ioutil.TempFile(os.TempDir(), "vpcpeer")
					Expect(err).To(BeNil())

					defer os.Remove(file.Name())

					data, _ := json.Marshal(config)
					err = ioutil.WriteFile(file.Name(), data, 0644)
					Expect(err).To(BeNil())

					command := exec.Command("./generate_vpc_peering_opsfile.rb", file.Name())
					var out bytes.Buffer
					command.Stdout = &out
					err = command.Run()
					Expect(err).ToNot(BeNil(), "Missing field: %s", field)
				}
			})
		})
	})
})
<|endoftext|>


<commit_before>package main

import (
	"io"
	"log"
	"net"
	"os"
	"sync"
	"sync/atomic"
	"time"

	"github.com/gholt/ring"
)

type msgMap struct {
	lock    sync.RWMutex
	mapping map[uint64]ring.MsgUnmarshaller
}

func newMsgMap() *msgMap {
	return &msgMap{mapping: make(map[uint64]ring.MsgUnmarshaller)}
}

func (mm *msgMap) set(t uint64, f ring.MsgUnmarshaller) ring.MsgUnmarshaller {
	mm.lock.Lock()
	p := mm.mapping[t]
	mm.mapping[t] = f
	mm.lock.Unlock()
	return p
}

func (mm *msgMap) get(t uint64) ring.MsgUnmarshaller {
	mm.lock.RLock()
	f := mm.mapping[t]
	mm.lock.RUnlock()
	return f
}

type node struct {
	id uint64
}

func (n *node) NodeID() uint64 {
	return n.id
}

func (n *node) Active() bool {
	return true
}

func (n *node) Capacity() uint32 {
	return 1
}

func (n *node) TierValues() []int {
	return nil
}

func (n *node) Address() string {
	return ""
}

type ringPipe struct {
	localNodeID     uint64
	conn            net.Conn
	lock            sync.RWMutex
	msgMap          *msgMap
	logError        *log.Logger
	logWarning      *log.Logger
	typeBytes       int
	lengthBytes     int
	writeChan       chan ring.Msg
	writingDoneChan chan struct{}
	sendDrops       uint32
}

func NewRingPipe(localNodeID uint64, c net.Conn) *ringPipe {
	rp := &ringPipe{
		localNodeID:     localNodeID,
		conn:            c,
		msgMap:          newMsgMap(),
		logError:        log.New(os.Stderr, "", log.LstdFlags),
		logWarning:      log.New(os.Stderr, "", log.LstdFlags),
		typeBytes:       1,
		lengthBytes:     3,
		writeChan:       make(chan ring.Msg, 40),
		writingDoneChan: make(chan struct{}, 1),
	}
	return rp
}

func (rp *ringPipe) Version() int64 {
	return 1
}

func (rp *ringPipe) PartitionBitCount() uint16 {
	return 8
}

func (rp *ringPipe) ReplicaCount() int {
	return 2
}

func (rp *ringPipe) Nodes() []ring.Node {
	return []ring.Node{&node{id: 0}, &node{id: 1}, &node{id: 2}}
}

func (rp *ringPipe) Node(id uint64) ring.Node {
	for _, node := range rp.Nodes() {
		if node.NodeID() == id {
			return node
		}
	}
	return nil
}

func (rp *ringPipe) LocalNode() ring.Node {
	return &node{id: rp.localNodeID}
}

func (rp *ringPipe) Responsible(partition uint32) bool {
	// TODO: Testing push replication, so node 2 is responsible for everything
	// but we're putting everything into node 1.
	return rp.localNodeID == 2
}

func (rp *ringPipe) ResponsibleNodes(partition uint32) []ring.Node {
	return []ring.Node{&node{id: 2}, &node{id: 2}}
}

func (rp *ringPipe) Start() {
	go rp.reading()
	go rp.writing()
}

const _GLH_SEND_MSG_TIMEOUT = 1

func (rp *ringPipe) MaxMsgLength() uint64 {
	return 16 * 1024 * 1024
}

func (rp *ringPipe) SetMsgHandler(t uint64, h ring.MsgUnmarshaller) {
	rp.msgMap.set(t, h)
}

func (rp *ringPipe) MsgToNode(localNodeID uint64, m ring.Msg) {
	select {
	case rp.writeChan <- m:
	case <-time.After(_GLH_SEND_MSG_TIMEOUT * time.Second):
		atomic.AddUint32(&rp.sendDrops, 1)
	}
	m.Done()
}

func (rp *ringPipe) MsgToOtherReplicas(ringVersion int64, partition uint32, m ring.Msg) {
	// TODO: If ringVersion has changed, partition invalid, etc. return false
	select {
	case rp.writeChan <- m:
	case <-time.After(_GLH_SEND_MSG_TIMEOUT * time.Second):
		atomic.AddUint32(&rp.sendDrops, 1)
	}
	m.Done()
}

func (rp *ringPipe) reading() {
	b := make([]byte, rp.typeBytes+rp.lengthBytes)
	d := make([]byte, 65536)
	for {
		var n int
		var sn int
		var err error
		for n != len(b) {
			if err != nil {
				if n != 0 || err != io.EOF {
					rp.logError.Print("error reading msg", err)
				}
				return
			}
			sn, err = rp.conn.Read(b[n:])
			n += sn
		}
		if err != nil {
			rp.logError.Print("error reading msg start", err)
			return
		}
		var t uint64
		for i := 0; i < rp.typeBytes; i++ {
			t = (t << 8) | uint64(b[i])
		}
		var l uint64
		for i := 0; i < rp.lengthBytes; i++ {
			l = (l << 8) | uint64(b[rp.typeBytes+i])
		}
		f := rp.msgMap.get(t)
		if f != nil {
			_, err = f(rp.conn, l)
			if err != nil {
				rp.logError.Print("error reading msg content", err)
				return
			}
		} else {
			rp.logWarning.Printf("unknown msg type %d", t)
			for l > 0 {
				if err != nil {
					rp.logError.Print("err reading unknown msg content", err)
					return
				}
				if l >= uint64(len(d)) {
					sn, err = rp.conn.Read(d)
				} else {
					sn, err = rp.conn.Read(d[:l])
				}
				l -= uint64(sn)
			}
		}
	}
}

func (rp *ringPipe) writing() {
	b := make([]byte, rp.typeBytes+rp.lengthBytes)
	for {
		m := <-rp.writeChan
		if m == nil {
			break
		}
		t := m.MsgType()
		for i := rp.typeBytes - 1; i >= 0; i-- {
			b[i] = byte(t)
			t >>= 8
		}
		l := m.MsgLength()
		for i := rp.lengthBytes - 1; i >= 0; i-- {
			b[rp.typeBytes+i] = byte(l)
			l >>= 8
		}
		_, err := rp.conn.Write(b)
		if err != nil {
			rp.logError.Print("err writing msg", err)
			break
		}
		_, err = m.WriteContent(rp.conn)
		if err != nil {
			rp.logError.Print("err writing msg content", err)
			break
		}
	}
	rp.writingDoneChan <- struct{}{}
}
<commit_msg>Stubbed Stats method<commit_after>package main

import (
	"io"
	"log"
	"net"
	"os"
	"sync"
	"sync/atomic"
	"time"

	"github.com/gholt/ring"
)

type msgMap struct {
	lock    sync.RWMutex
	mapping map[uint64]ring.MsgUnmarshaller
}

func newMsgMap() *msgMap {
	return &msgMap{mapping: make(map[uint64]ring.MsgUnmarshaller)}
}

func (mm *msgMap) set(t uint64, f ring.MsgUnmarshaller) ring.MsgUnmarshaller {
	mm.lock.Lock()
	p := mm.mapping[t]
	mm.mapping[t] = f
	mm.lock.Unlock()
	return p
}

func (mm *msgMap) get(t uint64) ring.MsgUnmarshaller {
	mm.lock.RLock()
	f := mm.mapping[t]
	mm.lock.RUnlock()
	return f
}

type node struct {
	id uint64
}

func (n *node) NodeID() uint64 {
	return n.id
}

func (n *node) Active() bool {
	return true
}

func (n *node) Capacity() uint32 {
	return 1
}

func (n *node) TierValues() []int {
	return nil
}

func (n *node) Address() string {
	return ""
}

type ringPipe struct {
	localNodeID     uint64
	conn            net.Conn
	lock            sync.RWMutex
	msgMap          *msgMap
	logError        *log.Logger
	logWarning      *log.Logger
	typeBytes       int
	lengthBytes     int
	writeChan       chan ring.Msg
	writingDoneChan chan struct{}
	sendDrops       uint32
}

func NewRingPipe(localNodeID uint64, c net.Conn) *ringPipe {
	rp := &ringPipe{
		localNodeID:     localNodeID,
		conn:            c,
		msgMap:          newMsgMap(),
		logError:        log.New(os.Stderr, "", log.LstdFlags),
		logWarning:      log.New(os.Stderr, "", log.LstdFlags),
		typeBytes:       1,
		lengthBytes:     3,
		writeChan:       make(chan ring.Msg, 40),
		writingDoneChan: make(chan struct{}, 1),
	}
	return rp
}

func (rp *ringPipe) Version() int64 {
	return 1
}

func (rp *ringPipe) PartitionBitCount() uint16 {
	return 8
}

func (rp *ringPipe) ReplicaCount() int {
	return 2
}

func (rp *ringPipe) Nodes() []ring.Node {
	return []ring.Node{&node{id: 0}, &node{id: 1}, &node{id: 2}}
}

func (rp *ringPipe) Node(id uint64) ring.Node {
	for _, node := range rp.Nodes() {
		if node.NodeID() == id {
			return node
		}
	}
	return nil
}

func (rp *ringPipe) LocalNode() ring.Node {
	return &node{id: rp.localNodeID}
}

func (rp *ringPipe) Responsible(partition uint32) bool {
	// TODO: Testing push replication, so node 2 is responsible for everything
	// but we're putting everything into node 1.
	return rp.localNodeID == 2
}

func (rp *ringPipe) ResponsibleNodes(partition uint32) []ring.Node {
	return []ring.Node{&node{id: 2}, &node{id: 2}}
}

func (rp *ringPipe) Stats() *ring.RingStats {
	return &ring.RingStats{}
}

func (rp *ringPipe) Start() {
	go rp.reading()
	go rp.writing()
}

const _GLH_SEND_MSG_TIMEOUT = 1

func (rp *ringPipe) MaxMsgLength() uint64 {
	return 16 * 1024 * 1024
}

func (rp *ringPipe) SetMsgHandler(t uint64, h ring.MsgUnmarshaller) {
	rp.msgMap.set(t, h)
}

func (rp *ringPipe) MsgToNode(localNodeID uint64, m ring.Msg) {
	select {
	case rp.writeChan <- m:
	case <-time.After(_GLH_SEND_MSG_TIMEOUT * time.Second):
		atomic.AddUint32(&rp.sendDrops, 1)
	}
	m.Done()
}

func (rp *ringPipe) MsgToOtherReplicas(ringVersion int64, partition uint32, m ring.Msg) {
	// TODO: If ringVersion has changed, partition invalid, etc. return false
	select {
	case rp.writeChan <- m:
	case <-time.After(_GLH_SEND_MSG_TIMEOUT * time.Second):
		atomic.AddUint32(&rp.sendDrops, 1)
	}
	m.Done()
}

func (rp *ringPipe) reading() {
	b := make([]byte, rp.typeBytes+rp.lengthBytes)
	d := make([]byte, 65536)
	for {
		var n int
		var sn int
		var err error
		for n != len(b) {
			if err != nil {
				if n != 0 || err != io.EOF {
					rp.logError.Print("error reading msg", err)
				}
				return
			}
			sn, err = rp.conn.Read(b[n:])
			n += sn
		}
		if err != nil {
			rp.logError.Print("error reading msg start", err)
			return
		}
		var t uint64
		for i := 0; i < rp.typeBytes; i++ {
			t = (t << 8) | uint64(b[i])
		}
		var l uint64
		for i := 0; i < rp.lengthBytes; i++ {
			l = (l << 8) | uint64(b[rp.typeBytes+i])
		}
		f := rp.msgMap.get(t)
		if f != nil {
			_, err = f(rp.conn, l)
			if err != nil {
				rp.logError.Print("error reading msg content", err)
				return
			}
		} else {
			rp.logWarning.Printf("unknown msg type %d", t)
			for l > 0 {
				if err != nil {
					rp.logError.Print("err reading unknown msg content", err)
					return
				}
				if l >= uint64(len(d)) {
					sn, err = rp.conn.Read(d)
				} else {
					sn, err = rp.conn.Read(d[:l])
				}
				l -= uint64(sn)
			}
		}
	}
}

func (rp *ringPipe) writing() {
	b := make([]byte, rp.typeBytes+rp.lengthBytes)
	for {
		m := <-rp.writeChan
		if m == nil {
			break
		}
		t := m.MsgType()
		for i := rp.typeBytes - 1; i >= 0; i-- {
			b[i] = byte(t)
			t >>= 8
		}
		l := m.MsgLength()
		for i := rp.lengthBytes - 1; i >= 0; i-- {
			b[rp.typeBytes+i] = byte(l)
			l >>= 8
		}
		_, err := rp.conn.Write(b)
		if err != nil {
			rp.logError.Print("err writing msg", err)
			break
		}
		_, err = m.WriteContent(rp.conn)
		if err != nil {
			rp.logError.Print("err writing msg content", err)
			break
		}
	}
	rp.writingDoneChan <- struct{}{}
}
<|endoftext|>