<commit_before>package worker

import (
	"errors"
	"flag"
	"fmt"
	"os"
	"strings"
	"time"

	"github.com/jinzhu/gorm"
	"github.com/qor/qor"
	"github.com/qor/qor/admin"
)

var (
	jobDB            *gorm.DB
	jobId            uint64
	defaultWorkerSet = &WorkerSet{Name: "default"}
	workerSets       = []*WorkerSet{defaultWorkerSet}
	queuers          = map[string]Queuer{}
	DefaultJobCli    = strings.Join(os.Args, " ")
)

func init() {
	flag.Uint64Var(&jobId, "job-id", 0, "qor job id")
}

func SetJobDB(db *gorm.DB) {
	jobDB = db
}

func NewWorker(name string, handle func(job *Job) error, queuer Queuer) (w *Worker) {
	return defaultWorkerSet.NewWorker(name, handle, queuer)
}

func SetAdmin(admin *admin.Admin) {
	ws := defaultWorkerSet.Workers
	defaultWorkerSet = NewWorkerSet(defaultWorkerSet.Name, "/workers", "", admin)
	defaultWorkerSet.Workers = ws
}

func RegisterQueuer(name string, queuer Queuer) {
	queuers[name] = queuer
}

func Listen() {
	flag.Parse()

	if jobId > 0 {
		var job Job
		if err := jobDB.Where("_id = ?", jobId).Find(&job).Error; err != nil {
			fmt.Fprintln(os.Stderr, err)
			os.Exit(1)
		}
		w, err := job.GetWorker()
		if err != nil {
			fmt.Fprintln(os.Stderr, err)
			os.Exit(1)
		}

		w.Run(&job)
	} else {
		for _, queuer := range queuers {
			go func() {
				for {
					jobId, err := queuer.Dequeue()
					if err != nil {
						fmt.Println("qor.worker.dequeue.error:", err)
					} else {
						go RunJob(jobId)
					}
				}
			}()
		}
	}
}

type WorkerSet struct {
	Name    string
	Workers []*Worker
}

func NewWorkerSet(name, router, tmplDir string, a *admin.Admin) (ws *WorkerSet) {
	ws = &WorkerSet{Name: name}
	workerSets = append(workerSets, ws)
	a.GetRouter().Get(router, func(ctx *admin.Context) {})
	// template register
	// menu register
	// Job resource register

	return
}

func (ws *WorkerSet) NewWorker(name string, handle func(job *Job) error, queuer Queuer) (w *Worker) {
	w = &Worker{
		Name:   name,
		Handle: handle,
		Queuer: queuer,
		set:    ws,
	}

	ws.Workers = append(ws.Workers, w)

	return
}

type Worker struct {
	Name   string
	Queuer Queuer
	Config *qor.Config
	set    *WorkerSet

	Handle    func(job *Job) error
	OnStart   func(job *Job) error
	OnKill    func(job *Job) error
	OnSuccess func(job *Job)
	OnFailed  func(job *Job)
}

// TODO: use docker
func (w *Worker) Run(job *Job) (err error) {
	if err = job.SavePID(); err != nil {
		fmt.Fprintln(os.Stderr, err)
		os.Exit(1)
	}
	logger, err := job.GetLogger()
	if err != nil {
		fmt.Println("can't get job logger")
		return
	}

	fmt.Fprintf(logger, "to run job (%d) with pid (%d)", job.Id, job.PID)

	if err = job.UpdateStatus(JobRunning); err != nil {
		fmt.Fprintln(os.Stderr, err)
		os.Exit(1)
	}

	if w.OnStart != nil {
		if err = w.OnStart(job); err != nil {
			logger.Write([]byte("worker.onstart: " + err.Error() + "\n"))

			if err = job.UpdateStatus(JobFailed); err != nil {
				fmt.Fprintln(os.Stderr, err)
				os.Exit(1)
			}

			fmt.Fprintln(os.Stderr, err)
			os.Exit(1)
		}
	}

	if err = w.Handle(job); err != nil {
		if err = job.UpdateStatus(JobFailed); err != nil {
			fmt.Fprintln(os.Stderr, err)
			os.Exit(1)
		}

		if w.OnFailed != nil {
			w.OnFailed(job)
		}

		logger.Write([]byte("worker.hanlde: " + err.Error() + "\n"))
		fmt.Fprintln(os.Stderr, err)
		os.Exit(1)
	} else if w.OnSuccess != nil {
		if err = job.UpdateStatus(JobRun); err != nil {
			fmt.Fprintln(os.Stderr, err)
		}

		w.OnSuccess(job)
	}

	fmt.Fprintf(logger, "finish job (%d) with pid (%d)", job.Id, job.PID)

	return
}

var ErrJobRun = errors.New("job is already run")

func (w *Worker) Kill(job *Job) (err error) {
	if w.OnKill != nil {
		if err = w.OnKill(job); err != nil {
			return
		}
	}

	switch job.Status {
	case JobToRun:
		err = w.Queuer.Purge(job)
	case JobRunning:
		if job.PID == 0 {
			return errors.New("pid is zero")
		}

		var process *os.Process
		process, err = os.FindProcess(job.PID)
		if err != nil {
			return
		}

		// err = process.Signal(syscall.SIGUSR1)
		err = process.Kill()
	case JobRun:
		return ErrJobRun
	}

	if err == nil {
		err = job.UpdateStatus(JobKilled)
	}

	return
}

func (w *Worker) NewJob(interval uint64, startAt time.Time) (job *Job, err error) {
	return w.NewJobWithCli(interval, startAt, DefaultJobCli)
}

func (w *Worker) NewJobWithCli(interval uint64, startAt time.Time, cli string) (job *Job, err error) {
	job = &Job{
		Interval:     interval,
		StartAt:      startAt,
		WokerSetName: w.set.Name,
		WorkerName:   w.Name,
		Cli:          DefaultJobCli,
	}
	if err = jobDB.Save(&job).Error; err != nil {
		return
	}

	err = w.Queuer.Enqueue(job)

	if job.QueueJobId != "" {
		if err = jobDB.Save(&job).Error; err != nil {
			return
		}
	}

	return
}
<commit_msg>worker: automigrate job table<commit_after>package worker

import (
	"errors"
	"flag"
	"fmt"
	"os"
	"strings"
	"time"

	"github.com/jinzhu/gorm"
	"github.com/qor/qor"
	"github.com/qor/qor/admin"
)

var (
	jobDB            *gorm.DB
	jobId            uint64
	defaultWorkerSet = &WorkerSet{Name: "default"}
	workerSets       = []*WorkerSet{defaultWorkerSet}
	queuers          = map[string]Queuer{}
	DefaultJobCli    = strings.Join(os.Args, " ")
)

func init() {
	flag.Uint64Var(&jobId, "job-id", 0, "qor job id")
}

// SetJobDB will run a auto migration for creating table jobs
func SetJobDB(db *gorm.DB) (err error) {
	jobDB = db
	err = db.AutoMigrate(&Job{}).Error
	return
}

func NewWorker(name string, handle func(job *Job) error, queuer Queuer) (w *Worker) {
	return defaultWorkerSet.NewWorker(name, handle, queuer)
}

func SetAdmin(admin *admin.Admin) {
	ws := defaultWorkerSet.Workers
	defaultWorkerSet = NewWorkerSet(defaultWorkerSet.Name, "/workers", "", admin)
	defaultWorkerSet.Workers = ws
}

func RegisterQueuer(name string, queuer Queuer) {
	queuers[name] = queuer
}

func Listen() {
	flag.Parse()

	if jobId > 0 {
		var job Job
		if err := jobDB.Where("_id = ?", jobId).Find(&job).Error; err != nil {
			fmt.Fprintln(os.Stderr, err)
			os.Exit(1)
		}
		w, err := job.GetWorker()
		if err != nil {
			fmt.Fprintln(os.Stderr, err)
			os.Exit(1)
		}

		w.Run(&job)
	} else {
		for _, queuer := range queuers {
			go func() {
				for {
					jobId, err := queuer.Dequeue()
					if err != nil {
						fmt.Println("qor.worker.dequeue.error:", err)
					} else {
						go RunJob(jobId)
					}
				}
			}()
		}
	}
}

type WorkerSet struct {
	Name    string
	Workers []*Worker
}

func NewWorkerSet(name, router, tmplDir string, a *admin.Admin) (ws *WorkerSet) {
	ws = &WorkerSet{Name: name}
	workerSets = append(workerSets, ws)
	a.GetRouter().Get(router, func(ctx *admin.Context) {})
	// template register
	// menu register
	// Job resource register

	return
}

func (ws *WorkerSet) NewWorker(name string, handle func(job *Job) error, queuer Queuer) (w *Worker) {
	w = &Worker{
		Name:   name,
		Handle: handle,
		Queuer: queuer,
		set:    ws,
	}

	ws.Workers = append(ws.Workers, w)

	return
}

type Worker struct {
	Name   string
	Queuer Queuer
	Config *qor.Config
	set    *WorkerSet

	Handle    func(job *Job) error
	OnStart   func(job *Job) error
	OnKill    func(job *Job) error
	OnSuccess func(job *Job)
	OnFailed  func(job *Job)
}

func (w *Worker) Run(job *Job) (err error) {
	if err = job.SavePID(); err != nil {
		fmt.Fprintln(os.Stderr, err)
		os.Exit(1)
	}
	logger, err := job.GetLogger()
	if err != nil {
		fmt.Println("can't get job logger")
		return
	}

	fmt.Fprintf(logger, "to run job (%d) with pid (%d)", job.Id, job.PID)

	if err = job.UpdateStatus(JobRunning); err != nil {
		fmt.Fprintln(os.Stderr, err)
		os.Exit(1)
	}

	if w.OnStart != nil {
		if err = w.OnStart(job); err != nil {
			logger.Write([]byte("worker.onstart: " + err.Error() + "\n"))

			if err = job.UpdateStatus(JobFailed); err != nil {
				fmt.Fprintln(os.Stderr, err)
				os.Exit(1)
			}

			fmt.Fprintln(os.Stderr, err)
			os.Exit(1)
		}
	}

	if err = w.Handle(job); err != nil {
		if err = job.UpdateStatus(JobFailed); err != nil {
			fmt.Fprintln(os.Stderr, err)
			os.Exit(1)
		}

		if w.OnFailed != nil {
			w.OnFailed(job)
		}

		logger.Write([]byte("worker.hanlde: " + err.Error() + "\n"))
		fmt.Fprintln(os.Stderr, err)
		os.Exit(1)
	} else if w.OnSuccess != nil {
		if err = job.UpdateStatus(JobRun); err != nil {
			fmt.Fprintln(os.Stderr, err)
		}

		w.OnSuccess(job)
	}

	fmt.Fprintf(logger, "finish job (%d) with pid (%d)", job.Id, job.PID)

	return
}

var ErrJobRun = errors.New("job is already run")

func (w *Worker) Kill(job *Job) (err error) {
	if w.OnKill != nil {
		if err = w.OnKill(job); err != nil {
			return
		}
	}

	switch job.Status {
	case JobToRun:
		err = w.Queuer.Purge(job)
	case JobRunning:
		if job.PID == 0 {
			return errors.New("pid is zero")
		}

		var process *os.Process
		process, err = os.FindProcess(job.PID)
		if err != nil {
			return
		}

		err = process.Kill()
	case JobRun:
		return ErrJobRun
	}

	if err == nil {
		err = job.UpdateStatus(JobKilled)
	}

	return
}

func (w *Worker) NewJob(interval uint64, startAt time.Time) (job *Job, err error) {
	return w.NewJobWithCli(interval, startAt, DefaultJobCli)
}

func (w *Worker) NewJobWithCli(interval uint64, startAt time.Time, cli string) (job *Job, err error) {
	job = &Job{
		Interval:     interval,
		StartAt:      startAt,
		WokerSetName: w.set.Name,
		WorkerName:   w.Name,
		Cli:          DefaultJobCli,
	}
	if err = jobDB.Save(&job).Error; err != nil {
		return
	}

	err = w.Queuer.Enqueue(job)

	if job.QueueJobId != "" {
		if err = jobDB.Save(&job).Error; err != nil {
			return
		}
	}

	return
}
<|endoftext|>


<commit_before>package main

import (
	"database/sql"
	"encoding/json"
	"flag"
	"fmt"
	"github.com/bmizerany/pq"
	"io/ioutil"
	"log"
	"net/http"
	"time"
)

const OPEN311_API_URI = "http://311api.cityofchicago.org/open311/v2/requests.json?extensions=true&page_size=500"

type Open311Request struct {
	Lat, Long                                                                                float64
	Ward, Police_district                                                                                int
	Service_request_id, Status, Service_name, Service_code, Agency_responsible, Address, Channel, Media_url string
	Requested_datetime, Updated_datetime                                                                    string // FIXME: should these be proper time objects?
	Extended_attributes                                                                                map[string]interface{}
	Notes                                                                                []map[string]interface{}
}

type Worker struct {
	Db        *sql.DB
	LastRunAt time.Time
}

var worker Worker
var sr_number string
var backfill bool
var backfill_date string

func init() {
	// open database
	db, err := sql.Open("postgres", "dbname=cwfy sslmode=disable")
	if err != nil {
		log.Fatal("Cannot open database connection", err)
	}
	worker.Db = db

	// fetch SR num from command line, if present
	flag.StringVar(&sr_number, "sr-number", "", "SR number to fetch")
	flag.BoolVar(&backfill, "backfill", false, "run in reverse and backfill data")
	flag.StringVar(&backfill_date, "backfill-from", time.Now().Format(time.RFC3339), "date to start backfilling data from. Use RFC3339 format")
}

func main() {
	defer worker.Db.Close()
	flag.Parse()

	if sr_number != "" {
		sr := fetchSingleRequest(sr_number)
		sr.Save()
		return
	}

	start_backfill_from := backfill_date
	for {
		switch {
		case backfill:
			requests := backFillRequests(start_backfill_from)
			for _, request := range requests {
				request.Save()
			}

			start_backfill_from = requests[len(requests)-1].Updated_datetime // FIXME: is it safe to assume the items are sorted?

		case time.Since(worker.LastRunAt) > (30 * time.Second):
			// load requests from open311
			for _, request := range fetchRequests() {
				request.Save()
			}
			worker.LastRunAt = time.Now()
		default:
			log.Print("sleeping for 10 seconds")
			time.Sleep(10 * time.Second)
		}
	}
}

func (req Open311Request) String() string {
	// pretty print SR information
	return fmt.Sprintf("%s: %s at %s %f,%f, last update %s", req.Service_request_id, req.Service_name, req.Address, req.Lat, req.Long, req.Updated_datetime)
}

func (req Open311Request) Save() (persisted bool) {
	// create or update a SR

	// open311 says we should always ignore a SR that does not have a SR# assigned
	if req.Service_request_id == "" {
		log.Printf("cowardly refusing to create a new SR record because of empty SR#. Request type is %s", req.Service_name)
		return false
	}

	persisted = false

	// find existing record if exists
	var existing_id int
	err := worker.Db.QueryRow("SELECT id FROM service_requests WHERE service_request_id = $1", req.Service_request_id).Scan(&existing_id)
	switch {
	case err == sql.ErrNoRows:
		// log.Printf("did not find existing record %s", req.Service_request_id)
	case err != nil:
		// log.Print("error searching for existing SR", err)
	default:
		persisted = true
		// log.Printf("found existing sr %s", req.Service_request_id)
	}

	var stmt *sql.Stmt

	if !persisted {
		// create new record
		stmt, err = worker.Db.Prepare("INSERT INTO service_requests(service_request_id," +
			"status, service_name, service_code, agency_responsible, " +
			"address, requested_datetime, updated_datetime, lat, long," +
			"ward, police_district, media_url, channel, duplicate, parent_service_request_id, closed_datetime) " +
			"VALUES ($1::varchar, $2, $3, $4, $5, $6, $7, $8, $9, $10, $11, $12, $13, $14, $15, $16, $17); ")

		// "WHERE NOT EXISTS (SELECT 1 FROM service_requests WHERE service_request_id = $1);")

		if err != nil {
			log.Fatal("error preparing database insert statement", err)
		}

	} else {
		// update existing record
		stmt, err = worker.Db.Prepare("UPDATE service_requests SET " +
			"status = $2, service_name = $3, service_code = $4, agency_responsible = $5, " +
			"address = $6, requested_datetime = $7, updated_datetime = $8, lat = $9, long = $10," +
			"ward = $11, police_district = $12, media_url = $13, channel = $14, duplicate = $15, " +
			"parent_service_request_id = $16, updated_at = NOW(), closed_datetime = $17 WHERE service_request_id = $1;")

		if err != nil {
			log.Fatal("error preparing database update statement", err)
		}
	}

	tx, err := worker.Db.Begin()

	if err != nil {
		log.Fatal("error beginning transaction", err)
	}

	t := req.ExtractClosedDatetime()
	closed_time := pq.NullTime{Time: t, Valid: !t.IsZero()}

	_, err = tx.Stmt(stmt).Exec(req.Service_request_id,
		req.Status,
		req.Service_name,
		req.Service_code,
		req.Agency_responsible,
		req.Address,
		req.Requested_datetime,
		req.Updated_datetime,
		req.Lat,
		req.Long,
		req.Extended_attributes["ward"],
		req.Extended_attributes["police_district"],
		req.Media_url,
		req.Extended_attributes["channel"],
		req.Extended_attributes["duplicate"],
		req.Extended_attributes["parent_service_request_id"],
		closed_time,
	)

	if err != nil {
		log.Fatalf("could not update %s because %s", req.Service_request_id, err)
	} else {
		var verb string
		switch {
		case !persisted && closed_time.Time.IsZero():
			verb = "CREATED"
		case !persisted && !closed_time.Time.IsZero():
			verb = "CREATED/CLOSED"
		case persisted && closed_time.Time.IsZero():
			verb = "UPDATED"
		case persisted && !closed_time.Time.IsZero():
			verb = "UPDATED/CLOSED"
		}

		log.Printf("[%s] %s", verb, req)
		persisted = true
	}

	err = tx.Commit()
	if err != nil {
		log.Fatal("error closing transaction", err)
	}

	return persisted

	// calculate closed time if necessary

}

func (req Open311Request) ExtractClosedDatetime() time.Time {
	// given an extended_attributes JSON blob, pluck out the closed time, if present
	// req.PrintNotes()

	var closed_at time.Time
	for _, note := range req.Notes {
		if note["type"] == "closed" {
			parsed_date, err := time.Parse("2006-01-02T15:04:05-07:00", note["datetime"].(string))
			if err != nil {
				log.Print("error parsing date", err)
			}
			log.Printf("SR %s closed at: %s", req, parsed_date)
			closed_at = parsed_date
		}
	}

	return closed_at
}

func (req Open311Request) PrintNotes() {
	fmt.Printf("Notes for SR %s:\n", req.Service_request_id)

	for _, note := range req.Notes {
		fmt.Printf("%+v\n", note)
	}
}

func fetchSingleRequest(sr_number string) (request Open311Request) {
	// given an SR, fetch the record
	log.Printf("fetching single SR %s", sr_number)
	open311_api_endpoint := fmt.Sprintf("http://311api.cityofchicago.org/open311/v2/requests/%s.json?extensions=true", sr_number)

	log.Printf("fetching from %s", open311_api_endpoint)
	resp, err := http.Get(open311_api_endpoint)
	defer resp.Body.Close()

	if err != nil {
		log.Fatalln("error fetching from Open311 endpoint", err)
	}

	// load response body
	body, err := ioutil.ReadAll(resp.Body)
	if err != nil {
		log.Fatal("error loading response body", err)
	}

	// parse JSON and load into an array of Open311Request objects
	var requests []Open311Request

	err = json.Unmarshal(body, &requests)
	if err != nil {
		log.Fatal("error parsing JSON:", err)
	}

	log.Printf("received %d requests from Open311", len(requests))

	return requests[0]
}

func fetchRequests() (requests []Open311Request) {
	// find the most recent SR that we know about in the database
	rows, err := worker.Db.Query("SELECT MAX(updated_datetime) FROM service_requests;")
	if err != nil {
		log.Fatal("error finding most recent service request", err)
	}

	last_updated_at := time.Now()
	for rows.Next() {
		if err := rows.Scan(&last_updated_at); err != nil {
			log.Print("error finding most recent SR", err)
		}

		log.Printf("most recent SR timestamp %s", last_updated_at)
	}

	// janky hack to transform the last updated timestamp into
	// a format that plays nicely with the Open311 API
	// FIXME: there HAS to be a better way to handle this.
	formatted_date_string := last_updated_at.Format(time.RFC3339)
	formatted_date_string_with_tz := formatted_date_string[0:len(formatted_date_string)-1] + "-0500" // trunc the trailing 'Z' and tack on timezone

	// construct the request URI using base params and the proper time
	open311_api_endpoint := OPEN311_API_URI + "&updated_after=" + formatted_date_string_with_tz

	log.Printf("fetching from %s", open311_api_endpoint)
	resp, err := http.Get(open311_api_endpoint)
	defer resp.Body.Close()

	if err != nil {
		log.Fatalln("error fetching from Open311 endpoint", err)
	}

	// load response body
	body, err := ioutil.ReadAll(resp.Body)
	if err != nil {
		log.Fatal("error loading response body", err)
	}

	// parse JSON and load into an array of Open311Request objects
	err = json.Unmarshal(body, &requests)
	if err != nil {
		log.Fatal("error parsing JSON:", err)
	}

	log.Printf("received %d requests from Open311", len(requests))

	return requests
}

func backFillRequests(start_from string) (requests []Open311Request) {
	formatted_date_string, err := time.Parse(time.RFC3339, start_from)
	if err != nil {
		log.Fatal("[backfill] error parsing date to start from", err)
	}
	formatted_date_string_with_tz := formatted_date_string.Format(time.RFC3339)

	// construct the request URI using base params and the proper time
	open311_api_endpoint := OPEN311_API_URI + "&updated_before=" + formatted_date_string_with_tz

	log.Printf("[backfill] fetching from %s", open311_api_endpoint)
	resp, err := http.Get(open311_api_endpoint)
	defer resp.Body.Close()

	if err != nil {
		log.Fatalln("[backfill] error fetching from Open311 endpoint", err)
	}

	// load response body
	body, err := ioutil.ReadAll(resp.Body)
	if err != nil {
		log.Fatal("[backfill] error loading response body", err)
	}

	// parse JSON and load into an array of Open311Request objects
	err = json.Unmarshal(body, &requests)
	if err != nil {
		log.Fatal("[backfill] error parsing JSON:", err)
	}

	log.Printf("[backfill] received %d requests from Open311", len(requests))

	return requests
}
<commit_msg>switch database driver to lib/pq<commit_after>package main

import (
	"database/sql"
	"encoding/json"
	"flag"
	"fmt"
	"github.com/lib/pq"
	"io/ioutil"
	"log"
	"net/http"
	"time"
)

const OPEN311_API_URI = "http://311api.cityofchicago.org/open311/v2/requests.json?extensions=true&page_size=500"

type Open311Request struct {
	Lat, Long                                                                                float64
	Ward, Police_district                                                                                int
	Service_request_id, Status, Service_name, Service_code, Agency_responsible, Address, Channel, Media_url string
	Requested_datetime, Updated_datetime                                                                    string // FIXME: should these be proper time objects?
	Extended_attributes                                                                                map[string]interface{}
	Notes                                                                                []map[string]interface{}
}

type Worker struct {
	Db        *sql.DB
	LastRunAt time.Time
}

var worker Worker
var sr_number string
var backfill bool
var backfill_date string

func init() {
	// open database
	db, err := sql.Open("postgres", "dbname=cwfy sslmode=disable")
	if err != nil {
		log.Fatal("Cannot open database connection", err)
	}
	worker.Db = db

	// fetch SR num from command line, if present
	flag.StringVar(&sr_number, "sr-number", "", "SR number to fetch")
	flag.BoolVar(&backfill, "backfill", false, "run in reverse and backfill data")
	flag.StringVar(&backfill_date, "backfill-from", time.Now().Format(time.RFC3339), "date to start backfilling data from. Use RFC3339 format")
}

func main() {
	defer worker.Db.Close()
	flag.Parse()

	if sr_number != "" {
		sr := fetchSingleRequest(sr_number)
		sr.Save()
		return
	}

	start_backfill_from := backfill_date
	for {
		switch {
		case backfill:
			requests := backFillRequests(start_backfill_from)
			for _, request := range requests {
				request.Save()
			}

			start_backfill_from = requests[len(requests)-1].Updated_datetime // FIXME: is it safe to assume the items are sorted?

		case time.Since(worker.LastRunAt) > (30 * time.Second):
			// load requests from open311
			for _, request := range fetchRequests() {
				request.Save()
			}
			worker.LastRunAt = time.Now()
		default:
			log.Print("sleeping for 10 seconds")
			time.Sleep(10 * time.Second)
		}
	}
}

func (req Open311Request) String() string {
	// pretty print SR information
	return fmt.Sprintf("%s: %s at %s %f,%f, last update %s", req.Service_request_id, req.Service_name, req.Address, req.Lat, req.Long, req.Updated_datetime)
}

func (req Open311Request) Save() (persisted bool) {
	// create or update a SR

	// open311 says we should always ignore a SR that does not have a SR# assigned
	if req.Service_request_id == "" {
		log.Printf("cowardly refusing to create a new SR record because of empty SR#. Request type is %s", req.Service_name)
		return false
	}

	persisted = false

	// find existing record if exists
	var existing_id int
	err := worker.Db.QueryRow("SELECT id FROM service_requests WHERE service_request_id = $1", req.Service_request_id).Scan(&existing_id)
	switch {
	case err == sql.ErrNoRows:
		// log.Printf("did not find existing record %s", req.Service_request_id)
	case err != nil:
		// log.Print("error searching for existing SR", err)
	default:
		persisted = true
		// log.Printf("found existing sr %s", req.Service_request_id)
	}

	var stmt *sql.Stmt

	if !persisted {
		// create new record
		stmt, err = worker.Db.Prepare("INSERT INTO service_requests(service_request_id," +
			"status, service_name, service_code, agency_responsible, " +
			"address, requested_datetime, updated_datetime, lat, long," +
			"ward, police_district, media_url, channel, duplicate, parent_service_request_id, closed_datetime) " +
			"VALUES ($1::varchar, $2, $3, $4, $5, $6, $7, $8, $9, $10, $11, $12, $13, $14, $15, $16, $17); ")

		// "WHERE NOT EXISTS (SELECT 1 FROM service_requests WHERE service_request_id = $1);")

		if err != nil {
			log.Fatal("error preparing database insert statement", err)
		}

	} else {
		// update existing record
		stmt, err = worker.Db.Prepare("UPDATE service_requests SET " +
			"status = $2, service_name = $3, service_code = $4, agency_responsible = $5, " +
			"address = $6, requested_datetime = $7, updated_datetime = $8, lat = $9, long = $10," +
			"ward = $11, police_district = $12, media_url = $13, channel = $14, duplicate = $15, " +
			"parent_service_request_id = $16, updated_at = NOW(), closed_datetime = $17 WHERE service_request_id = $1;")

		if err != nil {
			log.Fatal("error preparing database update statement", err)
		}
	}

	tx, err := worker.Db.Begin()

	if err != nil {
		log.Fatal("error beginning transaction", err)
	}

	t := req.ExtractClosedDatetime()
	closed_time := pq.NullTime{Time: t, Valid: !t.IsZero()}

	_, err = tx.Stmt(stmt).Exec(req.Service_request_id,
		req.Status,
		req.Service_name,
		req.Service_code,
		req.Agency_responsible,
		req.Address,
		req.Requested_datetime,
		req.Updated_datetime,
		req.Lat,
		req.Long,
		req.Extended_attributes["ward"],
		req.Extended_attributes["police_district"],
		req.Media_url,
		req.Extended_attributes["channel"],
		req.Extended_attributes["duplicate"],
		req.Extended_attributes["parent_service_request_id"],
		closed_time,
	)

	if err != nil {
		log.Fatalf("could not update %s because %s", req.Service_request_id, err)
	} else {
		var verb string
		switch {
		case !persisted && closed_time.Time.IsZero():
			verb = "CREATED"
		case !persisted && !closed_time.Time.IsZero():
			verb = "CREATED/CLOSED"
		case persisted && closed_time.Time.IsZero():
			verb = "UPDATED"
		case persisted && !closed_time.Time.IsZero():
			verb = "UPDATED/CLOSED"
		}

		log.Printf("[%s] %s", verb, req)
		persisted = true
	}

	err = tx.Commit()
	if err != nil {
		log.Fatal("error closing transaction", err)
	}

	return persisted

	// calculate closed time if necessary

}

func (req Open311Request) ExtractClosedDatetime() time.Time {
	// given an extended_attributes JSON blob, pluck out the closed time, if present
	// req.PrintNotes()

	var closed_at time.Time
	for _, note := range req.Notes {
		if note["type"] == "closed" {
			parsed_date, err := time.Parse("2006-01-02T15:04:05-07:00", note["datetime"].(string))
			if err != nil {
				log.Print("error parsing date", err)
			}
			log.Printf("SR %s closed at: %s", req, parsed_date)
			closed_at = parsed_date
		}
	}

	return closed_at
}

func (req Open311Request) PrintNotes() {
	fmt.Printf("Notes for SR %s:\n", req.Service_request_id)

	for _, note := range req.Notes {
		fmt.Printf("%+v\n", note)
	}
}

func fetchSingleRequest(sr_number string) (request Open311Request) {
	// given an SR, fetch the record
	log.Printf("fetching single SR %s", sr_number)
	open311_api_endpoint := fmt.Sprintf("http://311api.cityofchicago.org/open311/v2/requests/%s.json?extensions=true", sr_number)

	log.Printf("fetching from %s", open311_api_endpoint)
	resp, err := http.Get(open311_api_endpoint)
	defer resp.Body.Close()

	if err != nil {
		log.Fatalln("error fetching from Open311 endpoint", err)
	}

	// load response body
	body, err := ioutil.ReadAll(resp.Body)
	if err != nil {
		log.Fatal("error loading response body", err)
	}

	// parse JSON and load into an array of Open311Request objects
	var requests []Open311Request

	err = json.Unmarshal(body, &requests)
	if err != nil {
		log.Fatal("error parsing JSON:", err)
	}

	log.Printf("received %d requests from Open311", len(requests))

	return requests[0]
}

func fetchRequests() (requests []Open311Request) {
	// find the most recent SR that we know about in the database
	rows, err := worker.Db.Query("SELECT MAX(updated_datetime) FROM service_requests;")
	if err != nil {
		log.Fatal("error finding most recent service request", err)
	}

	last_updated_at := time.Now()
	for rows.Next() {
		if err := rows.Scan(&last_updated_at); err != nil {
			log.Print("error finding most recent SR", err)
		}

		log.Printf("most recent SR timestamp %s", last_updated_at)
	}

	// janky hack to transform the last updated timestamp into
	// a format that plays nicely with the Open311 API
	// FIXME: there HAS to be a better way to handle this.
	formatted_date_string := last_updated_at.Format(time.RFC3339)
	formatted_date_string_with_tz := formatted_date_string[0:len(formatted_date_string)-1] + "-0500" // trunc the trailing 'Z' and tack on timezone

	// construct the request URI using base params and the proper time
	open311_api_endpoint := OPEN311_API_URI + "&updated_after=" + formatted_date_string_with_tz

	log.Printf("fetching from %s", open311_api_endpoint)
	resp, err := http.Get(open311_api_endpoint)
	defer resp.Body.Close()

	if err != nil {
		log.Fatalln("error fetching from Open311 endpoint", err)
	}

	// load response body
	body, err := ioutil.ReadAll(resp.Body)
	if err != nil {
		log.Fatal("error loading response body", err)
	}

	// parse JSON and load into an array of Open311Request objects
	err = json.Unmarshal(body, &requests)
	if err != nil {
		log.Fatal("error parsing JSON:", err)
	}

	log.Printf("received %d requests from Open311", len(requests))

	return requests
}

func backFillRequests(start_from string) (requests []Open311Request) {
	formatted_date_string, err := time.Parse(time.RFC3339, start_from)
	if err != nil {
		log.Fatal("[backfill] error parsing date to start from", err)
	}
	formatted_date_string_with_tz := formatted_date_string.Format(time.RFC3339)

	// construct the request URI using base params and the proper time
	open311_api_endpoint := OPEN311_API_URI + "&updated_before=" + formatted_date_string_with_tz

	log.Printf("[backfill] fetching from %s", open311_api_endpoint)
	resp, err := http.Get(open311_api_endpoint)
	defer resp.Body.Close()

	if err != nil {
		log.Fatalln("[backfill] error fetching from Open311 endpoint", err)
	}

	// load response body
	body, err := ioutil.ReadAll(resp.Body)
	if err != nil {
		log.Fatal("[backfill] error loading response body", err)
	}

	// parse JSON and load into an array of Open311Request objects
	err = json.Unmarshal(body, &requests)
	if err != nil {
		log.Fatal("[backfill] error parsing JSON:", err)
	}

	log.Printf("[backfill] received %d requests from Open311", len(requests))

	return requests
}
<|endoftext|>


<commit_before>/* Copyright (C) 2015-2022 김운하 (unha.kim@ghts.org)

이 파일은 GHTS의 일부입니다.

이 프로그램은 자유 소프트웨어입니다.
소프트웨어의 피양도자는 자유 소프트웨어 재단이 공표한 GNU LGPL 2.1판
규정에 따라 프로그램을 개작하거나 재배포할 수 있습니다.

이 프로그램은 유용하게 사용될 수 있으리라는 희망에서 배포되고 있지만,
특정한 목적에 적합하다거나, 이익을 안겨줄 수 있다는 묵시적인 보증을 포함한
어떠한 형태의 보증도 제공하지 않습니다.
보다 자세한 사항에 대해서는 GNU LGPL 2.1판을 참고하시기 바랍니다.
GNU LGPL 2.1판은 이 프로그램과 함께 제공됩니다.
만약, 이 문서가 누락되어 있다면 자유 소프트웨어 재단으로 문의하시기 바랍니다.
(자유 소프트웨어 재단 : Free Software Foundation, Inc.,
59 Temple Place - Suite 330, Boston, MA 02111-1307, USA)

Copyright (C) 2015-2022년 UnHa Kim (unha.kim@ghts.org)

This file is part of GHTS.

GHTS is free software: you can redistribute it and/or modify
it under the terms of the GNU Lesser General Public License as published by
the Free Software Foundation, version 2.1 of the License.

GHTS is distributed in the hope that it will be useful,
but WITHOUT ANY WARRANTY; without even the implied warranty of
MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
GNU Lesser General Public License for more details.

You should have received a copy of the GNU Lesser General Public License
along with GHTS.  If not, see <http://www.gnu.org/licenses/>. */

package dll32

const (
	xing_dll = "xingAPI.dll"
	WM_USER  = 0x0400
)

const (
	XM_INIT                    = 1024
	XM_DISCONNECT              = XM_INIT + 1
	XM_RECEIVE_DATA            = XM_INIT + 3
	XM_RECEIVE_REAL_DATA       = XM_INIT + 4
	XM_LOGIN                   = XM_INIT + 5
	XM_LOGOUT                  = XM_INIT + 6
	XM_TIMEOUT                 = XM_INIT + 7
	XM_RECEIVE_LINK_DATA       = XM_INIT + 8
	XM_RECEIVE_REAL_DATA_CHART = XM_INIT + 10
)

const (
	RCV_TR_DATA      = 1
	RCV_MSG_DATA     = 2
	RCV_SYSTEM_ERROR = 3
	RCV_RELEASE      = 4
)

const (
	FALSE = 0
	TRUE  = 1
)
<commit_msg>Xing dll32 모듈 상수 정의 코드 정리<commit_after>/* Copyright (C) 2015-2022 김운하 (unha.kim@ghts.org)

이 파일은 GHTS의 일부입니다.

이 프로그램은 자유 소프트웨어입니다.
소프트웨어의 피양도자는 자유 소프트웨어 재단이 공표한 GNU LGPL 2.1판
규정에 따라 프로그램을 개작하거나 재배포할 수 있습니다.

이 프로그램은 유용하게 사용될 수 있으리라는 희망에서 배포되고 있지만,
특정한 목적에 적합하다거나, 이익을 안겨줄 수 있다는 묵시적인 보증을 포함한
어떠한 형태의 보증도 제공하지 않습니다.
보다 자세한 사항에 대해서는 GNU LGPL 2.1판을 참고하시기 바랍니다.
GNU LGPL 2.1판은 이 프로그램과 함께 제공됩니다.
만약, 이 문서가 누락되어 있다면 자유 소프트웨어 재단으로 문의하시기 바랍니다.
(자유 소프트웨어 재단 : Free Software Foundation, Inc.,
59 Temple Place - Suite 330, Boston, MA 02111-1307, USA)

Copyright (C) 2015-2022년 UnHa Kim (unha.kim@ghts.org)

This file is part of GHTS.

GHTS is free software: you can redistribute it and/or modify
it under the terms of the GNU Lesser General Public License as published by
the Free Software Foundation, version 2.1 of the License.

GHTS is distributed in the hope that it will be useful,
but WITHOUT ANY WARRANTY; without even the implied warranty of
MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
GNU Lesser General Public License for more details.

You should have received a copy of the GNU Lesser General Public License
along with GHTS.  If not, see <http://www.gnu.org/licenses/>. */

package dll32

const (
	xing_dll = "xingAPI.dll"

	FALSE = 0
	TRUE  = 1

	WM_USER                    = 0x0400
	XM_INIT                    = WM_USER
	XM_DISCONNECT              = XM_INIT + 1
	XM_RECEIVE_DATA            = XM_INIT + 3
	XM_RECEIVE_REAL_DATA       = XM_INIT + 4
	XM_LOGIN                   = XM_INIT + 5
	XM_LOGOUT                  = XM_INIT + 6
	XM_TIMEOUT                 = XM_INIT + 7
	XM_RECEIVE_LINK_DATA       = XM_INIT + 8
	XM_RECEIVE_REAL_DATA_CHART = XM_INIT + 10

	RCV_TR_DATA      = 1
	RCV_MSG_DATA     = 2
	RCV_SYSTEM_ERROR = 3
	RCV_RELEASE      = 4
)
<|endoftext|>


<commit_before>package main

import (
	"fmt"
)

type CloverInt struct {
	value int
}

type CloverFloat struct {
	value float32
}

type CloverBool struct {
	value bool
}

type CloverString struct {
	value string
}

type CloverObj interface {
	ShowValue() string
}

func (s CloverInt) ShowValue() string {
	return fmt.Sprint(s.value)
}

func (s CloverFloat) ShowValue() string {
	return fmt.Sprint(s.value)
}

func (s CloverBool) ShowValue() string {
	return fmt.Sprint(s.value)
}

func (s CloverString) ShowValue() string {
	return fmt.Sprint(s.value)
}
/*
func main() {
	ci := CloverInt{9}
	c := CloverInt{99}
	op := CloverInt{3}
	bl := CloverBool{true}
	fl := CloverFloat{5.55}
	fmt.Println(ci.ShowValue())
  fmt.Println(Plus(ci, op).ShowValue())
	fmt.Println(Or(bl, CloverBool{false}).ShowValue())
	fmt.Println(Mulf(fl, fl).ShowValue())
	fmt.Println(tes(ci, op).ShowValue())
	println(generated(c, op))
}
*/
func Tes(x, y int) int {
	return x + y
}

// built-in functions
// int
func Plus(o1, o2 CloverObj) CloverObj {
	v1 := o1.(CloverInt).value
	v2 := o2.(CloverInt).value
	o := CloverInt{v1 + v2}
	return o
}

func Minus(o1, o2 CloverObj) CloverObj {
	v1 := o1.(CloverInt).value
	v2 := o2.(CloverInt).value
	o := CloverInt{v1 - v2}
	return o
}

func Mul(o1, o2 CloverObj) CloverObj {
	v1 := o1.(CloverInt).value
	v2 := o2.(CloverInt).value
	o := CloverInt{v1 * v2}
	return o
}

// float
func Plusf(o1, o2 CloverObj) CloverObj {
	v1 := o1.(CloverFloat).value
	v2 := o2.(CloverFloat).value
	o := CloverFloat{v1 + v2}
	return o
}

func Minusf(o1, o2 CloverObj) CloverObj {
	v1 := o1.(CloverFloat).value
	v2 := o2.(CloverFloat).value
	o := CloverFloat{v1 - v2}
	return o
}

func Mulf(o1, o2 CloverObj) CloverObj {
	v1 := o1.(CloverFloat).value
	v2 := o2.(CloverFloat).value
	o := CloverFloat{v1 * v2}
	return o
}

// bool
func And(o1, o2 CloverObj) CloverObj {
	v1 := o1.(CloverBool).value
	v2 := o2.(CloverBool).value
	o := CloverBool{v1 && v2}
	return o
}

func Or(o1, o2 CloverObj) CloverObj {
	v1 := o1.(CloverBool).value
	v2 := o2.(CloverBool).value
	o := CloverBool{v1 || v2}
	return o
}

// util
func println(o CloverObj) {
	fmt.Println(o.ShowValue())
}

// user-define

func pluser(x, y CloverObj) CloverObj {
	return Plus(x, y)
}

func tes(x, y interface{}) CloverObj {
	v1 := x.(CloverInt)
	v2 := y.(CloverInt)
	return Plus(v1, v2)
}

// (defn f-name [x y] (+ x y))
// to
// List[Symbol "defn", Symbol "fname", Vector[Symbol "x", Symbol "y"],
//      List[Symbol "+", Symbol "x", Symbol "y"]
// ]
//
<commit_msg>refact base.go<commit_after>package main

import (
	"fmt"
)

type CloverInt struct {
	value int
}

type CloverFloat struct {
	value float32
}

type CloverBool struct {
	value bool
}

type CloverString struct {
	value string
}

type CloverObj interface {
	ShowValue() string
}

func (s CloverInt) ShowValue() string {
	return fmt.Sprint(s.value)
}

func (s CloverFloat) ShowValue() string {
	return fmt.Sprint(s.value)
}

func (s CloverBool) ShowValue() string {
	return fmt.Sprint(s.value)
}

func (s CloverString) ShowValue() string {
	return fmt.Sprint(s.value)
}

// built-in functions
// int
func Plus(o1, o2 CloverObj) CloverObj {
	v1 := o1.(CloverInt).value
	v2 := o2.(CloverInt).value
	o := CloverInt{v1 + v2}
	return o
}

func Minus(o1, o2 CloverObj) CloverObj {
	v1 := o1.(CloverInt).value
	v2 := o2.(CloverInt).value
	o := CloverInt{v1 - v2}
	return o
}

func Mul(o1, o2 CloverObj) CloverObj {
	v1 := o1.(CloverInt).value
	v2 := o2.(CloverInt).value
	o := CloverInt{v1 * v2}
	return o
}

// float
func Plusf(o1, o2 CloverObj) CloverObj {
	v1 := o1.(CloverFloat).value
	v2 := o2.(CloverFloat).value
	o := CloverFloat{v1 + v2}
	return o
}

func Minusf(o1, o2 CloverObj) CloverObj {
	v1 := o1.(CloverFloat).value
	v2 := o2.(CloverFloat).value
	o := CloverFloat{v1 - v2}
	return o
}

func Mulf(o1, o2 CloverObj) CloverObj {
	v1 := o1.(CloverFloat).value
	v2 := o2.(CloverFloat).value
	o := CloverFloat{v1 * v2}
	return o
}

// bool
func And(o1, o2 CloverObj) CloverObj {
	v1 := o1.(CloverBool).value
	v2 := o2.(CloverBool).value
	o := CloverBool{v1 && v2}
	return o
}

func Or(o1, o2 CloverObj) CloverObj {
	v1 := o1.(CloverBool).value
	v2 := o2.(CloverBool).value
	o := CloverBool{v1 || v2}
	return o
}

// util
func println(o CloverObj) {
	fmt.Println(o.ShowValue())
}

func tes(x, y interface{}) CloverObj {
	v1 := x.(CloverInt)
	v2 := y.(CloverInt)
	return Plus(v1, v2)
}

// (defn f-name [x y] (+ x y))
// to
// List[Symbol "defn", Symbol "fname", Vector[Symbol "x", Symbol "y"],
//      List[Symbol "+", Symbol "x", Symbol "y"]
// ]
//
<|endoftext|>


<commit_before>// Tideland Go REST Server Library - REST
//
// Copyright (C) 2009-2016 Frank Mueller / Tideland / Oldenburg / Germany
//
// All rights reserved. Use of this source code is governed
// by the new BSD license.

// The Tideland Go REST Server Library provides the package rest for the
// implementation of servers with a RESTful API. The business has to
// be implemented in types fullfilling the ResourceHandler interface.
// This basic interface only allows the initialization of the handler.
// More interesting are the other interfaces like GetResourceHandler
// which defines the Get() method for the HTTP request method GET.
// Others are for PUT, POST, HEAD, PATCH, DELETE, and OPTIONS. Their
// according methods get a Job as argument. It provides convenient
// helpers for the processing of the job.
//
// The processing methods return two values: a boolean and an error.
// The latter is pretty clear, it signals a job processing error. The
// boolean is more interesting. Registering a handler is based on a
// domain and a resource. The URL
//
// /<DOMAIN>/<RESOURCE>
//
// leads to a handler, or even better, to a list of handlers. All
// are used as long as the returned boolean value is true. E.g. the
// first handler can check the authentication, the second one the
// authorization, and the third one does the business. Additionally
// the URL
//
// /<DOMAIN>/<RESOURCE>/<ID>
//
// provides the resource identifier via Job.ResourceID().
package rest

//--------------------
// IMPORTS
//--------------------

import (
	"github.com/tideland/golib/version"
)

//--------------------
// VERSION
//--------------------

// PackageVersion returns the version of the version package.
func PackageVersion() version.Version {
	return version.New(2, 0, 0, "beta", "2016-09-02")
}

// EOF
<commit_msg>Added some documentation about handler deployment.<commit_after>// Tideland Go REST Server Library - REST
//
// Copyright (C) 2009-2016 Frank Mueller / Tideland / Oldenburg / Germany
//
// All rights reserved. Use of this source code is governed
// by the new BSD license.

// The Tideland Go REST Server Library provides the package rest for the
// implementation of servers with a RESTful API. The business has to
// be implemented in types fullfilling the ResourceHandler interface.
// This basic interface only allows the initialization of the handler.
// More interesting are the other interfaces like GetResourceHandler
// which defines the Get() method for the HTTP request method GET.
// Others are for PUT, POST, HEAD, PATCH, DELETE, and OPTIONS. Their
// according methods get a Job as argument. It provides convenient
// helpers for the processing of the job.
//
// The processing methods return two values: a boolean and an error.
// The latter is pretty clear, it signals a job processing error. The
// boolean is more interesting. Registering a handler is based on a
// domain and a resource. The URL
//
//     /<DOMAIN>/<RESOURCE>
//
// leads to a handler, or even better, to a list of handlers. All
// are used as long as the returned boolean value is true. E.g. the
// first handler can check the authentication, the second one the
// authorization, and the third one does the business. Additionally
// the URL
//
//     /<DOMAIN>/<RESOURCE>/<ID>
//
// provides the resource identifier via Job.ResourceID().
//
// The handlers then are deployed to the Multiplexer which implements
// the Handler interface of the net/http package. So the typical order
// is
//
//     mux := rest.NewMultiplexer()
//     mux.Register("domain", "resource-type-a", NewTypeAHandler("foo"))
//     mux.Register("domain", "resource-type-b", NewTypeBHandler("bar"))
//     mux.Register("admin", "user", NewUserManagementHandler())
//     http.ListenAndServe(":8000", mux)
//
// Additionally further handlers can be registered or running once
// removed during runtime.
package rest

//--------------------
// IMPORTS
//--------------------

import (
	"github.com/tideland/golib/version"
)

//--------------------
// VERSION
//--------------------

// PackageVersion returns the version of the version package.
func PackageVersion() version.Version {
	return version.New(2, 0, 0, "beta", "2016-09-12")
}

// EOF
<|endoftext|>


<commit_before>package restgate

/*
|--------------------------------------------------------------------------
| WARNING
|--------------------------------------------------------------------------
| Never Set HTTPSProtectionOff=true In Production.
| The Key and Password will be exposed and highly unsecure otherwise!
| The database server should also use HTTPS Connection and be hidden away
|
*/

/*
Thanks to Ido Ben-Natan ("IdoBn") for postgres fix.
Thanks to Jeremy Saenz & Brendon Murphy for timing-attack protection
*/

import (
	// "errors"
	"crypto/subtle"
	"database/sql"
	"fmt"
	"log"
	"net/http"
	"strings"

	e "github.com/pjebs/jsonerror"
	"gopkg.in/unrolled/render.v1"
)

type AuthenticationSource int

const (
	Static   AuthenticationSource = 0
	Database                      = 1
)

//When AuthenticationSource=Static, Key(s)=Actual Key and Secret(s)=Actual Secret.
//When AuthenticationSource=Database, Key[0]=Key_Column and Secret[0]=Secret_Column.
type Config struct {
	*sql.DB
	Key                []string
	Secret             []string //Can be "" but not recommended
	TableName          string
	ErrorMessages      map[int]map[string]string
	Context            func(r *http.Request, authenticatedKey string)
	Debug              bool
	Postgres           bool
	HTTPSProtectionOff bool //Default is HTTPS Protection On
}

type RESTGate struct {
	headerKeyLabel    string
	headerSecretLabel string
	source            AuthenticationSource
	config            Config
}

func New(headerKeyLabel string, headerSecretLabel string, as AuthenticationSource, config Config) *RESTGate {
	t := &RESTGate{headerKeyLabel: headerKeyLabel, headerSecretLabel: headerSecretLabel, source: as, config: config}
	log.Printf("RestGate initializing")

	numberKeys := len(t.config.Key)
	numberSecrets := len(t.config.Secret)

	if numberKeys == 0 { //Key is not set
		if t.config.Debug == true {
			log.Printf("RestGate: Key is not set")
		}
		return nil
	}

	if numberSecrets > numberKeys { //Too many Secret's defined
		if t.config.Debug == true {
			log.Printf("RestGate: Too many Secrets defined. At most there should be 1 secret per key")
		}
		return nil
	}

	if headerKeyLabel == "" { //headerKeyLabel must be defined
		if t.config.Debug == true {
			log.Printf("RestGate: headerKeyLabel is not defined.")
		}
		return nil
	}

	//Default Error Messages
	if t.config.ErrorMessages == nil {
		t.config.ErrorMessages = map[int]map[string]string{
			1:  e.New(1, "No Key Or Secret", "", "com.github.pjebs.restgate").Render(),
			2:  e.New(2, "Unauthorized Access", "", "com.github.pjebs.restgate").Render(),
			3:  e.New(3, "Please use HTTPS connection", "", "com.github.pjebs.restgate").Render(),
			99: e.New(99, "Software Developers have not setup authentication correctly", "", "com.github.pjebs.restgate").Render(),
		}
	} else {
		if _, ok := t.config.ErrorMessages[1]; !ok {
			t.config.ErrorMessages[1] = e.New(1, "No Key Or Secret", "", "com.github.pjebs.restgate").Render()
		}

		if _, ok := t.config.ErrorMessages[2]; !ok {
			t.config.ErrorMessages[2] = e.New(2, "Unauthorized Access", "", "com.github.pjebs.restgate").Render()
		}

		if _, ok := t.config.ErrorMessages[3]; !ok {
			t.config.ErrorMessages[3] = e.New(3, "Please use HTTPS connection", "", "com.github.pjebs.restgate").Render()
		}

		if _, ok := t.config.ErrorMessages[99]; !ok {
			t.config.ErrorMessages[99] = e.New(99, "Software Developers have not setup authentication correctly", "", "com.github.pjebs.restgate").Render()
		}
	}

	//Check if HTTPS Protection has been turned off
	if t.config.HTTPSProtectionOff {
		//HTTPS Protection is off
		log.Printf("\x1b[31mWARNING: HTTPS Protection is off. This is potentially insecure!\x1b[39;49m")
	}

	if as == Database {

		if numberKeys != 1 { //We need exactly 1 Key (it represents field name in database)
			if t.config.Debug == true {
				log.Printf("RestGate: For Database mode, we need exactly 1 Key which represents the field name in the database table")
			}
			return nil
		}

		//Check if database is set.
		//The developer should ensure a database has been selected (i.e. to prevent "No Database selected" error)
		if t.config.DB == nil { //DB is not set
			if t.config.Debug == true {
				log.Printf("RestGate: Database is not set. Be sure that a database name is selected")
			}
			return nil
		}

		//Check if table is set
		if t.config.TableName == "" { //Table name is not set
			if t.config.Debug == true {
				log.Printf("RestGate: For Database mode, a table name is required")
			}
			return nil
		}

	}

	return t
}

func (self *RESTGate) ServeHTTP(w http.ResponseWriter, req *http.Request, next http.HandlerFunc) {

	//Check if HTTPS Protection has been turned off
	if !self.config.HTTPSProtectionOff {
		//HTTPS Protection is on so we must check it
		if !(strings.EqualFold(req.URL.Scheme, "https") || req.TLS != nil) {
			r := render.New(render.Options{})
			r.JSON(w, http.StatusUnauthorized, self.config.ErrorMessages[3]) //"Please use HTTPS connection"
			return
		}
	}

	//Check key in Header
	key := req.Header.Get(self.headerKeyLabel)
	secret := req.Header.Get(self.headerSecretLabel)

	if key == "" {
		//Authentication Information not included in request
		r := render.New(render.Options{})
		r.JSON(w, http.StatusUnauthorized, self.config.ErrorMessages[1]) //"No Key Or Secret"
		return
	}

	if self.source == Static {

		secretDoesntExist := len(self.config.Secret) == 0
		authenticationPassed := false

		//First search through all keys
		for index, element := range self.config.Key {
			if secureCompare(key, element) { //Key matches

				//Now check if secret matches
				if secretDoesntExist {
					//Authentication PASSED
					authenticationPassed = true
					break
				} else if index > (len(self.config.Secret) - 1) { //Out of Range so corresponding secret doesn't exist
					//Authentication PASSED
					authenticationPassed = true
					break
				} else {
					//Corresponding Secret exists
					if secureCompare(secret, self.config.Secret[index]) {
						//Authentication PASSED
						authenticationPassed = true
						break
					} else {
						//Authentication FAILED
						authenticationPassed = false
						break
					}
				}
			}
		}

		//Authentication FAILED - No Key's matched
		if authenticationPassed == false {
			r := render.New(render.Options{})
			r.JSON(w, http.StatusUnauthorized, self.config.ErrorMessages[2]) //"Unauthorized Access"
			return
		} else { //Authentication PASSED
			if self.config.Context != nil {
				self.config.Context(req, key)
			}
			next(w, req)
		}

	} else if self.source == Database {
		db := self.config.DB

		secretDoesntExists := len(self.config.Secret) == 0 || self.config.Secret[0] == ""

		var preparedStatement string
		if secretDoesntExists {
			if self.config.Postgres == false { //COUNT(*) is definately faster on MYISAM and possibly InnoDB (MySQL engines)
				preparedStatement = fmt.Sprintf("SELECT COUNT(*) FROM `%v` WHERE `%v`=?", self.config.TableName, self.config.Key[0])
			} else {
				preparedStatement = fmt.Sprintf("SELECT COUNT(%v) FROM %v WHERE %v=$1", self.config.Key[0], self.config.TableName, self.config.Key[0])
			}
		} else {
			if self.config.Postgres == false {
				preparedStatement = fmt.Sprintf("SELECT COUNT(*) FROM `%v` WHERE `%v`=? AND `%v`=?", self.config.TableName, self.config.Key[0], self.config.Secret[0])
			} else {
				preparedStatement = fmt.Sprintf("SELECT COUNT(%v) FROM %v WHERE %v=$1 AND %v=$2", self.config.Key[0], self.config.TableName, self.config.Key[0], self.config.Secret[0])
			}
		}

		stmt, err := db.Prepare(preparedStatement)
		if err != nil {
			if self.config.Debug == true {
				log.Printf("RestGate: Run time database error: %+v", err)
			}
			r := render.New(render.Options{})
			r.JSON(w, http.StatusUnauthorized, self.config.ErrorMessages[99]) //"Software Developers have not setup authentication correctly"
			return
		}
		defer stmt.Close()

		var count int //stores query result
		if secretDoesntExists {
			err = stmt.QueryRow(key).Scan(&count)
		} else {
			err = stmt.QueryRow(key, secret).Scan(&count)
		}

		// log.Printf("result error: %+v", err)
		// log.Printf("count: %+v", count)

		if err == nil && count == 1 {
			//Authentication PASSED
			if self.config.Context != nil {
				self.config.Context(req, key)
			}
			next(w, req)
		} else { //==sql.ErrNoRows or count == 0
			//Something went wrong
			if self.config.Debug == true && count > 1 {
				log.Printf("RestGate: Database query returned more than 1 identical Key. Make sure the KEY field in the table is set to UNIQUE")
			}
			r := render.New(render.Options{})
			r.JSON(w, http.StatusUnauthorized, self.config.ErrorMessages[2]) //"Unauthorized Access"
			return
		}

	} else {
		r := render.New(render.Options{})
		r.JSON(w, http.StatusUnauthorized, self.config.ErrorMessages[99]) //"Software Developers have not setup authentication correctly"
		return
	}

}

// secureCompare performs a constant time compare of two strings to limit timing attacks.
func secureCompare(given string, actual string) bool {
	if subtle.ConstantTimeEq(int32(len(given)), int32(len(actual))) == 1 {
		return subtle.ConstantTimeCompare([]byte(given), []byte(actual)) == 1
	} else {
		/* Securely compare actual to itself to keep constant time, but always return false */
		return subtle.ConstantTimeCompare([]byte(actual), []byte(actual)) == 1 && false
	}
}
<commit_msg>Update HTTPS protection for GAE-Flexible Environment<commit_after>package restgate

/*
|--------------------------------------------------------------------------
| WARNING
|--------------------------------------------------------------------------
| Never Set HTTPSProtectionOff=true In Production.
| The Key and Password will be exposed and highly unsecure otherwise!
| The database server should also use HTTPS Connection and be hidden away
|
*/

/*
Thanks to Ido Ben-Natan ("IdoBn") for postgres fix.
Thanks to Jeremy Saenz & Brendon Murphy for timing-attack protection
*/

import (
	// "errors"
	"crypto/subtle"
	"database/sql"
	"fmt"
	"log"
	"net/http"
	"strings"

	e "github.com/pjebs/jsonerror"
	"gopkg.in/unrolled/render.v1"
)

type AuthenticationSource int

const (
	Static   AuthenticationSource = 0
	Database                      = 1
)

//When AuthenticationSource=Static, Key(s)=Actual Key and Secret(s)=Actual Secret.
//When AuthenticationSource=Database, Key[0]=Key_Column and Secret[0]=Secret_Column.
type Config struct {
	*sql.DB
	Key                    []string
	Secret                 []string //Can be "" but not recommended
	TableName              string
	ErrorMessages          map[int]map[string]string
	Context                func(r *http.Request, authenticatedKey string)
	Debug                  bool
	Postgres               bool
	HTTPSProtectionOff     bool //Default is HTTPS Protection On
	GAE_ManagedEnvironment bool //Default is false. ALWAYS KEEP THIS FALSE UNLESS you are using Google App Engine-Managed Environment
}

type RESTGate struct {
	headerKeyLabel    string
	headerSecretLabel string
	source            AuthenticationSource
	config            Config
}

func New(headerKeyLabel string, headerSecretLabel string, as AuthenticationSource, config Config) *RESTGate {
	t := &RESTGate{headerKeyLabel: headerKeyLabel, headerSecretLabel: headerSecretLabel, source: as, config: config}
	log.Printf("RestGate initializing")

	numberKeys := len(t.config.Key)
	numberSecrets := len(t.config.Secret)

	if numberKeys == 0 { //Key is not set
		if t.config.Debug == true {
			log.Printf("RestGate: Key is not set")
		}
		return nil
	}

	if numberSecrets > numberKeys { //Too many Secret's defined
		if t.config.Debug == true {
			log.Printf("RestGate: Too many Secrets defined. At most there should be 1 secret per key")
		}
		return nil
	}

	if headerKeyLabel == "" { //headerKeyLabel must be defined
		if t.config.Debug == true {
			log.Printf("RestGate: headerKeyLabel is not defined.")
		}
		return nil
	}

	//Default Error Messages
	if t.config.ErrorMessages == nil {
		t.config.ErrorMessages = map[int]map[string]string{
			1:  e.New(1, "No Key Or Secret", "", "com.github.pjebs.restgate").Render(),
			2:  e.New(2, "Unauthorized Access", "", "com.github.pjebs.restgate").Render(),
			3:  e.New(3, "Please use HTTPS connection", "", "com.github.pjebs.restgate").Render(),
			99: e.New(99, "Software Developers have not setup authentication correctly", "", "com.github.pjebs.restgate").Render(),
		}
	} else {
		if _, ok := t.config.ErrorMessages[1]; !ok {
			t.config.ErrorMessages[1] = e.New(1, "No Key Or Secret", "", "com.github.pjebs.restgate").Render()
		}

		if _, ok := t.config.ErrorMessages[2]; !ok {
			t.config.ErrorMessages[2] = e.New(2, "Unauthorized Access", "", "com.github.pjebs.restgate").Render()
		}

		if _, ok := t.config.ErrorMessages[3]; !ok {
			t.config.ErrorMessages[3] = e.New(3, "Please use HTTPS connection", "", "com.github.pjebs.restgate").Render()
		}

		if _, ok := t.config.ErrorMessages[99]; !ok {
			t.config.ErrorMessages[99] = e.New(99, "Software Developers have not setup authentication correctly", "", "com.github.pjebs.restgate").Render()
		}
	}

	//Check if HTTPS Protection has been turned off
	if t.config.HTTPSProtectionOff {
		//HTTPS Protection is off
		log.Printf("\x1b[31mWARNING: HTTPS Protection is off. This is potentially insecure!\x1b[39;49m")
	}

	if t.config.GAE_ManagedEnvironment {
		//HTTPS Protection is off
		log.Printf("\x1b[31mWARNING: Set GAE_ManagedEnvironment to false UNLESS you are using Google App Engine-Managed Environment. This is potentially insecure!\x1b[39;49m")
	}

	if as == Database {

		if numberKeys != 1 { //We need exactly 1 Key (it represents field name in database)
			if t.config.Debug == true {
				log.Printf("RestGate: For Database mode, we need exactly 1 Key which represents the field name in the database table")
			}
			return nil
		}

		//Check if database is set.
		//The developer should ensure a database has been selected (i.e. to prevent "No Database selected" error)
		if t.config.DB == nil { //DB is not set
			if t.config.Debug == true {
				log.Printf("RestGate: Database is not set. Be sure that a database name is selected")
			}
			return nil
		}

		//Check if table is set
		if t.config.TableName == "" { //Table name is not set
			if t.config.Debug == true {
				log.Printf("RestGate: For Database mode, a table name is required")
			}
			return nil
		}

	}

	return t
}

func (self *RESTGate) ServeHTTP(w http.ResponseWriter, req *http.Request, next http.HandlerFunc) {

	//Check if HTTPS Protection has been turned off
	if !self.config.HTTPSProtectionOff {
		//HTTPS Protection is on so we must check it

		if self.config.GAE_ManagedEnvironment == true {
			if req.Header.Get("X-AppEngine-Https") != "on" {
				r := render.New(render.Options{})
				r.JSON(w, http.StatusUnauthorized, self.config.ErrorMessages[3]) //"Please use HTTPS connection"
				return
			}
		} else {
			if !(strings.EqualFold(req.URL.Scheme, "https") || req.TLS != nil) {
				r := render.New(render.Options{})
				r.JSON(w, http.StatusUnauthorized, self.config.ErrorMessages[3]) //"Please use HTTPS connection"
				return
			}
		}
	}

	//Check key in Header
	key := req.Header.Get(self.headerKeyLabel)
	secret := req.Header.Get(self.headerSecretLabel)

	if key == "" {
		//Authentication Information not included in request
		r := render.New(render.Options{})
		r.JSON(w, http.StatusUnauthorized, self.config.ErrorMessages[1]) //"No Key Or Secret"
		return
	}

	if self.source == Static {

		secretDoesntExist := len(self.config.Secret) == 0
		authenticationPassed := false

		//First search through all keys
		for index, element := range self.config.Key {
			if secureCompare(key, element) { //Key matches

				//Now check if secret matches
				if secretDoesntExist {
					//Authentication PASSED
					authenticationPassed = true
					break
				} else if index > (len(self.config.Secret) - 1) { //Out of Range so corresponding secret doesn't exist
					//Authentication PASSED
					authenticationPassed = true
					break
				} else {
					//Corresponding Secret exists
					if secureCompare(secret, self.config.Secret[index]) {
						//Authentication PASSED
						authenticationPassed = true
						break
					} else {
						//Authentication FAILED
						authenticationPassed = false
						break
					}
				}
			}
		}

		//Authentication FAILED - No Key's matched
		if authenticationPassed == false {
			r := render.New(render.Options{})
			r.JSON(w, http.StatusUnauthorized, self.config.ErrorMessages[2]) //"Unauthorized Access"
			return
		} else { //Authentication PASSED
			if self.config.Context != nil {
				self.config.Context(req, key)
			}
			next(w, req)
		}

	} else if self.source == Database {
		db := self.config.DB

		secretDoesntExists := len(self.config.Secret) == 0 || self.config.Secret[0] == ""

		var preparedStatement string
		if secretDoesntExists {
			if self.config.Postgres == false { //COUNT(*) is definately faster on MYISAM and possibly InnoDB (MySQL engines)
				preparedStatement = fmt.Sprintf("SELECT COUNT(*) FROM `%v` WHERE `%v`=?", self.config.TableName, self.config.Key[0])
			} else {
				preparedStatement = fmt.Sprintf("SELECT COUNT(%v) FROM %v WHERE %v=$1", self.config.Key[0], self.config.TableName, self.config.Key[0])
			}
		} else {
			if self.config.Postgres == false {
				preparedStatement = fmt.Sprintf("SELECT COUNT(*) FROM `%v` WHERE `%v`=? AND `%v`=?", self.config.TableName, self.config.Key[0], self.config.Secret[0])
			} else {
				preparedStatement = fmt.Sprintf("SELECT COUNT(%v) FROM %v WHERE %v=$1 AND %v=$2", self.config.Key[0], self.config.TableName, self.config.Key[0], self.config.Secret[0])
			}
		}

		stmt, err := db.Prepare(preparedStatement)
		if err != nil {
			if self.config.Debug == true {
				log.Printf("RestGate: Run time database error: %+v", err)
			}
			r := render.New(render.Options{})
			r.JSON(w, http.StatusUnauthorized, self.config.ErrorMessages[99]) //"Software Developers have not setup authentication correctly"
			return
		}
		defer stmt.Close()

		var count int //stores query result
		if secretDoesntExists {
			err = stmt.QueryRow(key).Scan(&count)
		} else {
			err = stmt.QueryRow(key, secret).Scan(&count)
		}

		// log.Printf("result error: %+v", err)
		// log.Printf("count: %+v", count)

		if err == nil && count == 1 {
			//Authentication PASSED
			if self.config.Context != nil {
				self.config.Context(req, key)
			}
			next(w, req)
		} else { //==sql.ErrNoRows or count == 0
			//Something went wrong
			if self.config.Debug == true && count > 1 {
				log.Printf("RestGate: Database query returned more than 1 identical Key. Make sure the KEY field in the table is set to UNIQUE")
			}
			r := render.New(render.Options{})
			r.JSON(w, http.StatusUnauthorized, self.config.ErrorMessages[2]) //"Unauthorized Access"
			return
		}

	} else {
		r := render.New(render.Options{})
		r.JSON(w, http.StatusUnauthorized, self.config.ErrorMessages[99]) //"Software Developers have not setup authentication correctly"
		return
	}

}

// secureCompare performs a constant time compare of two strings to limit timing attacks.
func secureCompare(given string, actual string) bool {
	if subtle.ConstantTimeEq(int32(len(given)), int32(len(actual))) == 1 {
		return subtle.ConstantTimeCompare([]byte(given), []byte(actual)) == 1
	} else {
		/* Securely compare actual to itself to keep constant time, but always return false */
		return subtle.ConstantTimeCompare([]byte(actual), []byte(actual)) == 1 && false
	}
}
<|endoftext|>


<commit_before>// Copyright (c) 2016 Uber Technologies, Inc.
//
// Permission is hereby granted, free of charge, to any person obtaining a copy
// of this software and associated documentation files (the "Software"), to deal
// in the Software without restriction, including without limitation the rights
// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
// copies of the Software, and to permit persons to whom the Software is
// furnished to do so, subject to the following conditions:
//
// The above copyright notice and this permission notice shall be included in
// all copies or substantial portions of the Software.
//
// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
// THE SOFTWARE.

package retry

import (
	"errors"
	"math"
	"time"

	xerrors "github.com/m3db/m3/src/x/errors"

	"github.com/uber-go/tally"
)

var (
	// ErrWhileConditionFalse is returned when the while condition to a while retry
	// method evaluates false.
	ErrWhileConditionFalse = errors.New("retry while condition evaluated to false")
)

type retrier struct {
	initialBackoff time.Duration
	backoffFactor  float64
	maxBackoff     time.Duration
	maxRetries     int
	forever        bool
	jitter         bool
	rngFn          RngFn
	sleepFn        func(t time.Duration)
	metrics        retrierMetrics
}

type retrierMetrics struct {
	success            tally.Counter
	successLatency     tally.Timer
	errors             tally.Counter
	errorsNotRetryable tally.Counter
	errorsFinal        tally.Counter
	errorsLatency      tally.Timer
	retries            tally.Counter
}

// NewRetrier creates a new retrier.
func NewRetrier(opts Options) Retrier {
	scope := opts.MetricsScope()
	errorTags := struct {
		retryable    map[string]string
		notRetryable map[string]string
	}{
		map[string]string{
			"type": "retryable",
		},
		map[string]string{
			"type": "not-retryable",
		},
	}
	return &retrier{
		initialBackoff: opts.InitialBackoff(),
		backoffFactor:  opts.BackoffFactor(),
		maxBackoff:     opts.MaxBackoff(),
		maxRetries:     opts.MaxRetries(),
		forever:        opts.Forever(),
		jitter:         opts.Jitter(),
		rngFn:          opts.RngFn(),
		sleepFn:        time.Sleep,
		metrics: retrierMetrics{
			success:            scope.Counter("success"),
			successLatency:     scope.Timer("success-latency"),
			errors:             scope.Tagged(errorTags.retryable).Counter("errors"),
			errorsNotRetryable: scope.Tagged(errorTags.notRetryable).Counter("errors"),
			errorsFinal:        scope.Counter("errors-final"),
			errorsLatency:      scope.Timer("errors-latency"),
			retries:            scope.Counter("retries"),
		},
	}
}

func (r *retrier) Attempt(fn Fn) error {
	return r.attempt(nil, fn)
}

func (r *retrier) AttemptWhile(continueFn ContinueFn, fn Fn) error {
	return r.attempt(continueFn, fn)
}

func (r *retrier) attempt(continueFn ContinueFn, fn Fn) error {
	attempt := 0

	if continueFn != nil && !continueFn(attempt) {
		return ErrWhileConditionFalse
	}

	start := time.Now()
	err := fn()
	duration := time.Since(start)
	attempt++
	if err == nil {
		r.metrics.successLatency.Record(duration)
		r.metrics.success.Inc(1)
		return nil
	}
	r.metrics.errorsLatency.Record(duration)
	if xerrors.IsNonRetryableError(err) {
		r.metrics.errorsNotRetryable.Inc(1)
		return err
	}
	r.metrics.errors.Inc(1)

	for i := 1; r.forever || i <= r.maxRetries; i++ {
		r.sleepFn(time.Duration(BackoffNanos(
			i,
			r.jitter,
			r.backoffFactor,
			r.initialBackoff,
			r.maxBackoff,
			r.rngFn,
		)))

		if continueFn != nil && !continueFn(attempt) {
			return ErrWhileConditionFalse
		}

		r.metrics.retries.Inc(1)
		start := time.Now()
		err = fn()
		duration := time.Since(start)
		attempt++
		if err == nil {
			r.metrics.successLatency.Record(duration)
			r.metrics.success.Inc(1)
			return nil
		}
		r.metrics.errorsLatency.Record(duration)
		if xerrors.IsNonRetryableError(err) {
			r.metrics.errorsNotRetryable.Inc(1)
			return err
		}
		r.metrics.errors.Inc(1)
	}
	r.metrics.errorsFinal.Inc(1)

	return err
}

// BackoffNanos calculates the backoff for a retry in nanoseconds.
func BackoffNanos(
	retry int,
	jitter bool,
	backoffFactor float64,
	initialBackoff time.Duration,
	maxBackoff time.Duration,
	rngFn RngFn,
) int64 {
	backoff := initialBackoff.Nanoseconds()
	if retry >= 1 {
		backoffFloat64 := float64(backoff) * math.Pow(backoffFactor, float64(retry-1))
		// math.Inf is also larger than math.MaxInt64.
		if backoffFloat64 > math.MaxInt64 {
			return maxBackoff.Nanoseconds()
		}
		backoff = int64(backoffFloat64)
	}
	// Validate the value of backoff to make sure Int63n() does not panic.
	if jitter && backoff >= 2 {
		half := backoff / 2
		backoff = half + rngFn(half)
	}
	if maxBackoff := maxBackoff.Nanoseconds(); backoff > maxBackoff {
		backoff = maxBackoff
	}
	return backoff
}
<commit_msg>Convert Retrier interface metrics from timers to histograms (#1937)<commit_after>// Copyright (c) 2016 Uber Technologies, Inc.
//
// Permission is hereby granted, free of charge, to any person obtaining a copy
// of this software and associated documentation files (the "Software"), to deal
// in the Software without restriction, including without limitation the rights
// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
// copies of the Software, and to permit persons to whom the Software is
// furnished to do so, subject to the following conditions:
//
// The above copyright notice and this permission notice shall be included in
// all copies or substantial portions of the Software.
//
// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
// THE SOFTWARE.

package retry

import (
	"errors"
	"math"
	"time"

	xerrors "github.com/m3db/m3/src/x/errors"

	"github.com/uber-go/tally"
)

var (
	// ErrWhileConditionFalse is returned when the while condition to a while retry
	// method evaluates false.
	ErrWhileConditionFalse = errors.New("retry while condition evaluated to false")
)

type retrier struct {
	initialBackoff time.Duration
	backoffFactor  float64
	maxBackoff     time.Duration
	maxRetries     int
	forever        bool
	jitter         bool
	rngFn          RngFn
	sleepFn        func(t time.Duration)
	metrics        retrierMetrics
}

type retrierMetrics struct {
	success            tally.Counter
	successLatency     tally.Histogram
	errors             tally.Counter
	errorsNotRetryable tally.Counter
	errorsFinal        tally.Counter
	errorsLatency      tally.Histogram
	retries            tally.Counter
}

// NewRetrier creates a new retrier.
func NewRetrier(opts Options) Retrier {
	scope := opts.MetricsScope()
	errorTags := struct {
		retryable    map[string]string
		notRetryable map[string]string
	}{
		map[string]string{
			"type": "retryable",
		},
		map[string]string{
			"type": "not-retryable",
		},
	}

	return &retrier{
		initialBackoff: opts.InitialBackoff(),
		backoffFactor:  opts.BackoffFactor(),
		maxBackoff:     opts.MaxBackoff(),
		maxRetries:     opts.MaxRetries(),
		forever:        opts.Forever(),
		jitter:         opts.Jitter(),
		rngFn:          opts.RngFn(),
		sleepFn:        time.Sleep,
		metrics: retrierMetrics{
			success:            scope.Counter("success"),
			successLatency:     histogramWithDurationBuckets(scope, "success-latency"),
			errors:             scope.Tagged(errorTags.retryable).Counter("errors"),
			errorsNotRetryable: scope.Tagged(errorTags.notRetryable).Counter("errors"),
			errorsFinal:        scope.Counter("errors-final"),
			errorsLatency:      histogramWithDurationBuckets(scope, "errors-latency"),
			retries:            scope.Counter("retries"),
		},
	}
}

func (r *retrier) Attempt(fn Fn) error {
	return r.attempt(nil, fn)
}

func (r *retrier) AttemptWhile(continueFn ContinueFn, fn Fn) error {
	return r.attempt(continueFn, fn)
}

func (r *retrier) attempt(continueFn ContinueFn, fn Fn) error {
	attempt := 0

	if continueFn != nil && !continueFn(attempt) {
		return ErrWhileConditionFalse
	}

	start := time.Now()
	err := fn()
	duration := time.Since(start)
	attempt++
	if err == nil {
		r.metrics.successLatency.RecordDuration(duration)
		r.metrics.success.Inc(1)
		return nil
	}
	r.metrics.errorsLatency.RecordDuration(duration)
	if xerrors.IsNonRetryableError(err) {
		r.metrics.errorsNotRetryable.Inc(1)
		return err
	}
	r.metrics.errors.Inc(1)

	for i := 1; r.forever || i <= r.maxRetries; i++ {
		r.sleepFn(time.Duration(BackoffNanos(
			i,
			r.jitter,
			r.backoffFactor,
			r.initialBackoff,
			r.maxBackoff,
			r.rngFn,
		)))

		if continueFn != nil && !continueFn(attempt) {
			return ErrWhileConditionFalse
		}

		r.metrics.retries.Inc(1)
		start := time.Now()
		err = fn()
		duration := time.Since(start)
		attempt++
		if err == nil {
			r.metrics.successLatency.RecordDuration(duration)
			r.metrics.success.Inc(1)
			return nil
		}
		r.metrics.errorsLatency.RecordDuration(duration)
		if xerrors.IsNonRetryableError(err) {
			r.metrics.errorsNotRetryable.Inc(1)
			return err
		}
		r.metrics.errors.Inc(1)
	}
	r.metrics.errorsFinal.Inc(1)

	return err
}

// BackoffNanos calculates the backoff for a retry in nanoseconds.
func BackoffNanos(
	retry int,
	jitter bool,
	backoffFactor float64,
	initialBackoff time.Duration,
	maxBackoff time.Duration,
	rngFn RngFn,
) int64 {
	backoff := initialBackoff.Nanoseconds()
	if retry >= 1 {
		backoffFloat64 := float64(backoff) * math.Pow(backoffFactor, float64(retry-1))
		// math.Inf is also larger than math.MaxInt64.
		if backoffFloat64 > math.MaxInt64 {
			return maxBackoff.Nanoseconds()
		}
		backoff = int64(backoffFloat64)
	}
	// Validate the value of backoff to make sure Int63n() does not panic.
	if jitter && backoff >= 2 {
		half := backoff / 2
		backoff = half + rngFn(half)
	}
	if maxBackoff := maxBackoff.Nanoseconds(); backoff > maxBackoff {
		backoff = maxBackoff
	}
	return backoff
}

// histogramWithDurationBuckets returns a histogram with the standard duration buckets.
func histogramWithDurationBuckets(scope tally.Scope, name string) tally.Histogram {
	sub := scope.Tagged(map[string]string{
		// Bump the version if the histogram buckets need to be changed to avoid overlapping buckets
		// in the same query causing errors.
		"schema": "v1",
	})
	buckets := append(tally.DurationBuckets{0, time.Millisecond},
		tally.MustMakeExponentialDurationBuckets(2*time.Millisecond, 1.5, 30)...)
	return sub.Histogram(name, buckets)
}
<|endoftext|>


<commit_before>/*
// Copyright (c) 2016 Intel Corporation
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//      http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
*/

/* TODO

5. Install kernel
12. Make most output from osprepare optional
*/

package main

import (
	"context"
	"flag"
	"fmt"
	"os"
	"os/exec"
	"os/signal"
	"syscall"
)

func init() {
	flag.Usage = func() {
		fmt.Fprintf(os.Stderr, "Usage of %s:\n\n", os.Args[0])
		fmt.Fprintf(os.Stderr, "%s [prepare|start|stop|quit|status|connect|delete]\n\n", os.Args[0])
		fmt.Fprintln(os.Stderr, "- prepare : creates a new VM")
		fmt.Fprintln(os.Stderr, "- start : boots a stopped VM")
		fmt.Fprintln(os.Stderr, "- stop : cleanly powers down a running VM")
		fmt.Fprintln(os.Stderr, "- quit : quits a running VM")
		fmt.Fprintln(os.Stderr, "- status : prints status information about the ciao-down VM")
		fmt.Fprintln(os.Stderr, "- connect : connects to the VM via SSH")
		fmt.Fprintln(os.Stderr, "- delete : shuts down and deletes the VM")
	}
}

func vmFlags(fs *flag.FlagSet, memGB, CPUs *int) {
	*memGB, *CPUs = getMemAndCpus()
	fs.IntVar(memGB, "mem", *memGB, "Gigabytes of RAM allocated to VM")
	fs.IntVar(CPUs, "cpus", *CPUs, "VCPUs assignged to VM")
}

func prepareFlags() (memGB int, CPUs int, debug bool, err error) {
	fs := flag.NewFlagSet("prepare", flag.ExitOnError)
	vmFlags(fs, &memGB, &CPUs)
	fs.BoolVar(&debug, "debug", false, "Enables debug mode")
	if err := fs.Parse(flag.Args()[1:]); err != nil {
		return -1, -1, false, err
	}

	return memGB, CPUs, debug, nil
}

func startFlags() (memGB int, CPUs int, err error) {
	fs := flag.NewFlagSet("start", flag.ExitOnError)
	vmFlags(fs, &memGB, &CPUs)
	if err := fs.Parse(flag.Args()[1:]); err != nil {
		return -1, -1, err
	}

	return memGB, CPUs, nil
}

func downloadProgress(p progress) {
	if p.totalMB >= 0 {
		fmt.Printf("Downloaded %d MB of %d\n", p.downloadedMB, p.totalMB)
	} else {
		fmt.Printf("Downloaded %d MB\n", p.downloadedMB)
	}
}

func prepare(ctx context.Context, errCh chan error) {
	if !hostSupportsNestedKVM() {
		errCh <- fmt.Errorf("Nested KVM is not enabled.  Please enable try again.")
		return
	}

	fmt.Println("Checking environment")

	memGB, CPUs, debug, err := prepareFlags()
	if err != nil {
		errCh <- err
		return
	}

	ws, err := prepareEnv(ctx)
	if err != nil {
		errCh <- err
		return
	}

	_, err = os.Stat(ws.instanceDir)
	if err == nil {
		errCh <- fmt.Errorf("instance already exists")
		return
	}

	fmt.Println("Installing host dependencies")
	installDeps(ctx)

	err = os.MkdirAll(ws.instanceDir, 0755)
	if err != nil {
		errCh <- fmt.Errorf("unable to create cache dir: %v", err)
		return
	}

	err = prepareSSHKeys(ctx, ws)
	if err != nil {
		errCh <- err
		return
	}

	failed := true
	defer func() {
		if failed {
			_ = os.RemoveAll(ws.instanceDir)
		}
	}()

	qcowPath, err := downloadUbuntu(ctx, ws.ciaoDir, downloadProgress)
	if err != nil {
		errCh <- err
		return
	}

	err = buildISOImage(ctx, ws.instanceDir, ws, debug)
	if err != nil {
		errCh <- err
		return
	}

	err = createRootfs(ctx, qcowPath, ws.instanceDir)
	if err != nil {
		errCh <- err
		return
	}

	fmt.Printf("Booting VM with %d GB RAM and %d cpus\n", memGB, CPUs)

	err = bootVM(ctx, ws, memGB, CPUs)
	if err != nil {
		errCh <- err
		return
	}

	err = manageInstallation(ctx, ws.instanceDir, ws)
	if err != nil {
		errCh <- err
		return
	}
	errCh <- nil
	failed = false
	fmt.Println("VM successfully created!")
	fmt.Println("Type ciao-down connect to start using it.")
}

func start(ctx context.Context, errCh chan error) {
	if !hostSupportsNestedKVM() {
		errCh <- fmt.Errorf("Nested KVM is not enabled.  Please enable try again.")
		return
	}

	memGB, CPUs, err := startFlags()
	if err != nil {
		errCh <- err
		return
	}

	ws, err := prepareEnv(ctx)
	if err != nil {
		errCh <- err
		return
	}

	fmt.Printf("Booting VM with %d GB RAM and %d cpus\n", memGB, CPUs)

	err = bootVM(ctx, ws, memGB, CPUs)
	if err != nil {
		errCh <- err
		return
	}

	fmt.Println("VM Started")

	errCh <- err
}

func stop(ctx context.Context, errCh chan error) {
	ws, err := prepareEnv(ctx)
	if err != nil {
		errCh <- err
		return
	}

	err = stopVM(ctx, ws.instanceDir)
	if err != nil {
		errCh <- err
		return
	}

	fmt.Println("VM Stopped")

	errCh <- err
}

func quit(ctx context.Context, errCh chan error) {
	ws, err := prepareEnv(ctx)
	if err != nil {
		errCh <- err
		return
	}

	err = quitVM(ctx, ws.instanceDir)
	if err != nil {
		errCh <- err
		return
	}

	fmt.Println("VM Quit")

	errCh <- err
}

func status(ctx context.Context, errCh chan error) {
	ws, err := prepareEnv(ctx)
	if err != nil {
		errCh <- err
		return
	}

	statusVM(ctx, ws.instanceDir, ws.keyPath)
	errCh <- err
}

func connect(ctx context.Context, errCh chan error) {
	ws, err := prepareEnv(ctx)
	if err != nil {
		errCh <- err
		return
	}

	path, err := exec.LookPath("ssh")
	if err != nil {
		errCh <- fmt.Errorf("Unable to locate ssh binary")
		return
	}

	err = syscall.Exec(path, []string{path,
		"-q", "-o", "UserKnownHostsFile=/dev/null",
		"-o", "StrictHostKeyChecking=no",
		"-i", ws.keyPath,
		"127.0.0.1", "-p", "10022"},
		os.Environ())
	errCh <- err
}

func delete(ctx context.Context, errCh chan error) {
	ws, err := prepareEnv(ctx)
	if err != nil {
		errCh <- err
		return
	}

	_ = quitVM(ctx, ws.instanceDir)
	err = os.RemoveAll(ws.instanceDir)
	if err != nil {
		errCh <- fmt.Errorf("unable to delete instance: %v", err)
		return
	}

	errCh <- nil
}

func runCommand(signalCh <-chan os.Signal) error {
	var err error

	errCh := make(chan error)
	ctx, cancelFunc := context.WithCancel(context.Background())
	switch os.Args[1] {
	case "prepare":
		go prepare(ctx, errCh)
	case "start":
		go start(ctx, errCh)
	case "stop":
		go stop(ctx, errCh)
	case "quit":
		go quit(ctx, errCh)
	case "status":
		go status(ctx, errCh)
	case "connect":
		go connect(ctx, errCh)
	case "delete":
		go delete(ctx, errCh)
	}
	select {
	case <-signalCh:
		cancelFunc()
		err = <-errCh
	case err = <-errCh:
		cancelFunc()
	}

	return err
}

func main() {
	flag.Parse()
	if len(os.Args) < 2 ||
		!(os.Args[1] == "prepare" || os.Args[1] == "start" || os.Args[1] == "stop" ||
			os.Args[1] == "quit" || os.Args[1] == "status" ||
			os.Args[1] == "connect" || os.Args[1] == "delete") {
		flag.Usage()
		os.Exit(1)
	}

	signalCh := make(chan os.Signal, 1)
	signal.Notify(signalCh, syscall.SIGINT, syscall.SIGTERM, syscall.SIGHUP)

	if err := runCommand(signalCh); err != nil {
		fmt.Fprintf(os.Stderr, "%v\n", err)
		os.Exit(1)
	}
}
<commit_msg>ciao-down: Fix golint errors in ciao_down.go<commit_after>/*
// Copyright (c) 2016 Intel Corporation
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//      http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
*/

/* TODO

5. Install kernel
12. Make most output from osprepare optional
*/

package main

import (
	"context"
	"flag"
	"fmt"
	"os"
	"os/exec"
	"os/signal"
	"syscall"
)

func init() {
	flag.Usage = func() {
		fmt.Fprintf(os.Stderr, "Usage of %s:\n\n", os.Args[0])
		fmt.Fprintf(os.Stderr, "%s [prepare|start|stop|quit|status|connect|delete]\n\n", os.Args[0])
		fmt.Fprintln(os.Stderr, "- prepare : creates a new VM")
		fmt.Fprintln(os.Stderr, "- start : boots a stopped VM")
		fmt.Fprintln(os.Stderr, "- stop : cleanly powers down a running VM")
		fmt.Fprintln(os.Stderr, "- quit : quits a running VM")
		fmt.Fprintln(os.Stderr, "- status : prints status information about the ciao-down VM")
		fmt.Fprintln(os.Stderr, "- connect : connects to the VM via SSH")
		fmt.Fprintln(os.Stderr, "- delete : shuts down and deletes the VM")
	}
}

func vmFlags(fs *flag.FlagSet, memGB, CPUs *int) {
	*memGB, *CPUs = getMemAndCpus()
	fs.IntVar(memGB, "mem", *memGB, "Gigabytes of RAM allocated to VM")
	fs.IntVar(CPUs, "cpus", *CPUs, "VCPUs assignged to VM")
}

func prepareFlags() (memGB int, CPUs int, debug bool, err error) {
	fs := flag.NewFlagSet("prepare", flag.ExitOnError)
	vmFlags(fs, &memGB, &CPUs)
	fs.BoolVar(&debug, "debug", false, "Enables debug mode")
	if err := fs.Parse(flag.Args()[1:]); err != nil {
		return -1, -1, false, err
	}

	return memGB, CPUs, debug, nil
}

func startFlags() (memGB int, CPUs int, err error) {
	fs := flag.NewFlagSet("start", flag.ExitOnError)
	vmFlags(fs, &memGB, &CPUs)
	if err := fs.Parse(flag.Args()[1:]); err != nil {
		return -1, -1, err
	}

	return memGB, CPUs, nil
}

func downloadProgress(p progress) {
	if p.totalMB >= 0 {
		fmt.Printf("Downloaded %d MB of %d\n", p.downloadedMB, p.totalMB)
	} else {
		fmt.Printf("Downloaded %d MB\n", p.downloadedMB)
	}
}

func prepare(ctx context.Context, errCh chan error) {
	if !hostSupportsNestedKVM() {
		errCh <- fmt.Errorf("nested KVM is not enabled.  Please enable and try again")
		return
	}

	fmt.Println("Checking environment")

	memGB, CPUs, debug, err := prepareFlags()
	if err != nil {
		errCh <- err
		return
	}

	ws, err := prepareEnv(ctx)
	if err != nil {
		errCh <- err
		return
	}

	_, err = os.Stat(ws.instanceDir)
	if err == nil {
		errCh <- fmt.Errorf("instance already exists")
		return
	}

	fmt.Println("Installing host dependencies")
	installDeps(ctx)

	err = os.MkdirAll(ws.instanceDir, 0755)
	if err != nil {
		errCh <- fmt.Errorf("unable to create cache dir: %v", err)
		return
	}

	err = prepareSSHKeys(ctx, ws)
	if err != nil {
		errCh <- err
		return
	}

	failed := true
	defer func() {
		if failed {
			_ = os.RemoveAll(ws.instanceDir)
		}
	}()

	qcowPath, err := downloadUbuntu(ctx, ws.ciaoDir, downloadProgress)
	if err != nil {
		errCh <- err
		return
	}

	err = buildISOImage(ctx, ws.instanceDir, ws, debug)
	if err != nil {
		errCh <- err
		return
	}

	err = createRootfs(ctx, qcowPath, ws.instanceDir)
	if err != nil {
		errCh <- err
		return
	}

	fmt.Printf("Booting VM with %d GB RAM and %d cpus\n", memGB, CPUs)

	err = bootVM(ctx, ws, memGB, CPUs)
	if err != nil {
		errCh <- err
		return
	}

	err = manageInstallation(ctx, ws.instanceDir, ws)
	if err != nil {
		errCh <- err
		return
	}
	errCh <- nil
	failed = false
	fmt.Println("VM successfully created!")
	fmt.Println("Type ciao-down connect to start using it.")
}

func start(ctx context.Context, errCh chan error) {
	if !hostSupportsNestedKVM() {
		errCh <- fmt.Errorf("nested KVM is not enabled.  Please enable and try again")
		return
	}

	memGB, CPUs, err := startFlags()
	if err != nil {
		errCh <- err
		return
	}

	ws, err := prepareEnv(ctx)
	if err != nil {
		errCh <- err
		return
	}

	fmt.Printf("Booting VM with %d GB RAM and %d cpus\n", memGB, CPUs)

	err = bootVM(ctx, ws, memGB, CPUs)
	if err != nil {
		errCh <- err
		return
	}

	fmt.Println("VM Started")

	errCh <- err
}

func stop(ctx context.Context, errCh chan error) {
	ws, err := prepareEnv(ctx)
	if err != nil {
		errCh <- err
		return
	}

	err = stopVM(ctx, ws.instanceDir)
	if err != nil {
		errCh <- err
		return
	}

	fmt.Println("VM Stopped")

	errCh <- err
}

func quit(ctx context.Context, errCh chan error) {
	ws, err := prepareEnv(ctx)
	if err != nil {
		errCh <- err
		return
	}

	err = quitVM(ctx, ws.instanceDir)
	if err != nil {
		errCh <- err
		return
	}

	fmt.Println("VM Quit")

	errCh <- err
}

func status(ctx context.Context, errCh chan error) {
	ws, err := prepareEnv(ctx)
	if err != nil {
		errCh <- err
		return
	}

	statusVM(ctx, ws.instanceDir, ws.keyPath)
	errCh <- err
}

func connect(ctx context.Context, errCh chan error) {
	ws, err := prepareEnv(ctx)
	if err != nil {
		errCh <- err
		return
	}

	path, err := exec.LookPath("ssh")
	if err != nil {
		errCh <- fmt.Errorf("Unable to locate ssh binary")
		return
	}

	err = syscall.Exec(path, []string{path,
		"-q", "-o", "UserKnownHostsFile=/dev/null",
		"-o", "StrictHostKeyChecking=no",
		"-i", ws.keyPath,
		"127.0.0.1", "-p", "10022"},
		os.Environ())
	errCh <- err
}

func delete(ctx context.Context, errCh chan error) {
	ws, err := prepareEnv(ctx)
	if err != nil {
		errCh <- err
		return
	}

	_ = quitVM(ctx, ws.instanceDir)
	err = os.RemoveAll(ws.instanceDir)
	if err != nil {
		errCh <- fmt.Errorf("unable to delete instance: %v", err)
		return
	}

	errCh <- nil
}

func runCommand(signalCh <-chan os.Signal) error {
	var err error

	errCh := make(chan error)
	ctx, cancelFunc := context.WithCancel(context.Background())
	switch os.Args[1] {
	case "prepare":
		go prepare(ctx, errCh)
	case "start":
		go start(ctx, errCh)
	case "stop":
		go stop(ctx, errCh)
	case "quit":
		go quit(ctx, errCh)
	case "status":
		go status(ctx, errCh)
	case "connect":
		go connect(ctx, errCh)
	case "delete":
		go delete(ctx, errCh)
	}
	select {
	case <-signalCh:
		cancelFunc()
		err = <-errCh
	case err = <-errCh:
		cancelFunc()
	}

	return err
}

func main() {
	flag.Parse()
	if len(os.Args) < 2 ||
		!(os.Args[1] == "prepare" || os.Args[1] == "start" || os.Args[1] == "stop" ||
			os.Args[1] == "quit" || os.Args[1] == "status" ||
			os.Args[1] == "connect" || os.Args[1] == "delete") {
		flag.Usage()
		os.Exit(1)
	}

	signalCh := make(chan os.Signal, 1)
	signal.Notify(signalCh, syscall.SIGINT, syscall.SIGTERM, syscall.SIGHUP)

	if err := runCommand(signalCh); err != nil {
		fmt.Fprintf(os.Stderr, "%v\n", err)
		os.Exit(1)
	}
}
<|endoftext|>