<commit_before>package boxrunner

import (
	"fmt"
	"github.com/armon/consul-api"
	"github.com/christianberg/boxrunner/statemachine"
	"github.com/fsouza/go-dockerclient"
	"io"
	"log"
	"net/http"
	"os"
	"time"
)

type BoxRunner struct {
	Service string
	ID      string
	port    string
	consul  *consulapi.Client
	logger  *log.Logger
	dock    *docker.Client
	lock    *consulapi.KVPair
}

type BoxRunnerOptions struct {
	ConsulAddress string
	ConsulClient  *consulapi.Client
	Logger        *log.Logger
}

var DefaultOptions = &BoxRunnerOptions{
	ConsulAddress: "localhost:8500",
}

func NewBoxRunner(service_name string, options *BoxRunnerOptions) (runner *BoxRunner, err error) {
	if options == nil {
		options = DefaultOptions
	}
	completeOptions(options)

	logger := options.Logger
	dock, err := docker.NewClient("tcp://0.0.0.0:2375")
	if err != nil {
		logger.Printf("Could not initialize Docker client: %s", err)
		return
	}

	hostname, err := os.Hostname()
	if err != nil {
		logger.Printf("Could not determine hostname: %s", err)
		return
	}

	port := os.Getenv("PORT")
	runner_id := fmt.Sprintf("boxrunner-%v-%v", hostname, port)
	logger.Printf("This is boxrunner: %v", runner_id)

	runner = &BoxRunner{
		Service: service_name,
		ID:      runner_id,
		consul:  options.ConsulClient,
		logger:  options.Logger,
		dock:    dock,
		port:    port,
		lock: &consulapi.KVPair{
			Key:   service_name,
			Value: []byte(runner_id),
		},
	}
	return
}

func (b *BoxRunner) findOrCreateSession() (string, error) {
	sessions, _, err := b.consul.Session().List(nil)
	if err != nil {
		b.logger.Printf("Could not list existing sessions: %v", err)
		return "", err
	}
	for _, session := range sessions {
		if session.Name == b.ID {
			b.logger.Printf("Found existing session: %v\n", session.ID)
			return session.ID, nil
		}
	}

	err = b.consul.Agent().CheckRegister(&consulapi.AgentCheckRegistration{
		Name: b.ID,
		AgentServiceCheck: consulapi.AgentServiceCheck{
			Script:   fmt.Sprintf("curl -sf http://localhost:%v/health", b.port),
			Interval: "5s",
		},
	})
	if err != nil {
		b.logger.Printf("ERROR: Could not register boxrunner healthcheck: %v", err)
		return "", err
	}

	session_entry := &consulapi.SessionEntry{
		Name:      b.ID,
		LockDelay: 5 * time.Second,
		Checks:    []string{"serfHealth", b.ID},
	}
	session_id, _, err := b.consul.Session().Create(session_entry, nil)
	if err != nil {
		b.logger.Printf("ERROR: Could not create session: %v", err)
		return "", err
	}
	b.logger.Printf("INFO: Session created (ID: %v)\n", session_id)
	return session_id, nil
}

func (b *BoxRunner) waitForLockChange(predicate func(string) bool) (err error) {
	query_options := &consulapi.QueryOptions{
		WaitIndex: 0,
	}
	for {
		lock_status, meta, err := b.consul.KV().Get(b.lock.Key, query_options)
		if err != nil {
			b.logger.Printf("ERROR: Cannot check lock status: %s", err)
			return err
		}
		if lock_status == nil || predicate(lock_status.Session) {
			b.logger.Println("INFO: Lock was released")
			return nil
		}
		query_options.WaitIndex = meta.LastIndex
	}
}

func (b *BoxRunner) Run() (success bool, error error) {
	http.HandleFunc("/health", func(w http.ResponseWriter, _ *http.Request) {
		io.WriteString(w, "OK")
	})

	go http.ListenAndServe("127.0.0.1:"+b.port, nil)

	machine := &statemachine.Machine{
		Handlers: map[string]statemachine.Handler{},
		Logger:   b.logger,
	}

	machine.AddState("INIT", func() string {
		session_id, err := b.findOrCreateSession()
		if err != nil {
			return "FAILED"
		}
		b.lock.Session = session_id

		if err := b.dock.Ping(); err != nil {
			b.logger.Printf("ERROR: Cannot ping docker server: %v", err)
			return "FAILED"
		}
		return "DISCOVER"
	})

	machine.AddState("DISCOVER", func() string {
		return "COMPETE"
	})

	machine.AddState("COMPETE", func() string {
		b.logger.Printf("INFO: Trying to acquire lock for %s\n", b.Service)
		success, _, err := b.consul.KV().Acquire(b.lock, nil)
		if err != nil {
			b.logger.Printf("ERROR: Could not acquire lock: %s", err)
			return "FAILED"
		}
		if success {
			b.logger.Println("INFO: Lock acquired!")
			return "START"
		} else {
			lock_status, _, _ := b.consul.KV().Get(b.lock.Key, nil)
			if lock_status != nil {
				b.logger.Printf("INFO: Lock is already taken by: %s", lock_status.Value)
			}
			return "SLEEP"
		}
	})

	machine.AddState("START", func() string {<commit_msg>Check system status in DISCOVER state<commit_after>package boxrunner

import (
	"fmt"
	"github.com/armon/consul-api"
	"github.com/christianberg/boxrunner/statemachine"
	"github.com/fsouza/go-dockerclient"
	"io"
	"log"
	"net/http"
	"os"
	"time"
)

type BoxRunner struct {
	Service string
	ID      string
	port    string
	consul  *consulapi.Client
	logger  *log.Logger
	dock    *docker.Client
	lock    *consulapi.KVPair
}

type BoxRunnerOptions struct {
	ConsulAddress string
	ConsulClient  *consulapi.Client
	Logger        *log.Logger
}

var DefaultOptions = &BoxRunnerOptions{
	ConsulAddress: "localhost:8500",
}

func NewBoxRunner(service_name string, options *BoxRunnerOptions) (runner *BoxRunner, err error) {
	if options == nil {
		options = DefaultOptions
	}
	completeOptions(options)

	logger := options.Logger
	dock, err := docker.NewClient("tcp://0.0.0.0:2375")
	if err != nil {
		logger.Printf("Could not initialize Docker client: %s", err)
		return
	}

	hostname, err := os.Hostname()
	if err != nil {
		logger.Printf("Could not determine hostname: %s", err)
		return
	}

	port := os.Getenv("PORT")
	runner_id := fmt.Sprintf("boxrunner-%v-%v", hostname, port)
	logger.Printf("This is boxrunner: %v", runner_id)

	runner = &BoxRunner{
		Service: service_name,
		ID:      runner_id,
		consul:  options.ConsulClient,
		logger:  options.Logger,
		dock:    dock,
		port:    port,
		lock: &consulapi.KVPair{
			Key:   service_name,
			Value: []byte(runner_id),
		},
	}
	return
}

func (b *BoxRunner) findOrCreateSession() (string, error) {
	sessions, _, err := b.consul.Session().List(nil)
	if err != nil {
		b.logger.Printf("Could not list existing sessions: %v", err)
		return "", err
	}
	for _, session := range sessions {
		if session.Name == b.ID {
			b.logger.Printf("Found existing session: %v\n", session.ID)
			return session.ID, nil
		}
	}

	err = b.consul.Agent().CheckRegister(&consulapi.AgentCheckRegistration{
		Name: b.ID,
		AgentServiceCheck: consulapi.AgentServiceCheck{
			Script:   fmt.Sprintf("curl -sf http://localhost:%v/health", b.port),
			Interval: "5s",
		},
	})
	if err != nil {
		b.logger.Printf("ERROR: Could not register boxrunner healthcheck: %v", err)
		return "", err
	}

	session_entry := &consulapi.SessionEntry{
		Name:      b.ID,
		LockDelay: 5 * time.Second,
		Checks:    []string{"serfHealth", b.ID},
	}
	session_id, _, err := b.consul.Session().Create(session_entry, nil)
	if err != nil {
		b.logger.Printf("ERROR: Could not create session: %v", err)
		return "", err
	}
	b.logger.Printf("INFO: Session created (ID: %v)\n", session_id)
	return session_id, nil
}

func (b *BoxRunner) waitForLockChange(predicate func(string) bool) (err error) {
	query_options := &consulapi.QueryOptions{
		WaitIndex: 0,
	}
	for {
		lock_status, meta, err := b.consul.KV().Get(b.lock.Key, query_options)
		if err != nil {
			b.logger.Printf("ERROR: Cannot check lock status: %s", err)
			return err
		}
		if lock_status == nil || predicate(lock_status.Session) {
			b.logger.Println("INFO: Lock was released")
			return nil
		}
		query_options.WaitIndex = meta.LastIndex
	}
}

func (b *BoxRunner) Run() (success bool, error error) {
	http.HandleFunc("/health", func(w http.ResponseWriter, _ *http.Request) {
		io.WriteString(w, "OK")
	})

	go http.ListenAndServe("127.0.0.1:"+b.port, nil)

	machine := &statemachine.Machine{
		Handlers: map[string]statemachine.Handler{},
		Logger:   b.logger,
	}

	machine.AddState("INIT", func() string {
		return "DISCOVER"
	})

	machine.AddState("DISCOVER", func() string {
		session_id, err := b.findOrCreateSession()
		if err != nil {
			return "FAILED"
		}
		b.lock.Session = session_id

		if err := b.dock.Ping(); err != nil {
			b.logger.Printf("ERROR: Cannot ping docker server: %v", err)
			return "FAILED"
		}
		return "COMPETE"
	})

	machine.AddState("COMPETE", func() string {
		b.logger.Printf("INFO: Trying to acquire lock for %s\n", b.Service)
		success, _, err := b.consul.KV().Acquire(b.lock, nil)
		if err != nil {
			b.logger.Printf("ERROR: Could not acquire lock: %s", err)
			return "FAILED"
		}
		if success {
			b.logger.Println("INFO: Lock acquired!")
			return "START"
		} else {
			lock_status, _, _ := b.consul.KV().Get(b.lock.Key, nil)
			if lock_status != nil {
				b.logger.Printf("INFO: Lock is already taken by: %s", lock_status.Value)
			}
			return "SLEEP"
		}
	})

	machine.AddState("START", func() string {
{'fraction_non_alphanumeric': 0.15053518560692325, 'fraction_numerical': 0.004782509678888636, 'mean_word_length': 3.7475675675675677, 'pattern_counts': {'":': 0, '<': 3, '<?xml version=': 0, '>': 3, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 4, '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': '3789811', 'n_tokens_mistral': 3436, 'n_tokens_neox': 3201, 'n_words': 747}