<commit_before>package main

import (
	"errors"
	"flag"
	"fmt"
	golog "log"
	"os"
	"os/signal"
	"strings"
	"sync"
	"time"

	"strconv"
	"syscall"

	docker "github.com/fsouza/go-dockerclient"
	"github.com/litl/galaxy/commander"
	"github.com/litl/galaxy/config"
	"github.com/litl/galaxy/discovery"
	"github.com/litl/galaxy/log"
	"github.com/litl/galaxy/registry"
	"github.com/litl/galaxy/runtime"
	"github.com/litl/galaxy/utils"
)

var (
	stopCutoff      int64
	apps            []string
	env             string
	pool            string
	registryURL     string
	loop            bool
	hostIP          string
	dns             string
	shuttleAddr     string
	debug           bool
	runOnce         bool
	version         bool
	buildVersion    string
	serviceRegistry *registry.ServiceRegistry
	configStore     *config.Store
	serviceRuntime  *runtime.ServiceRuntime
	workerChans     map[string]chan string
	wg              sync.WaitGroup
	signalsChan     chan os.Signal
)

func initOrDie() {

	if registryURL == "" {
		log.Fatalf("ERROR: Registry URL not specified. Use '-registry redis://127.0.0.1:6379' or set 'GALAXY_REGISTRY_URL'")
	}

	serviceRegistry = registry.NewServiceRegistry(
		registry.DefaultTTL,
	)
	serviceRegistry.Connect(registryURL)

	configStore = config.NewStore(
		registry.DefaultTTL,
	)

	configStore.Connect(registryURL)

	serviceRuntime = runtime.NewServiceRuntime(serviceRegistry, dns, hostIP)

	apps, err := configStore.ListAssignments(env, pool)
	if err != nil {
		log.Fatalf("ERROR: Could not retrieve service configs for /%s/%s: %s", env, pool, err)
	}

	workerChans = make(map[string]chan string)
	for _, app := range apps {
		appCfg, err := configStore.GetApp(app, env)
		if err != nil {
			log.Fatalf("ERROR: Could not retrieve service config for /%s/%s: %s", env, pool, err)
		}

		workerChans[appCfg.Name] = make(chan string)
	}

	signalsChan = make(chan os.Signal, 1)
	signal.Notify(signalsChan, os.Interrupt, os.Kill, syscall.SIGTERM)
	go deregisterHost(signalsChan)
}

func ensureEnv() {
	envs, err := configStore.ListEnvs()
	if err != nil {
		log.Fatalf("ERROR: Could not check envs: %s", err)
	}

	if strings.TrimSpace(env) == "" {
		log.Fatalf("ERROR: Need an env.  Use '-env <env>'. Existing envs are: %s.", strings.Join(envs, ","))
	}
}

func ensurePool() {

	pools, err := configStore.ListPools(env)
	if err != nil {
		log.Fatalf("ERROR: Could not check pools: %s", err)
	}

	if strings.TrimSpace(pool) == "" {
		log.Fatalf("ERROR: Need a pool.  Use '-pool <pool>'. Existing pools are: %s", strings.Join(pools, ","))
	}
}

func pullImageAsync(appCfg config.AppConfig, errChan chan error) {
	// err logged via pullImage
	_, err := pullImage(&appCfg)
	if err != nil {
		errChan <- err
		return
	}
	errChan <- nil
}

func pullImage(appCfg *config.AppConfig) (*docker.Image, error) {

	image, err := serviceRuntime.InspectImage(appCfg.Version())
	if image != nil && image.ID == appCfg.VersionID() || appCfg.VersionID() == "" {
		return image, nil
	}

	log.Printf("Pulling %s version %s\n", appCfg.Name, appCfg.Version())
	image, err = serviceRuntime.PullImage(appCfg.Version(),
		appCfg.VersionID())
	if image == nil || err != nil {
		log.Errorf("ERROR: Could not pull image %s: %s",
			appCfg.Version(), err)
		return nil, err
	}

	if image.ID != appCfg.VersionID() && len(appCfg.VersionID()) > 12 {
		log.Errorf("ERROR: Pulled image for %s does not match expected ID. Expected: %s: Got: %s",
			appCfg.Version(),
			image.ID[0:12], appCfg.VersionID()[0:12])
		return nil, errors.New(fmt.Sprintf("failed to pull image ID %s", appCfg.VersionID()[0:12]))
	}

	log.Printf("Pulled %s\n", appCfg.Version())
	return image, nil
}

func startService(appCfg *config.AppConfig, logStatus bool) {

	desired, err := commander.Balanced(configStore, hostIP, appCfg.Name, env, pool)
	if err != nil {
		log.Errorf("ERROR: Could not determine instance count: %s", err)
		return
	}

	running, err := serviceRuntime.InstanceCount(appCfg.Name, strconv.FormatInt(appCfg.ID(), 10))
	if err != nil {
		log.Errorf("ERROR: Could not determine running instance count: %s", err)
		return
	}

	for i := 0; i < desired-running; i++ {
		container, err := serviceRuntime.Start(env, pool, appCfg)
		if err != nil {
			log.Errorf("ERROR: Could not start containers: %s", err)
			return
		}

		log.Printf("Started %s version %s as %s\n", appCfg.Name, appCfg.Version(), container.ID[0:12])

		err = serviceRuntime.StopOldVersion(appCfg, 1)
		if err != nil {
			log.Errorf("ERROR: Could not stop containers: %s", err)
		}
	}

	running, err = serviceRuntime.InstanceCount(appCfg.Name, strconv.FormatInt(appCfg.ID(), 10))
	if err != nil {
		log.Errorf("ERROR: Could not determine running instance count: %s", err)
		return
	}

	for i := 0; i < running-desired; i++ {
		err := serviceRuntime.Stop(appCfg)
		if err != nil {
			log.Errorf("ERROR: Could not stop container: %s", err)
		}
	}

}

func heartbeatHost() {
	wg.Add(1)

	_, err := configStore.CreatePool(pool, env)
	if err != nil {
		log.Fatalf("ERROR: Unabled to create pool %s: %s", pool, err)
	}

	defer wg.Done()
	for {
		configStore.UpdateHost(env, pool, config.HostInfo{
			HostIP: hostIP,
		})

		if !loop {<commit_msg>commander: fix old containers getting leaked<commit_after>package main

import (
	"errors"
	"flag"
	"fmt"
	golog "log"
	"os"
	"os/signal"
	"strings"
	"sync"
	"time"

	"strconv"
	"syscall"

	docker "github.com/fsouza/go-dockerclient"
	"github.com/litl/galaxy/commander"
	"github.com/litl/galaxy/config"
	"github.com/litl/galaxy/discovery"
	"github.com/litl/galaxy/log"
	"github.com/litl/galaxy/registry"
	"github.com/litl/galaxy/runtime"
	"github.com/litl/galaxy/utils"
)

var (
	stopCutoff      int64
	apps            []string
	env             string
	pool            string
	registryURL     string
	loop            bool
	hostIP          string
	dns             string
	shuttleAddr     string
	debug           bool
	runOnce         bool
	version         bool
	buildVersion    string
	serviceRegistry *registry.ServiceRegistry
	configStore     *config.Store
	serviceRuntime  *runtime.ServiceRuntime
	workerChans     map[string]chan string
	wg              sync.WaitGroup
	signalsChan     chan os.Signal
)

func initOrDie() {

	if registryURL == "" {
		log.Fatalf("ERROR: Registry URL not specified. Use '-registry redis://127.0.0.1:6379' or set 'GALAXY_REGISTRY_URL'")
	}

	serviceRegistry = registry.NewServiceRegistry(
		registry.DefaultTTL,
	)
	serviceRegistry.Connect(registryURL)

	configStore = config.NewStore(
		registry.DefaultTTL,
	)

	configStore.Connect(registryURL)

	serviceRuntime = runtime.NewServiceRuntime(serviceRegistry, dns, hostIP)

	apps, err := configStore.ListAssignments(env, pool)
	if err != nil {
		log.Fatalf("ERROR: Could not retrieve service configs for /%s/%s: %s", env, pool, err)
	}

	workerChans = make(map[string]chan string)
	for _, app := range apps {
		appCfg, err := configStore.GetApp(app, env)
		if err != nil {
			log.Fatalf("ERROR: Could not retrieve service config for /%s/%s: %s", env, pool, err)
		}

		workerChans[appCfg.Name] = make(chan string)
	}

	signalsChan = make(chan os.Signal, 1)
	signal.Notify(signalsChan, os.Interrupt, os.Kill, syscall.SIGTERM)
	go deregisterHost(signalsChan)
}

func ensureEnv() {
	envs, err := configStore.ListEnvs()
	if err != nil {
		log.Fatalf("ERROR: Could not check envs: %s", err)
	}

	if strings.TrimSpace(env) == "" {
		log.Fatalf("ERROR: Need an env.  Use '-env <env>'. Existing envs are: %s.", strings.Join(envs, ","))
	}
}

func ensurePool() {

	pools, err := configStore.ListPools(env)
	if err != nil {
		log.Fatalf("ERROR: Could not check pools: %s", err)
	}

	if strings.TrimSpace(pool) == "" {
		log.Fatalf("ERROR: Need a pool.  Use '-pool <pool>'. Existing pools are: %s", strings.Join(pools, ","))
	}
}

func pullImageAsync(appCfg config.AppConfig, errChan chan error) {
	// err logged via pullImage
	_, err := pullImage(&appCfg)
	if err != nil {
		errChan <- err
		return
	}
	errChan <- nil
}

func pullImage(appCfg *config.AppConfig) (*docker.Image, error) {

	image, err := serviceRuntime.InspectImage(appCfg.Version())
	if image != nil && image.ID == appCfg.VersionID() || appCfg.VersionID() == "" {
		return image, nil
	}

	log.Printf("Pulling %s version %s\n", appCfg.Name, appCfg.Version())
	image, err = serviceRuntime.PullImage(appCfg.Version(),
		appCfg.VersionID())
	if image == nil || err != nil {
		log.Errorf("ERROR: Could not pull image %s: %s",
			appCfg.Version(), err)
		return nil, err
	}

	if image.ID != appCfg.VersionID() && len(appCfg.VersionID()) > 12 {
		log.Errorf("ERROR: Pulled image for %s does not match expected ID. Expected: %s: Got: %s",
			appCfg.Version(),
			image.ID[0:12], appCfg.VersionID()[0:12])
		return nil, errors.New(fmt.Sprintf("failed to pull image ID %s", appCfg.VersionID()[0:12]))
	}

	log.Printf("Pulled %s\n", appCfg.Version())
	return image, nil
}

func startService(appCfg *config.AppConfig, logStatus bool) {

	desired, err := commander.Balanced(configStore, hostIP, appCfg.Name, env, pool)
	if err != nil {
		log.Errorf("ERROR: Could not determine instance count: %s", err)
		return
	}

	running, err := serviceRuntime.InstanceCount(appCfg.Name, strconv.FormatInt(appCfg.ID(), 10))
	if err != nil {
		log.Errorf("ERROR: Could not determine running instance count: %s", err)
		return
	}

	for i := 0; i < desired-running; i++ {
		container, err := serviceRuntime.Start(env, pool, appCfg)
		if err != nil {
			log.Errorf("ERROR: Could not start containers: %s", err)
			return
		}

		log.Printf("Started %s version %s as %s\n", appCfg.Name, appCfg.Version(), container.ID[0:12])

		err = serviceRuntime.StopOldVersion(appCfg, 1)
		if err != nil {
			log.Errorf("ERROR: Could not stop containers: %s", err)
		}
	}

	running, err = serviceRuntime.InstanceCount(appCfg.Name, strconv.FormatInt(appCfg.ID(), 10))
	if err != nil {
		log.Errorf("ERROR: Could not determine running instance count: %s", err)
		return
	}

	for i := 0; i < running-desired; i++ {
		err := serviceRuntime.Stop(appCfg)
		if err != nil {
			log.Errorf("ERROR: Could not stop container: %s", err)
		}
	}

	err = serviceRuntime.StopAllButCurrentVersion(appCfg)
	if err != nil {
		log.Errorf("ERROR: Could not stop old containers: %s", err)
	}

}

func heartbeatHost() {
	wg.Add(1)

	_, err := configStore.CreatePool(pool, env)
	if err != nil {
		log.Fatalf("ERROR: Unabled to create pool %s: %s", pool, err)
	}

	defer wg.Done()
	for {
		configStore.UpdateHost(env, pool, config.HostInfo{
			HostIP: hostIP,
		})

		if !loop {
{'fraction_non_alphanumeric': 0.14469055993915772, 'fraction_numerical': 0.006654624964350223, 'mean_word_length': 3.636403702071397, 'pattern_counts': {'":': 0, '<': 15, '<?xml version=': 0, '>': 9, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 2, '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': '7181389', 'n_tokens_mistral': 3966, 'n_tokens_neox': 3704, 'n_words': 936}