<commit_before>package registry

import (
	"encoding/json"
	"github.com/coreos/go-etcd/etcd"
	"github.com/jwilder/go-dockerclient"
	"github.com/litl/galaxy/utils"
	"path"
	"strings"
	"time"
)

const (
	ETCD_ENTRY_ALREADY_EXISTS = 105
	ETCD_ENTRY_NOT_EXISTS     = 100
)

type ServiceConfig struct {
	Name    string
	Version string
	Env     map[string]string
}

type ServiceRegistry struct {
	EctdClient   *etcd.Client
	Client       *docker.Client
	EtcdHosts    string
	Env          string
	Pool         string
	HostIp       string
	Hostname     string
	TTL          uint64
	OutputBuffer *utils.OutputBuffer
}

type ServiceRegistration struct {
	ExternalIp   string    `json:"EXTERNAL_IP"`
	ExternalPort string    `json:"EXTERNAL_PORT"`
	InternalIp   string    `json:"INTERNAL_IP"`
	InternalPort string    `json:"INTERNAL_PORT"`
	ContainerID  string    `json:"CONTAINER_ID"`
	StartedAt    time.Time `json:"STARTED_AT"`
	Expires      time.Time `json:"-"`
	Path         string    `json:"-"`
}

func (s *ServiceRegistration) Equals(other ServiceRegistration) bool {
	return s.ExternalIp == other.ExternalIp &&
		s.ExternalPort == other.ExternalPort &&
		s.InternalIp == other.InternalIp &&
		s.InternalPort == other.InternalPort
}

func (r *ServiceRegistry) setHostValue(service string, key string, value string) error {
	_, err := r.EctdClient.Set("/"+r.Env+"/"+r.Pool+"/hosts/"+r.Hostname+"/"+
		service+"/"+key, value, 0)
	return err
}

func (r *ServiceRegistry) makeServiceRegistration(container *docker.Container) *ServiceRegistration {
	//FIXME: We're using the first found port and assuming it's tcp.
	//How should we handle a service that exposes multiple ports
	//as well as tcp vs udp ports.
	var externalPort, internalPort string
	for k, _ := range container.NetworkSettings.Ports {
		externalPort = k.Port()
		internalPort = externalPort
		break
	}

	serviceRegistration := ServiceRegistration{
		ExternalIp:   r.HostIp,
		ExternalPort: externalPort,
		InternalIp:   container.NetworkSettings.IPAddress,
		InternalPort: internalPort,
		ContainerID:  container.ID,
		StartedAt:    container.Created,
	}
	return &serviceRegistration
}

func (r *ServiceRegistry) RegisterService(container *docker.Container, serviceConfig *ServiceConfig) error {

	machines := strings.Split(r.EtcdHosts, ",")
	r.EctdClient = etcd.NewClient(machines)

	_, err := r.EctdClient.CreateDir("/"+r.Env+"/"+r.Pool+"/hosts", 0)
	if err != nil && err.(*etcd.EtcdError).ErrorCode != ETCD_ENTRY_ALREADY_EXISTS {
		return err
	}

	registrationPath := "/" + r.Env + "/" + r.Pool + "/hosts/" + r.Hostname + "/" + serviceConfig.Name
	registration, err := r.EctdClient.CreateDir(registrationPath, r.TTL)
	if err != nil {

		if err.(*etcd.EtcdError).ErrorCode != ETCD_ENTRY_ALREADY_EXISTS {
			return err
		}

		registration, err = r.EctdClient.UpdateDir(registrationPath, r.TTL)
		if err != nil {
			return err
		}
	}

	var existingRegistration ServiceRegistration
	existingJson, err := r.EctdClient.Get(registrationPath+"/location", false, false)
	if err != nil {
		if err.(*etcd.EtcdError).ErrorCode != ETCD_ENTRY_NOT_EXISTS {
			return err
		}
	} else {
		err = json.Unmarshal([]byte(existingJson.Node.Value), &existingRegistration)
		if err != nil {
			return err
		}

		if existingRegistration.StartedAt.After(container.Created) {
			return nil
		}
	}

	serviceRegistration := r.makeServiceRegistration(container)

	jsonReg, err := json.Marshal(serviceRegistration)
	if err != nil {
		return err
	}

	err = r.setHostValue(serviceConfig.Name, "location", string(jsonReg))
	if err != nil {
		return err
	}

	jsonReg, err = json.Marshal(serviceConfig.Env)
	if err != nil {
		return err
	}

	err = r.setHostValue(serviceConfig.Name, "environment", string(jsonReg))
	if err != nil {
		return err
	}

	statusLine := strings.Join([]string{
		container.ID[0:12],
		registrationPath,
		container.Config.Image,
		serviceRegistration.ExternalIp + ":" + serviceRegistration.ExternalPort,<commit_msg>Don't re-register the same service if it already exists<commit_after>package registry

import (
	"encoding/json"
	"github.com/coreos/go-etcd/etcd"
	"github.com/jwilder/go-dockerclient"
	"github.com/litl/galaxy/utils"
	"path"
	"strings"
	"time"
)

const (
	ETCD_ENTRY_ALREADY_EXISTS = 105
	ETCD_ENTRY_NOT_EXISTS     = 100
)

type ServiceConfig struct {
	Name    string
	Version string
	Env     map[string]string
}

type ServiceRegistry struct {
	EctdClient   *etcd.Client
	Client       *docker.Client
	EtcdHosts    string
	Env          string
	Pool         string
	HostIp       string
	Hostname     string
	TTL          uint64
	OutputBuffer *utils.OutputBuffer
}

type ServiceRegistration struct {
	ExternalIp   string    `json:"EXTERNAL_IP"`
	ExternalPort string    `json:"EXTERNAL_PORT"`
	InternalIp   string    `json:"INTERNAL_IP"`
	InternalPort string    `json:"INTERNAL_PORT"`
	ContainerID  string    `json:"CONTAINER_ID"`
	StartedAt    time.Time `json:"STARTED_AT"`
	Expires      time.Time `json:"-"`
	Path         string    `json:"-"`
}

func (s *ServiceRegistration) Equals(other ServiceRegistration) bool {
	return s.ExternalIp == other.ExternalIp &&
		s.ExternalPort == other.ExternalPort &&
		s.InternalIp == other.InternalIp &&
		s.InternalPort == other.InternalPort
}

func (r *ServiceRegistry) setHostValue(service string, key string, value string) error {
	_, err := r.EctdClient.Set("/"+r.Env+"/"+r.Pool+"/hosts/"+r.Hostname+"/"+
		service+"/"+key, value, 0)
	return err
}

func (r *ServiceRegistry) makeServiceRegistration(container *docker.Container) *ServiceRegistration {
	//FIXME: We're using the first found port and assuming it's tcp.
	//How should we handle a service that exposes multiple ports
	//as well as tcp vs udp ports.
	var externalPort, internalPort string
	for k, _ := range container.NetworkSettings.Ports {
		externalPort = k.Port()
		internalPort = externalPort
		break
	}

	serviceRegistration := ServiceRegistration{
		ExternalIp:   r.HostIp,
		ExternalPort: externalPort,
		InternalIp:   container.NetworkSettings.IPAddress,
		InternalPort: internalPort,
		ContainerID:  container.ID,
		StartedAt:    container.Created,
	}
	return &serviceRegistration
}

func (r *ServiceRegistry) RegisterService(container *docker.Container, serviceConfig *ServiceConfig) error {

	machines := strings.Split(r.EtcdHosts, ",")
	r.EctdClient = etcd.NewClient(machines)

	_, err := r.EctdClient.CreateDir("/"+r.Env+"/"+r.Pool+"/hosts", 0)
	if err != nil && err.(*etcd.EtcdError).ErrorCode != ETCD_ENTRY_ALREADY_EXISTS {
		return err
	}

	registrationPath := "/" + r.Env + "/" + r.Pool + "/hosts/" + r.Hostname + "/" + serviceConfig.Name
	registration, err := r.EctdClient.CreateDir(registrationPath, r.TTL)
	if err != nil {

		if err.(*etcd.EtcdError).ErrorCode != ETCD_ENTRY_ALREADY_EXISTS {
			return err
		}

		registration, err = r.EctdClient.UpdateDir(registrationPath, r.TTL)
		if err != nil {
			return err
		}
	}

	var existingRegistration ServiceRegistration
	existingJson, err := r.EctdClient.Get(registrationPath+"/location", false, false)
	if err != nil {
		if err.(*etcd.EtcdError).ErrorCode != ETCD_ENTRY_NOT_EXISTS {
			return err
		}
	} else {
		err = json.Unmarshal([]byte(existingJson.Node.Value), &existingRegistration)
		if err != nil {
			return err
		}

		if existingRegistration.StartedAt.After(container.Created) {
			return nil
		}
	}

	serviceRegistration := r.makeServiceRegistration(container)
	if serviceRegistration.Equals(existingRegistration) {
		return nil
	}

	jsonReg, err := json.Marshal(serviceRegistration)
	if err != nil {
		return err
	}

	err = r.setHostValue(serviceConfig.Name, "location", string(jsonReg))
	if err != nil {
		return err
	}

	jsonReg, err = json.Marshal(serviceConfig.Env)
	if err != nil {
		return err
	}

	err = r.setHostValue(serviceConfig.Name, "environment", string(jsonReg))
	if err != nil {
		return err
	}

	statusLine := strings.Join([]string{
		container.ID[0:12],
		registrationPath,
		container.Config.Image,
		serviceRegistration.ExternalIp + ":" + serviceRegistration.ExternalPort,
{'fraction_non_alphanumeric': 0.12031852681348762, 'fraction_numerical': 0.003235037949483638, 'mean_word_length': 4.039498432601881, 'pattern_counts': {'":': 2, '<': 3, '<?xml version=': 0, '>': 3, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '5370796', 'n_tokens_mistral': 2825, 'n_tokens_neox': 2636, 'n_words': 587}