<commit_before>package main

import (
	"database/sql"
	"flag"
	"fmt"
	_ "github.com/go-sql-driver/mysql"
	"log"
	"os"
	"runtime"
)

const (
	// Timeout for DNS queries
	timeout = 3 * 1e9

	// maximum number of attempts for a query
	maxAttempts = 3
)

var (
	pending         = make(chan *job, 100)
	finished        = make(chan *job, 100)
	done            = make(chan bool)
	workersCount    = 1
	referenceServer = "8.8.8.8"
	connection      string
	domainArg       string
)

func main() {
	databaseArg := flag.String("database", "database.yml", "Path to file containing the database configuration")
	flag.StringVar(&domainArg, "domains", "domains.txt", "Path to file containing the domain list")
	flag.StringVar(&geoDbPath, "geodb", "GeoLite2-City.mmdb", "Path to GeoDB database")
	flag.StringVar(&referenceServer, "reference", referenceServer, "The nameserver that every other is compared with")
	workersPerCore := flag.Int("workers-per-core", 8, "Number of worker routines per CPU core")
	flag.Parse()

	dnsClient.ReadTimeout = timeout

	environment := os.Getenv("RAILS_ENV")
	if environment == "" {
		environment = "development"
	}

	// read domain list
	if err := readDomains(domainArg); err != nil {
		fmt.Println("unable to read domain list")
		panic(err)
	}

	// load database configuration
	connection = databasePath(*databaseArg, environment)

	// check the GeoDB
	location(referenceServer)

	// Use all cores
	cpuCount := runtime.NumCPU()
	runtime.GOMAXPROCS(cpuCount)
	workersCount = *workersPerCore * cpuCount
	fmt.Println("Starting", workersCount, "workers")

	// Get results from the reference nameserver
	res, err := resolveDomains(referenceServer)
	if err != nil {
		panic(err)
	}
	expectedResults = res

	// Start result writer
	go resultWriter()

	// Start workers
	for i := 0; i < workersCount; i++ {
		go worker()
	}

	createJobs()

	// wait for resultWriter to finish
	<-done
}

func createJobs() {
	currentId := 0
	batchSize := 1000

	// Open SQL connection
	db, err := sql.Open("mysql", connection)
	if err != nil {
		panic(err)
	}
	defer db.Close()

	for {
		// Read the next batch
		rows, err := db.Query("SELECT id, ip FROM nameservers WHERE id > ? LIMIT ?", currentId, batchSize)
		if err != nil {
			panic(err)
		}

		found := 0
		for rows.Next() {
			j := new(job)

			// get RawBytes from data
			err = rows.Scan(&j.id, &j.address)
			if err != nil {
				panic(err)
			}
			pending <- j
			currentId = j.id
			found += 1
		}
		rows.Close()

		// Last batch?
		if found < batchSize {
			close(pending)
			return
		}
	}
}

func worker() {
	for {
		job := <-pending
		if job != nil {
			executeJob(job)
			finished <- job
		} else {
			log.Println("received all jobs")
			finished <- nil
			return
		}
	}
}

func resultWriter() {
	// Open SQL connection
	db, err := sql.Open("mysql", connection)
	if err != nil {
		panic(err)
	}
	defer db.Close()

	stm, err := db.Prepare(
		"UPDATE nameservers SET name=?, state=?, error=?, version=?, checked_at=NOW(), country_id=?, city=?," +
			"state_changed_at = (CASE WHEN ? != state THEN NOW() ELSE state_changed_at END )" +
			"WHERE id=?")
	defer stm.Close()

	doneCount := 0
	for doneCount < workersCount {
		res := <-finished
		// log.Println("finished job", res.id)
		if res == nil {
			doneCount++
			log.Println("worker terminated")
		} else {
			log.Println(res)
			stm.Exec(res.name, res.state, res.err, res.version, res.country, res.city, res.state, res.id)
		}
	}
	done <- true
}

// consumes a job and writes the result in the given job
func executeJob(job *job) {
	// log.Println("received job", job.id)

	// GeoDB lookup
	job.country, job.city = location(job.address)

	// Run the check
	err := check(job)
	job.name = ptrName(job.address)

	// query the bind version
	if err == nil || err.Error() != "i/o timeout" {
		job.version = version(job.address)<commit_msg>Use wait groups<commit_after>package main

import (
	"database/sql"
	"flag"
	"fmt"
	_ "github.com/go-sql-driver/mysql"
	"log"
	"os"
	"runtime"
	"sync"
)

const (
	// Timeout for DNS queries
	timeout = 3 * 1e9

	// maximum number of attempts for a query
	maxAttempts = 3
)

var (
	pending         = make(chan *job, 100)
	finished        = make(chan *job, 100)
	done            sync.WaitGroup
	workersCount    = 1
	referenceServer = "8.8.8.8"
	connection      string
	domainArg       string
)

func main() {
	databaseArg := flag.String("database", "database.yml", "Path to file containing the database configuration")
	flag.StringVar(&domainArg, "domains", "domains.txt", "Path to file containing the domain list")
	flag.StringVar(&geoDbPath, "geodb", "GeoLite2-City.mmdb", "Path to GeoDB database")
	flag.StringVar(&referenceServer, "reference", referenceServer, "The nameserver that every other is compared with")
	workersPerCore := flag.Int("workers-per-core", 8, "Number of worker routines per CPU core")
	flag.Parse()

	dnsClient.ReadTimeout = timeout

	environment := os.Getenv("RAILS_ENV")
	if environment == "" {
		environment = "development"
	}

	// read domain list
	if err := readDomains(domainArg); err != nil {
		fmt.Println("unable to read domain list")
		panic(err)
	}

	// load database configuration
	connection = databasePath(*databaseArg, environment)

	// check the GeoDB
	location(referenceServer)

	// Use all cores
	cpuCount := runtime.NumCPU()
	runtime.GOMAXPROCS(cpuCount)

	workersCount = *workersPerCore * cpuCount
	fmt.Println("Starting", workersCount, "workers")
	done.Add(workersCount)

	// Get results from the reference nameserver
	res, err := resolveDomains(referenceServer)
	if err != nil {
		panic(err)
	}
	expectedResults = res

	// Start result writer
	go resultWriter()

	// Start workers
	for i := 0; i < workersCount; i++ {
		go worker()
	}

	createJobs()

	// wait for workers to finish
	done.Wait()

	close(finished)
}

func createJobs() {
	currentId := 0
	batchSize := 1000

	// Open SQL connection
	db, err := sql.Open("mysql", connection)
	if err != nil {
		panic(err)
	}
	defer db.Close()

	for {
		// Read the next batch
		rows, err := db.Query("SELECT id, ip FROM nameservers WHERE id > ? LIMIT ?", currentId, batchSize)
		if err != nil {
			panic(err)
		}

		found := 0
		for rows.Next() {
			j := new(job)

			// get RawBytes from data
			err = rows.Scan(&j.id, &j.address)
			if err != nil {
				panic(err)
			}
			pending <- j
			currentId = j.id
			found += 1
		}
		rows.Close()

		// Last batch?
		if found < batchSize {
			close(pending)
			return
		}
	}
}

func worker() {
	for job := range pending {
		executeJob(job)
		finished <- job
	}
	done.Done()
}

func resultWriter() {
	// Open SQL connection
	db, err := sql.Open("mysql", connection)
	if err != nil {
		panic(err)
	}
	defer db.Close()

	stm, err := db.Prepare(
		"UPDATE nameservers SET name=?, state=?, error=?, version=?, checked_at=NOW(), country_id=?, city=?," +
			"state_changed_at = (CASE WHEN ? != state THEN NOW() ELSE state_changed_at END )" +
			"WHERE id=?")
	defer stm.Close()

	for res := range finished {
		log.Println(res)
		stm.Exec(res.name, res.state, res.err, res.version, res.country, res.city, res.state, res.id)
	}
}

// consumes a job and writes the result in the given job
func executeJob(job *job) {
	// log.Println("received job", job.id)

	// GeoDB lookup
	job.country, job.city = location(job.address)

	// Run the check
	err := check(job)
	job.name = ptrName(job.address)

	// query the bind version
	if err == nil || err.Error() != "i/o timeout" {
		job.version = version(job.address)
{'fraction_non_alphanumeric': 0.14037200588786297, 'fraction_numerical': 0.006824568446407065, 'mean_word_length': 3.270857142857143, 'pattern_counts': {'":': 0, '<': 17, '<?xml version=': 0, '>': 5, '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': '2569246', 'n_tokens_mistral': 2820, 'n_tokens_neox': 2549, 'n_words': 736}