<commit_before>package main

import (
	"bytes"
	"encoding/json"
	"fmt"
	"github.com/colinmarc/hdfs"
	"io/ioutil"
	"log"
	"net"
	"net/http"
	"runtime"
	"sort"
	"strings"
	"sync"
	"time"
)

const (
	// How parallel should we be when polling the resource manager?
	runningJobWorkers = 3

	// How parallel should we be when polling the history server? This is
	// primarily useful during the inital data backfill.
	finishedJobWorkers = 3

	// Maximum number of jobs to keep track of. All data is retained in memory
	// on the server, and the details for each job are sent to the browser.
	jobLimit = 5000

	// How many hours of history should we ask for from the job server?
	jobHistoryDuration = time.Hour * 24 * 7

	// The server will only keep partial data for finished jobs older than this.
	// Keeping partial data reduces server memory usage. If someone tries to
	// fetch the job from the web the full data will be requested again from the
	// history server.
	fullDataDuration = time.Hour * 24

	// How many pairs of start/finish times should we keep around for each job?
	taskLimit = 500
)

var countersToKeep = map[string]string{
	"FileSystemCounter.HDFS_BYTES_READ":    "hdfs.bytes_read",
	"FileSystemCounter.HDFS_BYTES_WRITTEN": "hdfs.bytes_written",
	"TaskCounter.REDUCE_SHUFFLE_BYTES":     "hdfs.bytes_shuffled",
	"TaskCounter.MAP_INPUT_RECORDS":        "task.map_records",
	"TaskCounter.REDUCE_INPUT_RECORDS":     "task.reduce_records",
}

var configToKeep = map[string]string{
	"mapreduce.input.fileinputformat.inputdir":    "input",
	"mapreduce.output.fileoutputformat.outputdir": "output",
	"cascading.app.name":                          "cascading.app.name",
}

// This forces us to be consistent about the keys used in the Jobs map.
type jobID string

func hadoopIDs(id string) (string, jobID) {
	app := strings.Replace(id, "job_", "application_", 1)
	job := strings.Replace(id, "application_", "job_", 1)
	return app, jobID(job)
}

func checkRedirect(req *http.Request, via []*http.Request) error {
	// Redirects are banned for two reasons:
	// 1. When a job is in the accepted state the JSON API redirects to an HTML page.
	// 2. When a job is finished The RM redirects to the history server over an
	//    address we may not be able to follow.
	return fmt.Errorf("No redirects allowed. Blocked redirect to %s.", req.URL)
}

// This is cheesy. We're getting intermittent "use of closed network connection"
// errors when talking to the RM and history server. Whenever that happens we'll
// create a new http client and move on with our lives.
var httpClient http.Client
var httpClientMutex sync.Mutex

func init() {
	httpClientMutex = sync.Mutex{}
	generateNewHTTPClient()
}

func generateNewHTTPClient() {
	httpClientMutex.Lock()
	httpClient = http.Client{
		Transport: &http.Transport{
			Proxy: http.ProxyFromEnvironment,
			Dial: (&net.Dialer{
				Timeout:   *httpTimeout,
				KeepAlive: 30 * time.Second,
			}).Dial,
			TLSHandshakeTimeout: 10 * time.Second,
		},
		CheckRedirect: checkRedirect,
		Timeout:       *httpTimeout,
	}
	httpClientMutex.Unlock()
}

func getJSON(url string, data interface{}) error {
	resp, err := httpClient.Get(url)
	if err != nil {
		if strings.Index(err.Error(), "use of closed network connection") != -1 {
			generateNewHTTPClient()
		}
		return err
	}

	if resp.StatusCode != 200 {
		return fmt.Errorf("Failed getJSON: %v (%v)", url, resp.Status)
	}

	defer resp.Body.Close()
	jsonBytes, err := ioutil.ReadAll(resp.Body)
	if err != nil {
		return err
	}

	return json.Unmarshal(jsonBytes, data)
}

type jobTracker struct {
	Jobs     map[jobID]*job
	rm       string
	hs       string
	running  chan *job
	finished chan *job
	backfill chan *job
	updates  chan *job
}

func newJobTracker(rmHost string, historyHost string) jobTracker {
	generateNewHTTPClient()
	jt := jobTracker{
		Jobs:     make(map[jobID]*job),
		rm:       rmHost,
		hs:       historyHost,
		running:  make(chan *job, 100),
		finished: make(chan *job, 100),
		backfill: make(chan *job, 100),
		updates:  make(chan *job, 100),
	}
	return jt
}

func (jt *jobTracker) Loop() {
	go jt.runningJobLoop()
	go jt.finishedJobLoop()
	go jt.cleanupLoop()
}

func (jt *jobTracker) runningJobLoop() {
	for x := 1; x <= runningJobWorkers; x++ {
		go func() {
			for job := range jt.running {
				_, jobID := hadoopIDs(job.Details.ID)
				err := jt.UpdateJob(job, true)
				if err != nil {
					log.Println(err)
					// If a job is brand new we won't be able to fetch details from
					// the RM, so we'll get an error and end up here. If the job
					// doesn't exist in the jobs map then we'll assume that's
					// why we're here and let the mostly-empty job data
					// propagate anyways.
					if _, exists := jt.Jobs[jobID]; exists {
						continue
					}
				}
				jt.Jobs[jobID] = job
				jt.updates <- job
			}
		}()
	}

	for _ = range time.Tick(*pollInterval) {
		running := &appsResp{}
		if err := jt.ListJobs(jt.rm, running); err != nil {
			log.Println("Error listing running jobs:", err)
			continue
		}

		log.Println("Running jobs:", len(running.Apps.App))
		log.Println("Jobs in cache:", len(jt.Jobs))
		log.Println("Goroutines:", runtime.NumGoroutine())
		for i := range running.Apps.App {
			job := &job{Details: &running.Apps.App[i], host: jt.rm}
			jt.running <- job
		}
	}

}

func (jt *jobTracker) finishedJobLoop() {
	for x := 1; x <= finishedJobWorkers; x++ {
		go func() {
			for {
				var job *job

				select {
				case job = <-jt.finished:
				case job = <-jt.backfill:
				}

				full := job.Details.FinishTime/1000 > time.Now().Unix()-24*60*60
				err := jt.UpdateJob(job, full)
				if err != nil {
					log.Println(err)
					continue
				}
				_, jobID := hadoopIDs(job.Details.ID)
				jt.Jobs[jobID] = job
				jt.updates <- job
			}
		}()
	}

	go func() {
		finished := &jobResp{}
		for {
			if err := jt.FinishedJobs(jt.hs, finished); err != nil {
				log.Println("Error listing backfill jobs:", err)
				time.Sleep(time.Second * 5)
			} else {
				break
			}
		}

		// This sorts in reverse order of start time so that we'll fill in the
		// details for the newest jobs first.
		sort.Sort(sort.Reverse(jobDetails(finished.Jobs.Job)))
		log.Println("Backfill jobs:", len(finished.Jobs.Job))
		for i := range finished.Jobs.Job {
			if i > jobLimit {
				break
			}
			jt.finished <- &job{Details: &finished.Jobs.Job[i], host: jt.hs}
		}
	}()

	for _ = range time.Tick(*pollInterval) {
		finished := &jobResp{}
		if err := jt.ListJobs(jt.hs, finished); err != nil {
			log.Println("Error listing finished jobs:", err)
			continue
		}

		log.Println("Finished jobs:", len(finished.Jobs.Job))
		for i, details := range finished.Jobs.Job {
			// Only fill in details for jobs we don't already have.
			_, jobID := hadoopIDs(details.ID)
			j, exists := jt.Jobs[jobID]
			if exists && j.Details.State == details.State {
				continue
			}
			jt.finished <- &job{Details: &finished.Jobs.Job[i], host: jt.hs}
		}
	}
}

func (jt *jobTracker) cleanupLoop() {
	for _ = range time.Tick(time.Second * 60) {
		details := make(jobDetails, 0)
		for _, job := range jt.Jobs {
			if job.Details.State == "SUCCEEDED" {
				details = append(details, *job.Details)
			}
		}

		before := len(jt.Jobs)

		// Enforce the job limit.
		if len(details) > jobLimit {
			sort.Sort(sort.Reverse(details))
			for _, d := range details[jobLimit:] {
				_, jobID := hadoopIDs(d.ID)
				delete(jt.Jobs, jobID)
			}
		}
		log.Printf("Forgot about %d jobs to stay near the limit.\n", len(jt.Jobs)-before)

		// Drop tasks and counters for old jobs since those are only visible in
		// details pages (and unlikely to be viewed).
		counter := 0
		for _, d := range details {
			cutoff := time.Now().Add(-fullDataDuration).Unix()
			if d.FinishTime/1000 < cutoff {
				_, jobID := hadoopIDs(d.ID)
				job := jt.Jobs[jobID]
				if !job.complete {
					continue
				}
				job.complete = false
				job.Tasks = nil
				job.Counters = nil<commit_msg>Range over the map directly.<commit_after>package main

import (
	"bytes"
	"encoding/json"
	"fmt"
	"github.com/colinmarc/hdfs"
	"io/ioutil"
	"log"
	"net"
	"net/http"
	"runtime"
	"sort"
	"strings"
	"sync"
	"time"
)

const (
	// How parallel should we be when polling the resource manager?
	runningJobWorkers = 3

	// How parallel should we be when polling the history server? This is
	// primarily useful during the inital data backfill.
	finishedJobWorkers = 3

	// Maximum number of jobs to keep track of. All data is retained in memory
	// on the server, and the details for each job are sent to the browser.
	jobLimit = 5000

	// How many hours of history should we ask for from the job server?
	jobHistoryDuration = time.Hour * 24 * 7

	// The server will only keep partial data for finished jobs older than this.
	// Keeping partial data reduces server memory usage. If someone tries to
	// fetch the job from the web the full data will be requested again from the
	// history server.
	fullDataDuration = time.Hour * 24

	// How many pairs of start/finish times should we keep around for each job?
	taskLimit = 500
)

var countersToKeep = map[string]string{
	"FileSystemCounter.HDFS_BYTES_READ":    "hdfs.bytes_read",
	"FileSystemCounter.HDFS_BYTES_WRITTEN": "hdfs.bytes_written",
	"TaskCounter.REDUCE_SHUFFLE_BYTES":     "hdfs.bytes_shuffled",
	"TaskCounter.MAP_INPUT_RECORDS":        "task.map_records",
	"TaskCounter.REDUCE_INPUT_RECORDS":     "task.reduce_records",
}

var configToKeep = map[string]string{
	"mapreduce.input.fileinputformat.inputdir":    "input",
	"mapreduce.output.fileoutputformat.outputdir": "output",
	"cascading.app.name":                          "cascading.app.name",
}

// This forces us to be consistent about the keys used in the Jobs map.
type jobID string

func hadoopIDs(id string) (string, jobID) {
	app := strings.Replace(id, "job_", "application_", 1)
	job := strings.Replace(id, "application_", "job_", 1)
	return app, jobID(job)
}

func checkRedirect(req *http.Request, via []*http.Request) error {
	// Redirects are banned for two reasons:
	// 1. When a job is in the accepted state the JSON API redirects to an HTML page.
	// 2. When a job is finished The RM redirects to the history server over an
	//    address we may not be able to follow.
	return fmt.Errorf("No redirects allowed. Blocked redirect to %s.", req.URL)
}

// This is cheesy. We're getting intermittent "use of closed network connection"
// errors when talking to the RM and history server. Whenever that happens we'll
// create a new http client and move on with our lives.
var httpClient http.Client
var httpClientMutex sync.Mutex

func init() {
	httpClientMutex = sync.Mutex{}
	generateNewHTTPClient()
}

func generateNewHTTPClient() {
	httpClientMutex.Lock()
	httpClient = http.Client{
		Transport: &http.Transport{
			Proxy: http.ProxyFromEnvironment,
			Dial: (&net.Dialer{
				Timeout:   *httpTimeout,
				KeepAlive: 30 * time.Second,
			}).Dial,
			TLSHandshakeTimeout: 10 * time.Second,
		},
		CheckRedirect: checkRedirect,
		Timeout:       *httpTimeout,
	}
	httpClientMutex.Unlock()
}

func getJSON(url string, data interface{}) error {
	resp, err := httpClient.Get(url)
	if err != nil {
		if strings.Index(err.Error(), "use of closed network connection") != -1 {
			generateNewHTTPClient()
		}
		return err
	}

	if resp.StatusCode != 200 {
		return fmt.Errorf("Failed getJSON: %v (%v)", url, resp.Status)
	}

	defer resp.Body.Close()
	jsonBytes, err := ioutil.ReadAll(resp.Body)
	if err != nil {
		return err
	}

	return json.Unmarshal(jsonBytes, data)
}

type jobTracker struct {
	Jobs     map[jobID]*job
	rm       string
	hs       string
	running  chan *job
	finished chan *job
	backfill chan *job
	updates  chan *job
}

func newJobTracker(rmHost string, historyHost string) jobTracker {
	generateNewHTTPClient()
	jt := jobTracker{
		Jobs:     make(map[jobID]*job),
		rm:       rmHost,
		hs:       historyHost,
		running:  make(chan *job, 100),
		finished: make(chan *job, 100),
		backfill: make(chan *job, 100),
		updates:  make(chan *job, 100),
	}
	return jt
}

func (jt *jobTracker) Loop() {
	go jt.runningJobLoop()
	go jt.finishedJobLoop()
	go jt.cleanupLoop()
}

func (jt *jobTracker) runningJobLoop() {
	for x := 1; x <= runningJobWorkers; x++ {
		go func() {
			for job := range jt.running {
				_, jobID := hadoopIDs(job.Details.ID)
				err := jt.UpdateJob(job, true)
				if err != nil {
					log.Println(err)
					// If a job is brand new we won't be able to fetch details from
					// the RM, so we'll get an error and end up here. If the job
					// doesn't exist in the jobs map then we'll assume that's
					// why we're here and let the mostly-empty job data
					// propagate anyways.
					if _, exists := jt.Jobs[jobID]; exists {
						continue
					}
				}
				jt.Jobs[jobID] = job
				jt.updates <- job
			}
		}()
	}

	for _ = range time.Tick(*pollInterval) {
		running := &appsResp{}
		if err := jt.ListJobs(jt.rm, running); err != nil {
			log.Println("Error listing running jobs:", err)
			continue
		}

		log.Println("Running jobs:", len(running.Apps.App))
		log.Println("Jobs in cache:", len(jt.Jobs))
		log.Println("Goroutines:", runtime.NumGoroutine())
		for i := range running.Apps.App {
			job := &job{Details: &running.Apps.App[i], host: jt.rm}
			jt.running <- job
		}
	}

}

func (jt *jobTracker) finishedJobLoop() {
	for x := 1; x <= finishedJobWorkers; x++ {
		go func() {
			for {
				var job *job

				select {
				case job = <-jt.finished:
				case job = <-jt.backfill:
				}

				full := job.Details.FinishTime/1000 > time.Now().Unix()-24*60*60
				err := jt.UpdateJob(job, full)
				if err != nil {
					log.Println(err)
					continue
				}
				_, jobID := hadoopIDs(job.Details.ID)
				jt.Jobs[jobID] = job
				jt.updates <- job
			}
		}()
	}

	go func() {
		finished := &jobResp{}
		for {
			if err := jt.FinishedJobs(jt.hs, finished); err != nil {
				log.Println("Error listing backfill jobs:", err)
				time.Sleep(time.Second * 5)
			} else {
				break
			}
		}

		// This sorts in reverse order of start time so that we'll fill in the
		// details for the newest jobs first.
		sort.Sort(sort.Reverse(jobDetails(finished.Jobs.Job)))
		log.Println("Backfill jobs:", len(finished.Jobs.Job))
		for i := range finished.Jobs.Job {
			if i > jobLimit {
				break
			}
			jt.finished <- &job{Details: &finished.Jobs.Job[i], host: jt.hs}
		}
	}()

	for _ = range time.Tick(*pollInterval) {
		finished := &jobResp{}
		if err := jt.ListJobs(jt.hs, finished); err != nil {
			log.Println("Error listing finished jobs:", err)
			continue
		}

		log.Println("Finished jobs:", len(finished.Jobs.Job))
		for i, details := range finished.Jobs.Job {
			// Only fill in details for jobs we don't already have.
			_, jobID := hadoopIDs(details.ID)
			j, exists := jt.Jobs[jobID]
			if exists && j.Details.State == details.State {
				continue
			}
			jt.finished <- &job{Details: &finished.Jobs.Job[i], host: jt.hs}
		}
	}
}

func (jt *jobTracker) cleanupLoop() {
	for _ = range time.Tick(time.Second * 60) {
		details := make(jobDetails, 0)
		for _, job := range jt.Jobs {
			if job.Details.State == "SUCCEEDED" {
				details = append(details, *job.Details)
			}
		}

		before := len(jt.Jobs)

		// Enforce the job limit.
		if len(details) > jobLimit {
			sort.Sort(sort.Reverse(details))
			for _, d := range details[jobLimit:] {
				_, jobID := hadoopIDs(d.ID)
				delete(jt.Jobs, jobID)
			}
		}
		log.Printf("Forgot about %d jobs to stay near the limit.\n", len(jt.Jobs)-before)

		// Drop tasks and counters for old jobs since those are only visible in
		// details pages (and unlikely to be viewed).
		counter := 0
		for _, job := range jt.Jobs {
			cutoff := time.Now().Add(-fullDataDuration).Unix()
			if job.Details.FinishTime/1000 < cutoff && job.complete {
				job.complete = false
				job.Tasks = nil
				job.Counters = nil
{'fraction_non_alphanumeric': 0.12743239828152642, 'fraction_numerical': 0.007455142784938085, 'mean_word_length': 3.1383006535947713, 'pattern_counts': {'":': 16, '<': 23, '<?xml version=': 0, '>': 9, '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': '1819217', 'n_tokens_mistral': 6196, 'n_tokens_neox': 5389, 'n_words': 1627}