<commit_before>package store

import (
	"database/sql"
	"flag"
	"fmt"
	"log"
	"time"

	_ "github.com/lib/pq"
	"github.com/lib/pq/hstore"
	"github.com/rcrowley/go-metrics"
)

type postgresStore struct {
	db *sql.DB
}

var schedulingInterval = flag.Duration("scheduling_interval", time.Minute, "run the scheduler this often.")

var noPriorFunctionsScheduledCounter = metrics.GetOrRegisterCounter("NoPriorFunctionsScheduled", nil)
var lastMaxPriorityErrorCounter = metrics.GetOrRegisterCounter("LastMaxPriorityError", nil)
var deleteExpiredFunctionsErrorCounter = metrics.GetOrRegisterCounter("DeleteExpiredFunctionsError", nil)
var countSchedluedTasksErrorCounter = metrics.GetOrRegisterCounter("CountSchedluedTasksError", nil)
var countScheduledFunctionsErrorCounter = metrics.GetOrRegisterCounter("CountScheduledFunctionsError", nil)
var unfilledScheduleCounter = metrics.GetOrRegisterCounter("UnfilledSchedule", nil)
var insertScheduledFunctionsErrorCounter = metrics.GetOrRegisterCounter("InsertScheduledFunctionsError", nil)
var emptyTaskFunctionCounter = metrics.GetOrRegisterCounter("EmptyTaskFunction", nil)

func openPostgres(db *sql.DB) Store {
	return &postgresStore{
		db: db,
	}
}

func (store *postgresStore) Close() {
	store.db.Close()
}

func insertTaskFunctions(tx *sql.Tx) error {
	var minTaskFunction, minPriority int
	row := tx.QueryRow("SELECT task_function, priority FROM scheduled_functions ORDER BY scheduled_time DESC, priority DESC, task_function DESC LIMIT 1")
	if err := row.Scan(&minTaskFunction, &minPriority); err == sql.ErrNoRows {
		log.Printf("no prior functions scheduled")
		noPriorFunctionsScheduledCounter.Inc(1)
	} else if err != nil {
		log.Printf("error finding max last priority: %v", err)
		lastMaxPriorityErrorCounter.Inc(1)
		return err
	}

	log.Printf("min task function: %v, min priority: %v", minTaskFunction, minPriority)

	if _, err := tx.Exec("DELETE FROM scheduled_functions WHERE expiration_time < now() OR measurements_remaining <= 0"); err != nil {
		log.Printf("error deleting expired task functions: %v", err)
		deleteExpiredFunctionsErrorCounter.Inc(1)
		return err
	}

	var toSchedule int
	row = tx.QueryRow("SELECT concurrent_functions - scheduled FROM (SELECT count(1) scheduled FROM scheduled_functions) AS c, scheduler_configuration")
	if err := row.Scan(&toSchedule); err != nil {
		log.Printf("error counting scheduled tasks: %v", err)
		countSchedluedTasksErrorCounter.Inc(1)
		return err
	}

	result, err := tx.Exec("INSERT INTO scheduled_functions (task_function, expiration_time, measurements_remaining, priority, scheduled_time) SELECT id, now() + max_duration_seconds * interval '1 second', max_measurements, priority, now() FROM task_functions WHERE enabled AND ((priority = $1 AND id > $2) OR priority > $1) ORDER BY priority, id LIMIT $3", minPriority, minTaskFunction, toSchedule)
	if err != nil {
		log.Printf("error inserting new schedules: %v", err)
		insertScheduledFunctionsErrorCounter.Inc(1)
		return err
	}
	rowsAffected, err := result.RowsAffected()
	if err != nil {
		log.Printf("error discovering number of affected rows: %v", err)
		countScheduledFunctionsErrorCounter.Inc(1)
		return err
	}
	toSchedule -= int(rowsAffected)
	result, err = tx.Exec("INSERT INTO scheduled_functions (task_function, expiration_time, measurements_remaining, priority, scheduled_time) SELECT id, now() + max_duration_seconds * interval '1 second', max_measurements, priority, now() FROM task_functions WHERE enabled ORDER BY priority, id LIMIT $1", toSchedule)
	if err != nil {
		log.Printf("error inserting new schedules: %v", err)
		insertScheduledFunctionsErrorCounter.Inc(1)
		return err
	}
	rowsAffected, err = result.RowsAffected()
	if err != nil {
		log.Printf("error discovering number of affected rows: %v", err)
		countScheduledFunctionsErrorCounter.Inc(1)
		return err
	}
	toSchedule -= int(rowsAffected)
	if toSchedule > 0 {
		log.Printf("unable to fill schedule")
		unfilledScheduleCounter.Inc(1)
	}
	return nil
}

func (store *postgresStore) ScheduleTaskFunctions() {
	schedule := func() {
		tx, err := store.db.Begin()
		if err != nil {
			log.Printf("error starting transaction: %v", err)
			return
		}
		if err := insertTaskFunctions(tx); err != nil {
			tx.Rollback()
			return
		}
		if err := tx.Commit(); err != nil {
			log.Printf("error committing transaction: %v", err)
			return
		}
	}

	schedule()
	for _ = range time.Tick(*schedulingInterval) {
		schedule()
	}
}

func (store *postgresStore) Tasks(taskRequests <-chan *TaskRequest) {
	updateTicker := time.Tick(*schedulingInterval)

	selectStmt, err := store.db.Prepare("SELECT task_functions.task_function FROM scheduled_functions JOIN task_functions ON scheduled_functions.task_function = id ORDER BY id")
	if err != nil {
		log.Fatalf("error preparing schedules select statement: %v", err)
	}

	refreshTaskFunctions := func() []string {
		rows, err := selectStmt.Query()
		if err != nil {
			log.Fatalf("error selecting schedules: %v", err)
		}
		var taskFunctions []string
		for rows.Next() {
			var taskFunction string
			if err := rows.Scan(&taskFunction); err != nil {
				log.Fatalf("error scanning task function: %v", err)
			}
			taskFunctions = append(taskFunctions, taskFunction)
		}
		return taskFunctions
	}

	currentTaskFunctions := refreshTaskFunctions()
	functionIdx := 0
	var currentTaskFunction string
	for {
		if len(currentTaskFunctions) > 0 {
			functionIdx = (functionIdx + 1) % len(currentTaskFunctions)
			currentTaskFunction = currentTaskFunctions[functionIdx]
		}

		select {
		case taskRequest := <-taskRequests:
			queryString := fmt.Sprintf("SELECT id, parameters FROM task_functions.%[1]s($1) ORDER BY random() LIMIT 1", currentTaskFunction)
			log.Println(queryString)
			hints := hstore.Hstore{
				Map: make(map[string]sql.NullString),
			}
			for k, v := range taskRequest.Hints {
				hints.Map[k] = sql.NullString{
					String: v,
					Valid:  true,
				}
			}
			row := store.db.QueryRow(queryString, hints)
			var task *Task
			var id int
			var parameters hstore.Hstore
			if err := row.Scan(&id, &parameters); err != nil {
				log.Printf("error scanning task parameters: %v", err)
			} else {
				task = new(Task)
				task.Id = id
				task.Parameters = parameters.Map
			}
			select {
			case taskRequest.Response <- task:<commit_msg>remove debugging<commit_after>package store

import (
	"database/sql"
	"flag"
	"fmt"
	"log"
	"time"

	_ "github.com/lib/pq"
	"github.com/lib/pq/hstore"
	"github.com/rcrowley/go-metrics"
)

type postgresStore struct {
	db *sql.DB
}

var schedulingInterval = flag.Duration("scheduling_interval", time.Minute, "run the scheduler this often.")

var noPriorFunctionsScheduledCounter = metrics.GetOrRegisterCounter("NoPriorFunctionsScheduled", nil)
var lastMaxPriorityErrorCounter = metrics.GetOrRegisterCounter("LastMaxPriorityError", nil)
var deleteExpiredFunctionsErrorCounter = metrics.GetOrRegisterCounter("DeleteExpiredFunctionsError", nil)
var countSchedluedTasksErrorCounter = metrics.GetOrRegisterCounter("CountSchedluedTasksError", nil)
var countScheduledFunctionsErrorCounter = metrics.GetOrRegisterCounter("CountScheduledFunctionsError", nil)
var unfilledScheduleCounter = metrics.GetOrRegisterCounter("UnfilledSchedule", nil)
var insertScheduledFunctionsErrorCounter = metrics.GetOrRegisterCounter("InsertScheduledFunctionsError", nil)
var emptyTaskFunctionCounter = metrics.GetOrRegisterCounter("EmptyTaskFunction", nil)

func openPostgres(db *sql.DB) Store {
	return &postgresStore{
		db: db,
	}
}

func (store *postgresStore) Close() {
	store.db.Close()
}

func insertTaskFunctions(tx *sql.Tx) error {
	var minTaskFunction, minPriority int
	row := tx.QueryRow("SELECT task_function, priority FROM scheduled_functions ORDER BY scheduled_time DESC, priority DESC, task_function DESC LIMIT 1")
	if err := row.Scan(&minTaskFunction, &minPriority); err == sql.ErrNoRows {
		log.Printf("no prior functions scheduled")
		noPriorFunctionsScheduledCounter.Inc(1)
	} else if err != nil {
		log.Printf("error finding max last priority: %v", err)
		lastMaxPriorityErrorCounter.Inc(1)
		return err
	}

	log.Printf("min task function: %v, min priority: %v", minTaskFunction, minPriority)

	if _, err := tx.Exec("DELETE FROM scheduled_functions WHERE expiration_time < now() OR measurements_remaining <= 0"); err != nil {
		log.Printf("error deleting expired task functions: %v", err)
		deleteExpiredFunctionsErrorCounter.Inc(1)
		return err
	}

	var toSchedule int
	row = tx.QueryRow("SELECT concurrent_functions - scheduled FROM (SELECT count(1) scheduled FROM scheduled_functions) AS c, scheduler_configuration")
	if err := row.Scan(&toSchedule); err != nil {
		log.Printf("error counting scheduled tasks: %v", err)
		countSchedluedTasksErrorCounter.Inc(1)
		return err
	}

	result, err := tx.Exec("INSERT INTO scheduled_functions (task_function, expiration_time, measurements_remaining, priority, scheduled_time) SELECT id, now() + max_duration_seconds * interval '1 second', max_measurements, priority, now() FROM task_functions WHERE enabled AND ((priority = $1 AND id > $2) OR priority > $1) ORDER BY priority, id LIMIT $3", minPriority, minTaskFunction, toSchedule)
	if err != nil {
		log.Printf("error inserting new schedules: %v", err)
		insertScheduledFunctionsErrorCounter.Inc(1)
		return err
	}
	rowsAffected, err := result.RowsAffected()
	if err != nil {
		log.Printf("error discovering number of affected rows: %v", err)
		countScheduledFunctionsErrorCounter.Inc(1)
		return err
	}
	toSchedule -= int(rowsAffected)
	result, err = tx.Exec("INSERT INTO scheduled_functions (task_function, expiration_time, measurements_remaining, priority, scheduled_time) SELECT id, now() + max_duration_seconds * interval '1 second', max_measurements, priority, now() FROM task_functions WHERE enabled ORDER BY priority, id LIMIT $1", toSchedule)
	if err != nil {
		log.Printf("error inserting new schedules: %v", err)
		insertScheduledFunctionsErrorCounter.Inc(1)
		return err
	}
	rowsAffected, err = result.RowsAffected()
	if err != nil {
		log.Printf("error discovering number of affected rows: %v", err)
		countScheduledFunctionsErrorCounter.Inc(1)
		return err
	}
	toSchedule -= int(rowsAffected)
	if toSchedule > 0 {
		log.Printf("unable to fill schedule")
		unfilledScheduleCounter.Inc(1)
	}
	return nil
}

func (store *postgresStore) ScheduleTaskFunctions() {
	schedule := func() {
		tx, err := store.db.Begin()
		if err != nil {
			log.Printf("error starting transaction: %v", err)
			return
		}
		if err := insertTaskFunctions(tx); err != nil {
			tx.Rollback()
			return
		}
		if err := tx.Commit(); err != nil {
			log.Printf("error committing transaction: %v", err)
			return
		}
	}

	schedule()
	for _ = range time.Tick(*schedulingInterval) {
		schedule()
	}
}

func (store *postgresStore) Tasks(taskRequests <-chan *TaskRequest) {
	updateTicker := time.Tick(*schedulingInterval)

	selectStmt, err := store.db.Prepare("SELECT task_functions.task_function FROM scheduled_functions JOIN task_functions ON scheduled_functions.task_function = id ORDER BY id")
	if err != nil {
		log.Fatalf("error preparing schedules select statement: %v", err)
	}

	refreshTaskFunctions := func() []string {
		rows, err := selectStmt.Query()
		if err != nil {
			log.Fatalf("error selecting schedules: %v", err)
		}
		var taskFunctions []string
		for rows.Next() {
			var taskFunction string
			if err := rows.Scan(&taskFunction); err != nil {
				log.Fatalf("error scanning task function: %v", err)
			}
			taskFunctions = append(taskFunctions, taskFunction)
		}
		return taskFunctions
	}

	currentTaskFunctions := refreshTaskFunctions()
	functionIdx := 0
	var currentTaskFunction string
	for {
		if len(currentTaskFunctions) > 0 {
			functionIdx = (functionIdx + 1) % len(currentTaskFunctions)
			currentTaskFunction = currentTaskFunctions[functionIdx]
		}

		select {
		case taskRequest := <-taskRequests:
			queryString := fmt.Sprintf("SELECT id, parameters FROM task_functions.%[1]s($1) ORDER BY random() LIMIT 1", currentTaskFunction)
			hints := hstore.Hstore{
				Map: make(map[string]sql.NullString),
			}
			for k, v := range taskRequest.Hints {
				hints.Map[k] = sql.NullString{
					String: v,
					Valid:  true,
				}
			}
			row := store.db.QueryRow(queryString, hints)
			var task *Task
			var id int
			var parameters hstore.Hstore
			if err := row.Scan(&id, &parameters); err != nil {
				log.Printf("error scanning task parameters: %v", err)
			} else {
				task = new(Task)
				task.Id = id
				task.Parameters = parameters.Map
			}
			select {
			case taskRequest.Response <- task:
{'fraction_non_alphanumeric': 0.1117987780687138, 'fraction_numerical': 0.004126001745616123, 'mean_word_length': 4.97629208155524, 'pattern_counts': {'":': 0, '<': 13, '<?xml version=': 0, '>': 11, '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': '6155062', 'n_tokens_mistral': 4181, 'n_tokens_neox': 3811, 'n_words': 1041}