<commit_before>package concurrency_sqlite_test

import (
	"fmt"
	"math/rand"
	"os"
	"path/filepath"
	"runtime"
	"strconv"
	"time"

	. "github.com/onsi/ginkgo"
	. "github.com/onsi/gomega"

	_ "github.com/mattn/go-sqlite3"

	boshlog "github.com/cloudfoundry/bosh-utils/logger"
	common "github.com/maximilien/bosh-softlayer-cpi/common"
	bslcvmpool "github.com/maximilien/bosh-softlayer-cpi/softlayer/vm/pool"
	testhelpers "github.com/maximilien/softlayer-go/test_helpers"
)

var (
	SQLITE_DB_FOLDER    = "/tmp/concurrency_sqlite_test"
	SQLITE_DB_FILE      = "vm_pool.sqlite"
	SQLITE_DB_FILE_PATH = filepath.Join(SQLITE_DB_FOLDER, SQLITE_DB_FILE)
	stemcellUuid        = "fake_stemcell_uuid"
	domain              = "softlayer.com"
	c                   = make(chan string, 0)
	logger              = boshlog.NewLogger(boshlog.LevelInfo)
)

var _ = Describe("Concurrency test for Sqlite DB", func() {
	var (
		err error
	)

	BeforeEach(func() {
		common.SetOSEnvVariable("OS_RELOAD_ENABLED", "TRUE")
		common.SetOSEnvVariable("SQLITE_DB_FOLDER", SQLITE_DB_FOLDER)
		common.SetOSEnvVariable("SQLITE_DB_FILE", SQLITE_DB_FILE)

		testhelpers.TIMEOUT = 35 * time.Minute
		testhelpers.POLLING_INTERVAL = 10 * time.Second

		err = os.RemoveAll(SQLITE_DB_FOLDER)
		Expect(err).ToNot(HaveOccurred())

		err = bslcvmpool.InitVMPoolDB(bslcvmpool.DB_RETRY_TIMEOUT, bslcvmpool.DB_RETRY_INTERVAL, logger)
		Expect(err).ToNot(HaveOccurred())

		populateDB()

	})

	AfterEach(func() {
		err = os.RemoveAll(SQLITE_DB_FOLDER)
		Expect(err).ToNot(HaveOccurred())

	})

	Context("Manipulate DB concurrently", func() {

		It("Manipulate DB concurrently", func() {
			runtime.GOMAXPROCS(1)

			for i := 1000; i <= 2000; i += 1000 {
				go insertVMInfo(i)
				go updateVMInfoByID()
				go queryVMInfobyID()
				go queryVMInfobyAgentID()
			}

			time.Sleep(480 * time.Second)

		})
	})

})

func populateDB() {
	db, err := bslcvmpool.OpenDB(SQLITE_DB_FILE_PATH)
	Expect(err).ToNot(HaveOccurred())

	for i := 0; i < 1000; i++ {
		vmID := i
		hostname := fmt.Sprintf("concurrency_test_%d", vmID)
		agentID := strconv.Itoa(i)
		vmInfoDB := bslcvmpool.NewVMInfoDB(vmID, hostname+"."+domain, "t", stemcellUuid, agentID, logger, db)
		defer vmInfoDB.CloseDB()

		err = vmInfoDB.InsertVMInfo(bslcvmpool.DB_RETRY_TIMEOUT, bslcvmpool.DB_RETRY_INTERVAL)
		Expect(err).ToNot(HaveOccurred())
	}
}

func insertVMInfo(init int) {
	db, err := bslcvmpool.OpenDB(SQLITE_DB_FILE_PATH)
	Expect(err).ToNot(HaveOccurred())

	for i := init; i < init+1000; i++ {
		vmID := i
		hostname := fmt.Sprintf("concurrency_test_%d", vmID)
		agentID := strconv.Itoa(i)
		vmInfoDB := bslcvmpool.NewVMInfoDB(vmID, hostname+"."+domain, "t", stemcellUuid, agentID, logger, db)
		defer vmInfoDB.CloseDB()

		err = vmInfoDB.InsertVMInfo(bslcvmpool.DB_RETRY_TIMEOUT, bslcvmpool.DB_RETRY_INTERVAL)
		Expect(err).ToNot(HaveOccurred())

		time.Sleep(time.Duration(rand.Int63n(100)) * time.Millisecond)
	}

}

func updateVMInfoByID() {
	db, err := bslcvmpool.OpenDB(SQLITE_DB_FILE_PATH)
	Expect(err).ToNot(HaveOccurred())

	for i := 0; i < 1000; i++ {
		vmID := rand.Intn(1000)
		hostname := fmt.Sprintf("concurrency_test_%d", vmID)
		agentID := strconv.Itoa(i)
		vmInfoDB := bslcvmpool.NewVMInfoDB(vmID, hostname+"."+domain, "f", stemcellUuid, agentID, logger, db)
		defer vmInfoDB.CloseDB()

		err = vmInfoDB.UpdateVMInfoByID(bslcvmpool.DB_RETRY_TIMEOUT, bslcvmpool.DB_RETRY_INTERVAL)
		Expect(err).ToNot(HaveOccurred())

		time.Sleep(time.Duration(rand.Int63n(100)) * time.Millisecond)
	}

}

func queryVMInfobyID() {
	db, err := bslcvmpool.OpenDB(SQLITE_DB_FILE_PATH)
	Expect(err).ToNot(HaveOccurred())

	for i := 0; i < 1000; i++ {
		vmID := rand.Intn(1000)
		hostname := fmt.Sprintf("concurrency_test_%d", vmID)
		agentID := strconv.Itoa(i)
		vmInfoDB := bslcvmpool.NewVMInfoDB(vmID, hostname+"."+domain, "t", stemcellUuid, agentID, logger, db)
		defer vmInfoDB.CloseDB()

		err = vmInfoDB.QueryVMInfobyID(bslcvmpool.DB_RETRY_TIMEOUT, bslcvmpool.DB_RETRY_INTERVAL)
		Expect(err).ToNot(HaveOccurred())

		time.Sleep(time.Duration(rand.Int63n(100)) * time.Millisecond)
	}

}

func queryVMInfobyAgentID() {
	db, err := bslcvmpool.OpenDB(SQLITE_DB_FILE_PATH)
	Expect(err).ToNot(HaveOccurred())

	for i := 0; i < 1000; i++ {
		vmID := rand.Intn(1000)
		hostname := fmt.Sprintf("concurrency_test_%d", vmID)
		agentID := strconv.Itoa(i)
		vmInfoDB := bslcvmpool.NewVMInfoDB(vmID, hostname+"."+domain, "t", stemcellUuid, agentID, logger, db)
		defer vmInfoDB.CloseDB()

		err = vmInfoDB.QueryVMInfobyAgentID(bslcvmpool.DB_RETRY_TIMEOUT, bslcvmpool.DB_RETRY_INTERVAL)
		Expect(err).ToNot(HaveOccurred())

		time.Sleep(time.Duration(rand.Int63n(100)) * time.Millisecond)
	}

}
<commit_msg>Reduce the amount of request due to low garden performance<commit_after>package concurrency_sqlite_test

import (
	"fmt"
	"math/rand"
	"os"
	"path/filepath"
	"runtime"
	"strconv"
	"time"

	. "github.com/onsi/ginkgo"
	. "github.com/onsi/gomega"

	_ "github.com/mattn/go-sqlite3"

	boshlog "github.com/cloudfoundry/bosh-utils/logger"
	common "github.com/maximilien/bosh-softlayer-cpi/common"
	bslcvmpool "github.com/maximilien/bosh-softlayer-cpi/softlayer/vm/pool"
	testhelpers "github.com/maximilien/softlayer-go/test_helpers"
)

var (
	SQLITE_DB_FOLDER    = "/tmp/concurrency_sqlite_test"
	SQLITE_DB_FILE      = "vm_pool.sqlite"
	SQLITE_DB_FILE_PATH = filepath.Join(SQLITE_DB_FOLDER, SQLITE_DB_FILE)
	stemcellUuid        = "fake_stemcell_uuid"
	domain              = "softlayer.com"
	c                   = make(chan string, 0)
	logger              = boshlog.NewLogger(boshlog.LevelInfo)
)

var _ = Describe("Concurrency test for Sqlite DB", func() {
	var (
		err error
	)

	BeforeEach(func() {
		common.SetOSEnvVariable("OS_RELOAD_ENABLED", "TRUE")
		common.SetOSEnvVariable("SQLITE_DB_FOLDER", SQLITE_DB_FOLDER)
		common.SetOSEnvVariable("SQLITE_DB_FILE", SQLITE_DB_FILE)

		testhelpers.TIMEOUT = 35 * time.Minute
		testhelpers.POLLING_INTERVAL = 10 * time.Second

		err = os.RemoveAll(SQLITE_DB_FOLDER)
		Expect(err).ToNot(HaveOccurred())

		err = bslcvmpool.InitVMPoolDB(bslcvmpool.DB_RETRY_TIMEOUT, bslcvmpool.DB_RETRY_INTERVAL, logger)
		Expect(err).ToNot(HaveOccurred())

		populateDB()

	})

	AfterEach(func() {
		err = os.RemoveAll(SQLITE_DB_FOLDER)
		Expect(err).ToNot(HaveOccurred())

	})

	Context("Manipulate DB concurrently", func() {

		It("Manipulate DB concurrently", func() {
			runtime.GOMAXPROCS(1)

			for i := 100; i <= 500; i += 100 {
				go insertVMInfo(i)
				go updateVMInfoByID()
				go queryVMInfobyID()
				go queryVMInfobyAgentID()
			}

			time.Sleep(240 * time.Second)

		})
	})

})

func populateDB() {
	db, err := bslcvmpool.OpenDB(SQLITE_DB_FILE_PATH)
	Expect(err).ToNot(HaveOccurred())

	for i := 0; i < 100; i++ {
		vmID := i
		hostname := fmt.Sprintf("concurrency_test_%d", vmID)
		agentID := strconv.Itoa(i)
		vmInfoDB := bslcvmpool.NewVMInfoDB(vmID, hostname+"."+domain, "t", stemcellUuid, agentID, logger, db)
		defer vmInfoDB.CloseDB()

		err = vmInfoDB.InsertVMInfo(bslcvmpool.DB_RETRY_TIMEOUT, bslcvmpool.DB_RETRY_INTERVAL)
		Expect(err).ToNot(HaveOccurred())
	}
}

func insertVMInfo(init int) {
	db, err := bslcvmpool.OpenDB(SQLITE_DB_FILE_PATH)
	Expect(err).ToNot(HaveOccurred())

	for i := init; i < init+100; i++ {
		vmID := i
		hostname := fmt.Sprintf("concurrency_test_%d", vmID)
		agentID := strconv.Itoa(i)
		vmInfoDB := bslcvmpool.NewVMInfoDB(vmID, hostname+"."+domain, "t", stemcellUuid, agentID, logger, db)
		defer vmInfoDB.CloseDB()

		err = vmInfoDB.InsertVMInfo(bslcvmpool.DB_RETRY_TIMEOUT, bslcvmpool.DB_RETRY_INTERVAL)
		Expect(err).ToNot(HaveOccurred())

		time.Sleep(time.Duration(rand.Int63n(100)) * time.Millisecond)
	}

}

func updateVMInfoByID() {
	db, err := bslcvmpool.OpenDB(SQLITE_DB_FILE_PATH)
	Expect(err).ToNot(HaveOccurred())

	for i := 0; i < 100; i++ {
		vmID := rand.Intn(100)
		hostname := fmt.Sprintf("concurrency_test_%d", vmID)
		agentID := strconv.Itoa(i)
		vmInfoDB := bslcvmpool.NewVMInfoDB(vmID, hostname+"."+domain, "f", stemcellUuid, agentID, logger, db)
		defer vmInfoDB.CloseDB()

		err = vmInfoDB.UpdateVMInfoByID(bslcvmpool.DB_RETRY_TIMEOUT, bslcvmpool.DB_RETRY_INTERVAL)
		Expect(err).ToNot(HaveOccurred())

		time.Sleep(time.Duration(rand.Int63n(100)) * time.Millisecond)
	}

}

func queryVMInfobyID() {
	db, err := bslcvmpool.OpenDB(SQLITE_DB_FILE_PATH)
	Expect(err).ToNot(HaveOccurred())

	for i := 0; i < 100; i++ {
		vmID := rand.Intn(100)
		hostname := fmt.Sprintf("concurrency_test_%d", vmID)
		agentID := strconv.Itoa(i)
		vmInfoDB := bslcvmpool.NewVMInfoDB(vmID, hostname+"."+domain, "t", stemcellUuid, agentID, logger, db)
		defer vmInfoDB.CloseDB()

		err = vmInfoDB.QueryVMInfobyID(bslcvmpool.DB_RETRY_TIMEOUT, bslcvmpool.DB_RETRY_INTERVAL)
		Expect(err).ToNot(HaveOccurred())

		time.Sleep(time.Duration(rand.Int63n(100)) * time.Millisecond)
	}

}

func queryVMInfobyAgentID() {
	db, err := bslcvmpool.OpenDB(SQLITE_DB_FILE_PATH)
	Expect(err).ToNot(HaveOccurred())

	for i := 0; i < 100; i++ {
		vmID := rand.Intn(100)
		hostname := fmt.Sprintf("concurrency_test_%d", vmID)
		agentID := strconv.Itoa(i)
		vmInfoDB := bslcvmpool.NewVMInfoDB(vmID, hostname+"."+domain, "t", stemcellUuid, agentID, logger, db)
		defer vmInfoDB.CloseDB()

		err = vmInfoDB.QueryVMInfobyAgentID(bslcvmpool.DB_RETRY_TIMEOUT, bslcvmpool.DB_RETRY_INTERVAL)
		Expect(err).ToNot(HaveOccurred())

		time.Sleep(time.Duration(rand.Int63n(100)) * time.Millisecond)
	}

}
<|endoftext|>


<commit_before>package main

import (
	"flag"
	"fmt"
	"os"
	"strconv"
	"strings"

	"github.com/dokku/dokku/plugins/common"
	"github.com/dokku/dokku/plugins/logs"
)

const (
	helpHeader = `Usage: dokku logs[:COMMAND]

Manage log integration for an app

Additional commands:`

	helpContent = `
    logs [-h] [-t] [-n num] [-q] [-p process] <app>, Display recent log output
    logs:failed [<app>], Shows the last failed deploy logs
`
)

func main() {
	flag.Usage = usage
	flag.Parse()

	cmd := flag.Arg(0)
	switch cmd {
	case "logs":
		args := flag.NewFlagSet("logs", flag.ExitOnError)
		help := args.Bool("h", false, "-h: print help for the command")
		var num int64
		args.Int64Var(&num, "n", 0, "-n: the number of lines to display")
		args.Int64Var(&num, "num", 0, "--num: the number of lines to display")
		var ps string
		args.StringVar(&ps, "p", "", "-p: only display logs from the given process")
		args.StringVar(&ps, "ps", "", "--p: only display logs from the given process")
		var tail bool
		args.BoolVar(&tail, "t", true, "-t: continually stream logs")
		args.BoolVar(&tail, "tail", true, "--tail: continually stream logs")
		var quiet bool
		args.BoolVar(&quiet, "q", true, "-q: display raw logs without colors, time and names")
		args.BoolVar(&quiet, "quiet", true, "--quiet: display raw logs without colors, time and names")
		args.Parse(flag.Args()[1:])
		if *help {
			usage()
			return
		}
		appName := args.Arg(0)
		err := logs.CommandDefault(appName, num, ps, tail, quiet)
		if err != nil {
			common.LogFail(err.Error())
		}
	case "logs:help":
		usage()
	case "help":
		command := common.NewShellCmd(fmt.Sprintf("ps -o command= %d", os.Getppid()))
		command.ShowOutput = false
		output, err := command.Output()

		if err == nil && strings.Contains(string(output), "--all") {
			fmt.Println(helpContent)
		} else {
			fmt.Print("\n    logs, Manage log integration for an app\n")
		}
	default:
		dokkuNotImplementExitCode, err := strconv.Atoi(os.Getenv("DOKKU_NOT_IMPLEMENTED_EXIT"))
		if err != nil {
			fmt.Println("failed to retrieve DOKKU_NOT_IMPLEMENTED_EXIT environment variable")
			dokkuNotImplementExitCode = 10
		}
		os.Exit(dokkuNotImplementExitCode)
	}
}

func usage() {
	common.CommandUsage(helpHeader, helpContent)
}
<commit_msg>chore: formatting<commit_after>package main

import (
	"flag"
	"fmt"
	"os"
	"strconv"
	"strings"

	"github.com/dokku/dokku/plugins/common"
	"github.com/dokku/dokku/plugins/logs"
)

const (
	helpHeader = `Usage: dokku logs[:COMMAND]

Manage log integration for an app

Additional commands:`

	helpContent = `
    logs [-h] [-t] [-n num] [-q] [-p process] <app>, Display recent log output
    logs:failed [<app>], Shows the last failed deploy logs
`
)

func main() {
	flag.Usage = usage
	flag.Parse()

	cmd := flag.Arg(0)
	switch cmd {
	case "logs":
		args := flag.NewFlagSet("logs", flag.ExitOnError)
		var num int64
		var ps string
		var tail bool
		var quiet bool
		help := args.Bool("h", false, "-h: print help for the command")
		args.Int64Var(&num, "n", 0, "-n: the number of lines to display")
		args.Int64Var(&num, "num", 0, "--num: the number of lines to display")
		args.StringVar(&ps, "p", "", "-p: only display logs from the given process")
		args.StringVar(&ps, "ps", "", "--p: only display logs from the given process")
		args.BoolVar(&tail, "t", true, "-t: continually stream logs")
		args.BoolVar(&tail, "tail", true, "--tail: continually stream logs")
		args.BoolVar(&quiet, "q", true, "-q: display raw logs without colors, time and names")
		args.BoolVar(&quiet, "quiet", true, "--quiet: display raw logs without colors, time and names")
		args.Parse(flag.Args()[1:])
		if *help {
			usage()
			return
		}
		appName := args.Arg(0)
		err := logs.CommandDefault(appName, num, ps, tail, quiet)
		if err != nil {
			common.LogFail(err.Error())
		}
	case "logs:help":
		usage()
	case "help":
		command := common.NewShellCmd(fmt.Sprintf("ps -o command= %d", os.Getppid()))
		command.ShowOutput = false
		output, err := command.Output()

		if err == nil && strings.Contains(string(output), "--all") {
			fmt.Println(helpContent)
		} else {
			fmt.Print("\n    logs, Manage log integration for an app\n")
		}
	default:
		dokkuNotImplementExitCode, err := strconv.Atoi(os.Getenv("DOKKU_NOT_IMPLEMENTED_EXIT"))
		if err != nil {
			fmt.Println("failed to retrieve DOKKU_NOT_IMPLEMENTED_EXIT environment variable")
			dokkuNotImplementExitCode = 10
		}
		os.Exit(dokkuNotImplementExitCode)
	}
}

func usage() {
	common.CommandUsage(helpHeader, helpContent)
}
<|endoftext|>


<commit_before>// Copyright 2015 Comcast Cable Communications Management, LLC

// 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.

// Started from https://raw.githubusercontent.com/jordan-wright/gophish

package routes

import (
	"github.com/Comcast/traffic_control/traffic_ops/goto2/api"
	"github.com/Comcast/traffic_control/traffic_ops/goto2/auth"
	"github.com/Comcast/traffic_control/traffic_ops/goto2/crconfig"
	"github.com/Comcast/traffic_control/traffic_ops/goto2/csconfig"
	output "github.com/Comcast/traffic_control/traffic_ops/goto2/output_format"

	"encoding/json"
	"github.com/gorilla/mux"
	"io/ioutil"
	"log"
	"net/http"
)

const apiPath = "/api/2.0/"

// CreateAdminRouter creates the routes for handling requests to the web interface.
// This function returns an http.Handler to be used in http.ListenAndServe().
func CreateRouter() http.Handler {
	router := mux.NewRouter().StrictSlash(true)
	router.HandleFunc("/login", auth.Login).Methods("POST")
	router.HandleFunc(apiPath+"{table}", auth.Use(optionsHandler, auth.DONTRequireLogin)).Methods("OPTIONS")
	router.HandleFunc(apiPath+"{table}/{id}", auth.Use(optionsHandler, auth.DONTRequireLogin)).Methods("OPTIONS")
	router.HandleFunc("/config/cr/{cdn}/CRConfig.json", auth.Use(handleCRConfig, auth.RequireLogin))
	router.HandleFunc("/config/csconfig/{hostname}", auth.Use(handleCSConfig, auth.RequireLogin))
	addApiHandlers(router)
	return auth.Use(router.ServeHTTP, auth.GetContext)
}

func setHeaders(w http.ResponseWriter, methods api.ApiMethods) {
	w.Header().Set("Access-Control-Allow-Origin", "*")
	w.Header().Set("Access-Control-Allow-Methods", methods.String())
	w.Header().Set("Access-Control-Allow-Headers", "Origin, Authorization, X-Requested-With, Content-Type")
}

func optionsHandler(w http.ResponseWriter, r *http.Request) {
	vars := mux.Vars(r)
	table := vars["table"]
	route := table
	if _, ok := vars["id"]; ok {
		route += "/{id}"
	}

	apiHandlers := api.ApiHandlers()
	if tableHandlers, ok := apiHandlers[route]; !ok {
		w.WriteHeader(http.StatusNotFound)
		return
	} else {
		setHeaders(w, tableHandlers.Methods())
		w.WriteHeader(http.StatusOK)
	}
}

func addApiHandlers(router *mux.Router) {
	for route, funcs := range api.ApiHandlers() {
		wrapRouter := func(f api.ApiHandlerFunc) func(w http.ResponseWriter, r *http.Request) {
			return func(w http.ResponseWriter, r *http.Request) {
				setHeaders(w, funcs.Methods())
				body, err := ioutil.ReadAll(r.Body)
				response, err := f(mux.Vars(r), body)
				if err != nil {
					log.Println(err)
				}
				jresponse := output.MakeApiResponse(response, nil, err)
				w.Header().Set("Content-Type", "application/json")
				enc := json.NewEncoder(w)
				enc.Encode(jresponse)
			}
		}
		for method, f := range funcs {
			router.HandleFunc(apiPath+route, auth.Use(wrapRouter(f), auth.RequireLogin)).Methods(method.String())
		}
	}
}

func handleCRConfig(w http.ResponseWriter, r *http.Request) {
	vars := mux.Vars(r)
	cdn := vars["cdn"]
	resp, _ := crconfig.GetCRConfig(cdn)
	enc := json.NewEncoder(w)
	enc.Encode(resp)
}

func handleCSConfig(w http.ResponseWriter, r *http.Request) {
	vars := mux.Vars(r)
	hostName := vars["hostname"]
	resp, _ := csconfig.GetCSConfig(hostName)
	enc := json.NewEncoder(w)
	enc.Encode(resp)
}
<commit_msg>Move routes wrap lambda to its own function<commit_after>// Copyright 2015 Comcast Cable Communications Management, LLC

// 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.

// Started from https://raw.githubusercontent.com/jordan-wright/gophish

package routes

import (
	"github.com/Comcast/traffic_control/traffic_ops/goto2/api"
	"github.com/Comcast/traffic_control/traffic_ops/goto2/auth"
	"github.com/Comcast/traffic_control/traffic_ops/goto2/crconfig"
	"github.com/Comcast/traffic_control/traffic_ops/goto2/csconfig"
	output "github.com/Comcast/traffic_control/traffic_ops/goto2/output_format"

	"encoding/json"
	"github.com/gorilla/mux"
	"io/ioutil"
	"log"
	"net/http"
)

const apiPath = "/api/2.0/"

// CreateRouter creates the routes for handling requests to the web interface.
// This function returns an http.Handler to be used in http.ListenAndServe().
func CreateRouter() http.Handler {
	router := mux.NewRouter().StrictSlash(true)
	router.HandleFunc("/login", auth.Login).Methods("POST")
	router.HandleFunc(apiPath+"{table}", auth.Use(optionsHandler, auth.DONTRequireLogin)).Methods("OPTIONS")
	router.HandleFunc(apiPath+"{table}/{id}", auth.Use(optionsHandler, auth.DONTRequireLogin)).Methods("OPTIONS")
	router.HandleFunc("/config/cr/{cdn}/CRConfig.json", auth.Use(handleCRConfig, auth.RequireLogin))
	router.HandleFunc("/config/csconfig/{hostname}", auth.Use(handleCSConfig, auth.RequireLogin))
	addApiHandlers(router)
	return auth.Use(router.ServeHTTP, auth.GetContext)
}

// setHeaders writes the universal headers needed by all routes,
// along with the given accepted HTTP Methods.
func setHeaders(w http.ResponseWriter, methods api.ApiMethods) {
	w.Header().Set("Access-Control-Allow-Origin", "*")
	w.Header().Set("Access-Control-Allow-Methods", methods.String())
	w.Header().Set("Access-Control-Allow-Headers", "Origin, Authorization, X-Requested-With, Content-Type")
}

// optionsHandler handles HTTP OPTIONS requests, writing the
// appropriate options and an HTTP OK.
func optionsHandler(w http.ResponseWriter, r *http.Request) {
	vars := mux.Vars(r)
	table := vars["table"]
	route := table
	if _, ok := vars["id"]; ok {
		route += "/{id}"
	}

	apiHandlers := api.ApiHandlers()
	if tableHandlers, ok := apiHandlers[route]; !ok {
		w.WriteHeader(http.StatusNotFound)
		return
	} else {
		setHeaders(w, tableHandlers.Methods())
		w.WriteHeader(http.StatusOK)
	}
}

// wrapApiHandler takes an api.ApiHandlerFunc and returns a func with the
// signature expected by http.HandleFunc.
//
// The returned func sets the headers, reads the params, calls the
// handler with the params, encodes the handlers response, and writes it.
func wrapApiHandler(f api.ApiHandlerFunc, methods api.ApiMethods) func(w http.ResponseWriter, r *http.Request) {
	return func(w http.ResponseWriter, r *http.Request) {
		setHeaders(w, methods)
		body, err := ioutil.ReadAll(r.Body)
		response, err := f(mux.Vars(r), body)
		if err != nil {
			log.Println(err)
		}
		jresponse := output.MakeApiResponse(response, nil, err)
		w.Header().Set("Content-Type", "application/json")
		enc := json.NewEncoder(w)
		enc.Encode(jresponse)
	}
}

// addApiHandlers adds each API handler to the given router.
func addApiHandlers(router *mux.Router) {
	for route, funcs := range api.ApiHandlers() {
		for method, f := range funcs {
			router.HandleFunc(apiPath+route, auth.Use(wrapApiHandler(f, funcs.Methods()), auth.RequireLogin)).Methods(method.String())
		}
	}
}

// handleCRConfig handles requests to the CRConfig endpoint,
// returning the encoded CRConfig data for the requested CDN.
func handleCRConfig(w http.ResponseWriter, r *http.Request) {
	vars := mux.Vars(r)
	cdn := vars["cdn"]
	resp, _ := crconfig.GetCRConfig(cdn)
	enc := json.NewEncoder(w)
	enc.Encode(resp)
}

// handleCSConfig handles requests to the CSConfig endpoint,
// returning the encoded CSConfig data for the requested host.
func handleCSConfig(w http.ResponseWriter, r *http.Request) {
	vars := mux.Vars(r)
	hostName := vars["hostname"]
	resp, _ := csconfig.GetCSConfig(hostName)
	enc := json.NewEncoder(w)
	enc.Encode(resp)
}
<|endoftext|>


<commit_before>/* This Source Code Form is subject to the terms of the Mozilla Public
 * License, v. 2.0. If a copy of the MPL was not distributed with this
 * file, You can obtain one at http://mozilla.org/MPL/2.0/. */

package components

import (
	"log"
	"reflect"

	"github.com/ernestio/ernestprovider/event"
	"github.com/ernestio/ernestprovider/providers/azure/lb"
	graph "gopkg.in/r3labs/graph.v2"
)

// LB : ..
type LB struct {
	ID string `json:"id"`
	lb.Event
	Base
}

// GetID : returns the component's ID
func (i *LB) GetID() string {
	return i.ComponentID
}

// GetName returns a components name
func (i *LB) GetName() string {
	return i.Name
}

// GetProvider : returns the provider type
func (i *LB) GetProvider() string {
	return i.ProviderType
}

// GetProviderID returns a components provider id
func (i *LB) GetProviderID() string {
	return i.ID
}

// GetType : returns the type of the component
func (i *LB) GetType() string {
	return i.ComponentType
}

// GetState : returns the state of the component
func (i *LB) GetState() string {
	return i.State
}

// SetState : sets the state of the component
func (i *LB) SetState(s string) {
	i.State = s
}

// GetAction : returns the action of the component
func (i *LB) GetAction() string {
	return i.Action
}

// SetAction : Sets the action of the component
func (i *LB) SetAction(s string) {
	i.Action = s
}

// GetGroup : returns the components group
func (i *LB) GetGroup() string {
	return ""
}

// GetTags returns a components tags
func (i *LB) GetTags() map[string]string {
	return i.Tags
}

// GetTag returns a components tag
func (i *LB) GetTag(tag string) string {
	return ""
}

// Diff : diff's the component against another component of the same type
func (i *LB) Diff(c graph.Component) bool {
	cs, ok := c.(*LB)
	if ok {
		if len(i.Tags) != 0 && len(cs.Tags) != 0 {
			if !reflect.DeepEqual(i.Tags, cs.Tags) {
				return true
			}
		}
		if i.Location != cs.Location {
			return true
		}
		if len(i.FrontendIPConfigurations) != len(cs.FrontendIPConfigurations) {
			return true
		}
		for x, cfg := range i.FrontendIPConfigurations {
			if cfg.Name != cs.FrontendIPConfigurations[x].Name {
				return true
			}
			if cfg.PrivateIPAddress != cs.FrontendIPConfigurations[x].PrivateIPAddress {
				return true
			}
			if cfg.PrivateIPAddressAllocation != cs.FrontendIPConfigurations[x].PrivateIPAddressAllocation {
				return true
			}
			if cfg.PublicIPAddress != cs.FrontendIPConfigurations[x].PublicIPAddress {
				return true
			}
			if cfg.Subnet != cs.FrontendIPConfigurations[x].Subnet {
				return true
			}
		}
	}

	return false
}

// Update : updates the provider returned values of a component
func (i *LB) Update(c graph.Component) {
	cs, ok := c.(*LB)
	if ok {
		i.ID = cs.ID
	}
	i.SetDefaultVariables()
}

// Rebuild : rebuilds the component's internal state, such as templated values
func (i *LB) Rebuild(g *graph.Graph) {
	for x := 0; x < len(i.FrontendIPConfigurations); x++ {
		if i.FrontendIPConfigurations[x].PublicIPAddress == "" && i.FrontendIPConfigurations[x].PublicIPAddressID != "" {
			ip := g.GetComponents().ByProviderID(i.FrontendIPConfigurations[x].PublicIPAddressID)
			if ip != nil {
				i.FrontendIPConfigurations[x].PublicIPAddress = ip.GetName()
			}
		}

		if i.FrontendIPConfigurations[x].PublicIPAddressID == "" && i.FrontendIPConfigurations[x].PublicIPAddress != "" {
			i.FrontendIPConfigurations[x].PublicIPAddressID = templPublicIPAddressID(i.FrontendIPConfigurations[x].PublicIPAddress)
		}

		if i.FrontendIPConfigurations[x].Subnet == "" && i.FrontendIPConfigurations[x].SubnetID != "" {
			sub := g.GetComponents().ByProviderID(i.FrontendIPConfigurations[x].SubnetID)
			if sub != nil {
				i.FrontendIPConfigurations[x].Subnet = sub.GetName()
			}
		}

		if i.FrontendIPConfigurations[x].SubnetID == "" && i.FrontendIPConfigurations[x].Subnet != "" {
			i.FrontendIPConfigurations[x].SubnetID = templSubnetID(i.FrontendIPConfigurations[x].Subnet)
		}
	}

	i.SetDefaultVariables()
}

// Dependencies : returns a list of component id's upon which the component depends
func (i *LB) Dependencies() (deps []string) {
	deps = append(deps, TYPERESOURCEGROUP+TYPEDELIMITER+i.ResourceGroupName)
	for _, ip := range i.FrontendIPConfigurations {
		if ip.PublicIPAddress != "" {
			deps = append(deps, TYPEPUBLICIP+TYPEDELIMITER+ip.PublicIPAddress)
		}
		if ip.Subnet != "" {
			deps = append(deps, TYPESUBNET+TYPEDELIMITER+ip.Subnet)
		}
	}
	return
}

// Validate : validates the components values
func (i *LB) Validate() error {
	log.Println("Validating LB")
	val := event.NewValidator()
	return val.Validate(i)
}

// IsStateful : returns true if the component needs to be actioned to be removed.
func (i *LB) IsStateful() bool {
	return true
}

// SetDefaultVariables : sets up the default template variables for a component
func (i *LB) SetDefaultVariables() {
	i.ProviderType = PROVIDERTYPE
	i.ComponentType = TYPELB
	i.ComponentID = TYPELB + TYPEDELIMITER + i.Name
	i.DatacenterName = DATACENTERNAME
	i.DatacenterType = DATACENTERTYPE
	i.DatacenterRegion = DATACENTERREGION
	i.ClientID = CLIENTID
	i.ClientSecret = CLIENTSECRET
	i.TenantID = TENANTID
	i.SubscriptionID = SUBSCRIPTIONID
	i.Environment = ENVIRONMENT
}
<commit_msg>Update frontendipconfigurations<commit_after>/* This Source Code Form is subject to the terms of the Mozilla Public
 * License, v. 2.0. If a copy of the MPL was not distributed with this
 * file, You can obtain one at http://mozilla.org/MPL/2.0/. */

package components

import (
	"log"
	"reflect"

	"github.com/ernestio/ernestprovider/event"
	"github.com/ernestio/ernestprovider/providers/azure/lb"
	graph "gopkg.in/r3labs/graph.v2"
)

// LB : ..
type LB struct {
	ID string `json:"id"`
	lb.Event
	Base
}

// GetID : returns the component's ID
func (i *LB) GetID() string {
	return i.ComponentID
}

// GetName returns a components name
func (i *LB) GetName() string {
	return i.Name
}

// GetProvider : returns the provider type
func (i *LB) GetProvider() string {
	return i.ProviderType
}

// GetProviderID returns a components provider id
func (i *LB) GetProviderID() string {
	return i.ID
}

// GetType : returns the type of the component
func (i *LB) GetType() string {
	return i.ComponentType
}

// GetState : returns the state of the component
func (i *LB) GetState() string {
	return i.State
}

// SetState : sets the state of the component
func (i *LB) SetState(s string) {
	i.State = s
}

// GetAction : returns the action of the component
func (i *LB) GetAction() string {
	return i.Action
}

// SetAction : Sets the action of the component
func (i *LB) SetAction(s string) {
	i.Action = s
}

// GetGroup : returns the components group
func (i *LB) GetGroup() string {
	return ""
}

// GetTags returns a components tags
func (i *LB) GetTags() map[string]string {
	return i.Tags
}

// GetTag returns a components tag
func (i *LB) GetTag(tag string) string {
	return ""
}

// Diff : diff's the component against another component of the same type
func (i *LB) Diff(c graph.Component) bool {
	cs, ok := c.(*LB)
	if ok {
		if len(i.Tags) != 0 && len(cs.Tags) != 0 {
			if !reflect.DeepEqual(i.Tags, cs.Tags) {
				return true
			}
		}
		if i.Location != cs.Location {
			return true
		}
		if len(i.FrontendIPConfigurations) != len(cs.FrontendIPConfigurations) {
			return true
		}
		for x, cfg := range i.FrontendIPConfigurations {
			if cfg.Name != cs.FrontendIPConfigurations[x].Name {
				return true
			}
			if cfg.PrivateIPAddress != cs.FrontendIPConfigurations[x].PrivateIPAddress {
				return true
			}
			if cfg.PrivateIPAddressAllocation != cs.FrontendIPConfigurations[x].PrivateIPAddressAllocation {
				return true
			}
			if cfg.PublicIPAddress != cs.FrontendIPConfigurations[x].PublicIPAddress {
				return true
			}
			if cfg.Subnet != cs.FrontendIPConfigurations[x].Subnet {
				return true
			}
		}
	}

	return false
}

// Update : updates the provider returned values of a component
func (i *LB) Update(c graph.Component) {
	cs, ok := c.(*LB)
	if ok {
		i.ID = cs.ID
		i.FrontendIPConfigurations = cs.FrontendIPConfigurations
	}
	i.SetDefaultVariables()
}

// Rebuild : rebuilds the component's internal state, such as templated values
func (i *LB) Rebuild(g *graph.Graph) {
	for x := 0; x < len(i.FrontendIPConfigurations); x++ {
		if i.FrontendIPConfigurations[x].PublicIPAddress == "" && i.FrontendIPConfigurations[x].PublicIPAddressID != "" {
			ip := g.GetComponents().ByProviderID(i.FrontendIPConfigurations[x].PublicIPAddressID)
			if ip != nil {
				i.FrontendIPConfigurations[x].PublicIPAddress = ip.GetName()
			}
		}

		if i.FrontendIPConfigurations[x].PublicIPAddressID == "" && i.FrontendIPConfigurations[x].PublicIPAddress != "" {
			i.FrontendIPConfigurations[x].PublicIPAddressID = templPublicIPAddressID(i.FrontendIPConfigurations[x].PublicIPAddress)
		}

		if i.FrontendIPConfigurations[x].Subnet == "" && i.FrontendIPConfigurations[x].SubnetID != "" {
			sub := g.GetComponents().ByProviderID(i.FrontendIPConfigurations[x].SubnetID)
			if sub != nil {
				i.FrontendIPConfigurations[x].Subnet = sub.GetName()
			}
		}

		if i.FrontendIPConfigurations[x].SubnetID == "" && i.FrontendIPConfigurations[x].Subnet != "" {
			i.FrontendIPConfigurations[x].SubnetID = templSubnetID(i.FrontendIPConfigurations[x].Subnet)
		}
	}

	i.SetDefaultVariables()
}

// Dependencies : returns a list of component id's upon which the component depends
func (i *LB) Dependencies() (deps []string) {
	deps = append(deps, TYPERESOURCEGROUP+TYPEDELIMITER+i.ResourceGroupName)
	for _, ip := range i.FrontendIPConfigurations {
		if ip.PublicIPAddress != "" {
			deps = append(deps, TYPEPUBLICIP+TYPEDELIMITER+ip.PublicIPAddress)
		}
		if ip.Subnet != "" {
			deps = append(deps, TYPESUBNET+TYPEDELIMITER+ip.Subnet)
		}
	}
	return
}

// Validate : validates the components values
func (i *LB) Validate() error {
	log.Println("Validating LB")
	val := event.NewValidator()
	return val.Validate(i)
}

// IsStateful : returns true if the component needs to be actioned to be removed.
func (i *LB) IsStateful() bool {
	return true
}

// SetDefaultVariables : sets up the default template variables for a component
func (i *LB) SetDefaultVariables() {
	i.ProviderType = PROVIDERTYPE
	i.ComponentType = TYPELB
	i.ComponentID = TYPELB + TYPEDELIMITER + i.Name
	i.DatacenterName = DATACENTERNAME
	i.DatacenterType = DATACENTERTYPE
	i.DatacenterRegion = DATACENTERREGION
	i.ClientID = CLIENTID
	i.ClientSecret = CLIENTSECRET
	i.TenantID = TENANTID
	i.SubscriptionID = SUBSCRIPTIONID
	i.Environment = ENVIRONMENT
}
<|endoftext|>


<commit_before>package server

import (
	"fmt"
	"io/ioutil"
	"strings"

	"github.com/satori/go.uuid"

	"go.pachyderm.com/pachyderm/src/pkg/container"
	"go.pachyderm.com/pachyderm/src/pps"
	"go.pachyderm.com/pachyderm/src/pps/persist"
	"go.pedge.io/google-protobuf"
	"go.pedge.io/protolog"
	"golang.org/x/net/context"
)

var (
	emptyInstance = &google_protobuf.Empty{}
)

type apiServer struct {
	persistAPIClient persist.APIClient
	containerClient  container.Client
}

func newAPIServer(persistAPIClient persist.APIClient, containerClient container.Client) *apiServer {
	return &apiServer{persistAPIClient, containerClient}
}

func (a *apiServer) CreateJob(ctx context.Context, request *pps.CreateJobRequest) (response *pps.Job, err error) {
	persistJob, err := a.persistAPIClient.CreateJob(ctx, jobToPersist(request.Job))
	if err != nil {
		return nil, err
	}
	return persistToJob(persistJob), nil
}

func (a *apiServer) GetJob(ctx context.Context, request *pps.GetJobRequest) (response *pps.Job, err error) {
	persistJob, err := a.persistAPIClient.GetJobByID(ctx, &google_protobuf.StringValue{Value: request.JobId})
	if err != nil {
		return nil, err
	}
	return persistToJob(persistJob), nil
}

func (a *apiServer) GetJobsByPipelineName(ctx context.Context, request *pps.GetJobsByPipelineNameRequest) (response *pps.Jobs, err error) {
	persistPipelines, err := a.persistAPIClient.GetPipelinesByName(ctx, &google_protobuf.StringValue{Value: request.PipelineName})
	if err != nil {
		return nil, err
	}
	var jobs []*pps.Job
	for _, persistPipeline := range persistPipelines.Pipeline {
		persistJobs, err := a.persistAPIClient.GetJobsByPipelineID(ctx, &google_protobuf.StringValue{Value: persistPipeline.Id})
		if err != nil {
			return nil, err
		}
		iJobs := persistToJobs(persistJobs)
		jobs = append(jobs, iJobs.Job...)
	}
	return &pps.Jobs{
		Job: jobs,
	}, nil
}

func (a *apiServer) StartJob(ctx context.Context, request *pps.StartJobRequest) (response *google_protobuf.Empty, err error) {
	persistJob, err := a.persistAPIClient.GetJobByID(ctx, &google_protobuf.StringValue{Value: request.JobId})
	if err != nil {
		return nil, err
	}
	if err := a.startPersistJob(persistJob); err != nil {
		return nil, err
	}
	return emptyInstance, nil
}

func (a *apiServer) GetJobStatus(ctx context.Context, request *pps.GetJobStatusRequest) (response *pps.JobStatus, err error) {
	persistJobStatuses, err := a.persistAPIClient.GetJobStatusesByJobID(ctx, &google_protobuf.StringValue{Value: request.JobId})
	if err != nil {
		return nil, err
	}
	if len(persistJobStatuses.JobStatus) == 0 {
		return nil, fmt.Errorf("pachyderm.pps.server: no job statuses for %s", request.JobId)
	}
	return persistToJobStatus(persistJobStatuses.JobStatus[0]), nil
}

func (a *apiServer) GetJobLogs(request *pps.GetJobLogsRequest, responseServer pps.API_GetJobLogsServer) (err error) {
	persistJobLogs, err := a.persistAPIClient.GetJobLogsByJobID(context.Background(), &google_protobuf.StringValue{Value: request.JobId})
	if err != nil {
		return err
	}
	for _, persistJobLog := range persistJobLogs.JobLog {
		if persistJobLog.OutputStream == request.OutputStream {
			if err := responseServer.Send(&google_protobuf.BytesValue{Value: persistJobLog.Value}); err != nil {
				return err
			}
		}
	}
	return nil
}

func (a *apiServer) CreatePipeline(ctx context.Context, request *pps.CreatePipelineRequest) (response *pps.Pipeline, err error) {
	persistPipeline, err := a.persistAPIClient.CreatePipeline(ctx, pipelineToPersist(request.Pipeline))
	if err != nil {
		return nil, err
	}
	return persistToPipeline(persistPipeline), nil
}

func (a *apiServer) GetPipeline(ctx context.Context, request *pps.GetPipelineRequest) (response *pps.Pipeline, err error) {
	persistPipelines, err := a.persistAPIClient.GetPipelinesByName(ctx, &google_protobuf.StringValue{Value: request.PipelineName})
	if err != nil {
		return nil, err
	}
	if len(persistPipelines.Pipeline) == 0 {
		return nil, fmt.Errorf("pachyderm.pps.server: no piplines for name %s", request.PipelineName)
	}
	return persistToPipeline(persistPipelines.Pipeline[0]), nil
}

func (a *apiServer) GetAllPipelines(ctx context.Context, request *google_protobuf.Empty) (response *pps.Pipelines, err error) {
	persistPipelines, err := a.persistAPIClient.GetAllPipelines(ctx, request)
	if err != nil {
		return nil, err
	}
	pipelineMap := make(map[string]*pps.Pipeline)
	for _, persistPipeline := range persistPipelines.Pipeline {
		// pipelines are ordered newest to oldest, so if we have already
		// seen a pipeline with the same name, it is newer
		if _, ok := pipelineMap[persistPipeline.Name]; !ok {
			pipelineMap[persistPipeline.Name] = persistToPipeline(persistPipeline)
		}
	}
	pipelines := make([]*pps.Pipeline, len(pipelineMap))
	i := 0
	for _, pipeline := range pipelineMap {
		pipelines[i] = pipeline
		i++
	}
	return &pps.Pipelines{
		Pipeline: pipelines,
	}, nil
}

func (a *apiServer) startPersistJob(persistJob *persist.Job) error {
	if _, err := a.persistAPIClient.CreateJobStatus(
		context.Background(),
		&persist.JobStatus{
			JobId: persistJob.Id,
			Type:  pps.JobStatusType_JOB_STATUS_TYPE_STARTED,
		},
	); err != nil {
		return err
	}
	// TODO(pedge): throttling? worker pool?
	go func() {
		if err := a.runJob(persistJob); err != nil {
			protolog.Errorln(err.Error())
			// TODO(pedge): how to handle the error?
			if _, err = a.persistAPIClient.CreateJobStatus(
				context.Background(),
				&persist.JobStatus{
					JobId: persistJob.Id,
					Type:  pps.JobStatusType_JOB_STATUS_TYPE_ERROR,
				},
			); err != nil {
				protolog.Errorln(err.Error())
			}
		} else {
			// TODO(pedge): how to handle the error?
			if _, err = a.persistAPIClient.CreateJobStatus(
				context.Background(),
				&persist.JobStatus{
					JobId: persistJob.Id,
					Type:  pps.JobStatusType_JOB_STATUS_TYPE_SUCCESS,
				},
			); err != nil {
				protolog.Errorln(err.Error())
			}
		}
	}()
	return nil
}

func (a *apiServer) runJob(persistJob *persist.Job) error {
	switch {
	case persistJob.GetTransform() != nil:
		return a.reallyRunJob(strings.Replace(uuid.NewV4().String(), "-", "", -1), persistJob.GetTransform(), persistJob.JobInput, persistJob.JobOutput)
	case persistJob.GetPipelineId() != "":
		persistPipeline, err := a.persistAPIClient.GetPipelineByID(
			context.Background(),
			&google_protobuf.StringValue{Value: persistJob.GetPipelineId()},
		)
		if err != nil {
			return err
		}
		if persistPipeline.Transform == nil {
			return fmt.Errorf("pachyderm.pps.server: transform not set on pipeline %v", persistPipeline)
		}
		return a.reallyRunJob(persistPipeline.Name, persistPipeline.Transform, persistJob.JobInput, persistJob.JobOutput)
	default:
		return fmt.Errorf("pachyderm.pps.server: neither transform or pipeline id set on job %v", persistJob)
	}
}

func (a *apiServer) reallyRunJob(
	name string,
	transform *pps.Transform,
	jobInputs []*pps.JobInput,
	jobOutputs []*pps.JobOutput,
) error {
	return nil
}

// return image name
func (a *apiServer) buildOrPull(name string, transform *pps.Transform) (string, error) {
	image := transform.Image
	if transform.Build != "" {
		image = fmt.Sprintf("ppspipelines/%s", name)
		if err := a.containerClient.Build(
			image,
			transform.Build,
			// TODO(pedge): this will not work, the path to a dockerfile is not real
			container.BuildOptions{
				Dockerfile:   transform.Dockerfile,
				OutputStream: ioutil.Discard,
			},
		); err != nil {
			return "", err
		}
	} else if err := a.containerClient.Pull(
		transform.Image,
		container.PullOptions{},
	); err != nil {
		return "", err
	}
	return image, nil
}
<commit_msg>work on pps api server start job<commit_after>package server

import (
	"fmt"
	"io/ioutil"
	"strings"

	"github.com/satori/go.uuid"

	"go.pachyderm.com/pachyderm/src/pkg/container"
	"go.pachyderm.com/pachyderm/src/pps"
	"go.pachyderm.com/pachyderm/src/pps/persist"
	"go.pedge.io/google-protobuf"
	"go.pedge.io/protolog"
	"golang.org/x/net/context"
)

var (
	emptyInstance = &google_protobuf.Empty{}
)

type apiServer struct {
	persistAPIClient persist.APIClient
	containerClient  container.Client
}

func newAPIServer(persistAPIClient persist.APIClient, containerClient container.Client) *apiServer {
	return &apiServer{persistAPIClient, containerClient}
}

func (a *apiServer) CreateJob(ctx context.Context, request *pps.CreateJobRequest) (response *pps.Job, err error) {
	persistJob, err := a.persistAPIClient.CreateJob(ctx, jobToPersist(request.Job))
	if err != nil {
		return nil, err
	}
	return persistToJob(persistJob), nil
}

func (a *apiServer) GetJob(ctx context.Context, request *pps.GetJobRequest) (response *pps.Job, err error) {
	persistJob, err := a.persistAPIClient.GetJobByID(ctx, &google_protobuf.StringValue{Value: request.JobId})
	if err != nil {
		return nil, err
	}
	return persistToJob(persistJob), nil
}

func (a *apiServer) GetJobsByPipelineName(ctx context.Context, request *pps.GetJobsByPipelineNameRequest) (response *pps.Jobs, err error) {
	persistPipelines, err := a.persistAPIClient.GetPipelinesByName(ctx, &google_protobuf.StringValue{Value: request.PipelineName})
	if err != nil {
		return nil, err
	}
	var jobs []*pps.Job
	for _, persistPipeline := range persistPipelines.Pipeline {
		persistJobs, err := a.persistAPIClient.GetJobsByPipelineID(ctx, &google_protobuf.StringValue{Value: persistPipeline.Id})
		if err != nil {
			return nil, err
		}
		iJobs := persistToJobs(persistJobs)
		jobs = append(jobs, iJobs.Job...)
	}
	return &pps.Jobs{
		Job: jobs,
	}, nil
}

func (a *apiServer) StartJob(ctx context.Context, request *pps.StartJobRequest) (response *google_protobuf.Empty, err error) {
	persistJob, err := a.persistAPIClient.GetJobByID(ctx, &google_protobuf.StringValue{Value: request.JobId})
	if err != nil {
		return nil, err
	}
	if err := a.startPersistJob(persistJob); err != nil {
		return nil, err
	}
	return emptyInstance, nil
}

func (a *apiServer) GetJobStatus(ctx context.Context, request *pps.GetJobStatusRequest) (response *pps.JobStatus, err error) {
	persistJobStatuses, err := a.persistAPIClient.GetJobStatusesByJobID(ctx, &google_protobuf.StringValue{Value: request.JobId})
	if err != nil {
		return nil, err
	}
	if len(persistJobStatuses.JobStatus) == 0 {
		return nil, fmt.Errorf("pachyderm.pps.server: no job statuses for %s", request.JobId)
	}
	return persistToJobStatus(persistJobStatuses.JobStatus[0]), nil
}

func (a *apiServer) GetJobLogs(request *pps.GetJobLogsRequest, responseServer pps.API_GetJobLogsServer) (err error) {
	persistJobLogs, err := a.persistAPIClient.GetJobLogsByJobID(context.Background(), &google_protobuf.StringValue{Value: request.JobId})
	if err != nil {
		return err
	}
	for _, persistJobLog := range persistJobLogs.JobLog {
		if persistJobLog.OutputStream == request.OutputStream {
			if err := responseServer.Send(&google_protobuf.BytesValue{Value: persistJobLog.Value}); err != nil {
				return err
			}
		}
	}
	return nil
}

func (a *apiServer) CreatePipeline(ctx context.Context, request *pps.CreatePipelineRequest) (response *pps.Pipeline, err error) {
	persistPipeline, err := a.persistAPIClient.CreatePipeline(ctx, pipelineToPersist(request.Pipeline))
	if err != nil {
		return nil, err
	}
	return persistToPipeline(persistPipeline), nil
}

func (a *apiServer) GetPipeline(ctx context.Context, request *pps.GetPipelineRequest) (response *pps.Pipeline, err error) {
	persistPipelines, err := a.persistAPIClient.GetPipelinesByName(ctx, &google_protobuf.StringValue{Value: request.PipelineName})
	if err != nil {
		return nil, err
	}
	if len(persistPipelines.Pipeline) == 0 {
		return nil, fmt.Errorf("pachyderm.pps.server: no piplines for name %s", request.PipelineName)
	}
	return persistToPipeline(persistPipelines.Pipeline[0]), nil
}

func (a *apiServer) GetAllPipelines(ctx context.Context, request *google_protobuf.Empty) (response *pps.Pipelines, err error) {
	persistPipelines, err := a.persistAPIClient.GetAllPipelines(ctx, request)
	if err != nil {
		return nil, err
	}
	pipelineMap := make(map[string]*pps.Pipeline)
	for _, persistPipeline := range persistPipelines.Pipeline {
		// pipelines are ordered newest to oldest, so if we have already
		// seen a pipeline with the same name, it is newer
		if _, ok := pipelineMap[persistPipeline.Name]; !ok {
			pipelineMap[persistPipeline.Name] = persistToPipeline(persistPipeline)
		}
	}
	pipelines := make([]*pps.Pipeline, len(pipelineMap))
	i := 0
	for _, pipeline := range pipelineMap {
		pipelines[i] = pipeline
		i++
	}
	return &pps.Pipelines{
		Pipeline: pipelines,
	}, nil
}

func (a *apiServer) startPersistJob(persistJob *persist.Job) error {
	if _, err := a.persistAPIClient.CreateJobStatus(
		context.Background(),
		&persist.JobStatus{
			JobId: persistJob.Id,
			Type:  pps.JobStatusType_JOB_STATUS_TYPE_STARTED,
		},
	); err != nil {
		return err
	}
	// TODO(pedge): throttling? worker pool?
	go func() {
		if err := a.runJob(persistJob); err != nil {
			protolog.Errorln(err.Error())
			// TODO(pedge): how to handle the error?
			if _, err = a.persistAPIClient.CreateJobStatus(
				context.Background(),
				&persist.JobStatus{
					JobId: persistJob.Id,
					Type:  pps.JobStatusType_JOB_STATUS_TYPE_ERROR,
				},
			); err != nil {
				protolog.Errorln(err.Error())
			}
		} else {
			// TODO(pedge): how to handle the error?
			if _, err = a.persistAPIClient.CreateJobStatus(
				context.Background(),
				&persist.JobStatus{
					JobId: persistJob.Id,
					Type:  pps.JobStatusType_JOB_STATUS_TYPE_SUCCESS,
				},
			); err != nil {
				protolog.Errorln(err.Error())
			}
		}
	}()
	return nil
}

func (a *apiServer) runJob(persistJob *persist.Job) error {
	switch {
	case persistJob.GetTransform() != nil:
		return a.reallyRunJob(strings.Replace(uuid.NewV4().String(), "-", "", -1), persistJob.GetTransform(), persistJob.JobInput, persistJob.JobOutput, 1)
	case persistJob.GetPipelineId() != "":
		persistPipeline, err := a.persistAPIClient.GetPipelineByID(
			context.Background(),
			&google_protobuf.StringValue{Value: persistJob.GetPipelineId()},
		)
		if err != nil {
			return err
		}
		if persistPipeline.Transform == nil {
			return fmt.Errorf("pachyderm.pps.server: transform not set on pipeline %v", persistPipeline)
		}
		return a.reallyRunJob(persistPipeline.Name, persistPipeline.Transform, persistJob.JobInput, persistJob.JobOutput, 1)
	default:
		return fmt.Errorf("pachyderm.pps.server: neither transform or pipeline id set on job %v", persistJob)
	}
}

func (a *apiServer) reallyRunJob(
	name string,
	transform *pps.Transform,
	jobInputs []*pps.JobInput,
	jobOutputs []*pps.JobOutput,
	numContainers int,
) error {
	image, err := a.buildOrPull(name, transform)
	if err != nil {
		return err
	}
	var containers []string
	defer func() {
		for _, containerID := range containers {
			_ = a.containerClient.Kill(containerID, container.KillOptions{})
			_ = a.containerClient.Remove(containerID, container.RemoveOptions{})
		}
	}()
	for i := 0; i < numContainers; i++ {
		container, err := a.containerClient.Create(
			image,
			container.CreateOptions{
				Binds:      append(getInputBinds(jobInputs), getOutputBinds(jobOutputs)...),
				HasCommand: len(transform.Cmd) > 0,
			},
		)
		if err != nil {
			return err
		}
		containers = append(containers, container)
	}
	return nil
}

// return image name
func (a *apiServer) buildOrPull(name string, transform *pps.Transform) (string, error) {
	image := transform.Image
	if transform.Build != "" {
		image = fmt.Sprintf("ppspipelines/%s", name)
		if err := a.containerClient.Build(
			image,
			transform.Build,
			// TODO(pedge): this will not work, the path to a dockerfile is not real
			container.BuildOptions{
				Dockerfile:   transform.Dockerfile,
				OutputStream: ioutil.Discard,
			},
		); err != nil {
			return "", err
		}
	} else if err := a.containerClient.Pull(
		transform.Image,
		container.PullOptions{},
	); err != nil {
		return "", err
	}
	return image, nil
}

func getInputBinds(jobInputs []*pps.JobInput) []string {
	var binds []string
	for _, jobInput := range jobInputs {
		if jobInput.GetHostDir() != "" {
			binds = append(binds, getBinds(jobInput.GetHostDir(), fmt.Sprintf("/var/lib/pps/host/%s", jobInput.GetHostDir()), "ro"))
		}
	}
	return binds
}

func getOutputBinds(jobOutputs []*pps.JobOutput) []string {
	var binds []string
	for _, jobOutput := range jobOutputs {
		if jobOutput.GetHostDir() != "" {
			binds = append(binds, getBinds(jobOutput.GetHostDir(), fmt.Sprintf("/var/lib/pps/host/%s", jobOutput.GetHostDir()), "rw"))
		}
	}
	return binds
}

func getBinds(from string, to string, postfix string) string {
	return fmt.Sprintf("%s:%s:%s", from, to, postfix)
}
<|endoftext|>


<commit_before>//this holds the tui functions and information
package main

import (
	"fmt"
	"github.com/jroimartin/gocui"
	"strings"
)

var (
	yCursorOffset   int      = 0
	selectedPodcast *Podcast = nil
	cachedPodcast   []PodcastEntry
)

func listLayout(g *gocui.Gui) error {
	maxX, maxY := g.Size()
	g.Cursor = true
	v, err := g.SetView("list", -1, -1, maxX+1, maxY-2)
	if err != nil {
		if err != gocui.ErrUnknownView {
			return err
		}
	}
	//now print subscribed
	err = printSubscribed(v)
	if err != nil {
		return err
	}
	//now print the player
	err = printPlayer(g)
	//now set current view to main view
	if err := g.SetCurrentView("list"); err != nil {
		return err
	}
	if err := v.SetCursor(0, 1+yCursorOffset); err != nil {
		return err
	}
	return err
}

func printSubscribed(v *gocui.View) error {
	//first clear
	v.Clear()
	//then set properties
	//TODO add in colors from settings
	setProperties(v)
	xMax, _ := v.Size()
	spacing := (xMax - 34) / 3 //43 chracters
	space := strings.Repeat("-", spacing)
	fmt.Fprintf(v, "Podcast Name %s Artist %s Description %s\n", space, space, space)
	for _, item := range globals.Config.Subscribed {
		strin := item.CollectionName + " - " + item.ArtistName + " - " + item.Description
		if len(item.Description+item.CollectionName+item.ArtistName)+6 < xMax {
			//do nothing
		} else { //else truncate string
			strin = strin[0:xMax]
		}
		fmt.Fprintf(v, "%s\n", strin)
	}
	return nil
}

//the audio player
func printPlayer(g *gocui.Gui) error {
	maxX, maxY := g.Size()
	v, err := g.SetView("player", -1, maxY-2, maxX+1, maxY)
	if err != nil {
		if err != gocui.ErrUnknownView {
			return err
		}
		err = nil
	}
	setProperties(v)
	fmt.Fprintf(v, "Play Something: [%s]", strings.Repeat("=", maxX-18))
	return nil
}

//cursor movement functions, should consolodate
//TODO reduce number
func cursorUpList(g *gocui.Gui, v *gocui.View) error {
	if v != nil {
		x, y := v.Cursor()
		//if Y is 1 at the top, so don't move up again
		if y == 1 {
			return nil
		}
		yCursorOffset--
		if err := v.SetCursor(x, y-1); err != nil {
			return err
		}
	}
	return nil
}

func cursorDownList(g *gocui.Gui, v *gocui.View) error {
	if v != nil {
		x, y := v.Cursor()
		//if y is equal to number subscribed+1 is at bottom
		if y == len(globals.Config.Subscribed) {
			return nil
		}
		yCursorOffset++
		if err := v.SetCursor(x, y+1); err != nil {
			return err
		}
	}
	return nil
}

//the second view a list podcast

func switchListPodcast(g *gocui.Gui, v *gocui.View) error {
	_, position := v.Cursor()                                  //get cursor position to select
	yCursorOffset = 0                                          //reset cursor
	selectedPodcast = &(globals.Config.Subscribed[position-1]) //minus 1 is needed
	cachedPodcast = nil                                        //now delete the cache from the last time
	g.SetLayout(listPodcast)
	return nil
}

func listPodcast(g *gocui.Gui) error {
	maxX, maxY := g.Size()
	//first 5 rows reserved for description
	v, err := g.SetView("podcast", -1, 5, maxX+1, maxY-1)
	if err != nil {
		if err != gocui.ErrUnknownView { //if not created yet cool we make it
			return err
		}
	}
	d, err := g.SetView("podcastDescription", -1, -1, maxX+1, 5)
	if err != nil {
		if err != gocui.ErrUnknownView { //if not created yet cool we make it
			return err
		}
	}
	//set current view to podcast
	if err := g.SetCurrentView("podcast"); err != nil {
		return err
	}
	if err := v.SetCursor(0, 1+yCursorOffset); err != nil {
		return err
	}
	//first print the podcast description
	err = printPodcastDescription(d)
	//first print the list
	err = printList(v)
	//now print the player
	err = printPlayer(g)
	return err
}

//this function will print the podcast information when it goes to a podcast
func printPodcastDescription(v *gocui.View) error {
	setProperties(v)
	//now actually print
	return nil
}
func printList(v *gocui.View) error {
	setProperties(v)
	//now actually print
	fmt.Fprintf(v, "BLEH! \n %v", *selectedPodcast)
	return nil
}

func setProperties(v *gocui.View) {
	//First clear
	v.Clear()
	//set properties
	v.BgColor = gocui.ColorWhite
	v.FgColor = gocui.ColorBlack
	v.Wrap = false
	v.Frame = false
}
func switchToDownload(g *gocui.Gui, v *gocui.View) error {
	return nil
}

func playSelected(g *gocui.Gui, v *gocui.View) error {
	return nil
}
func quitGui(g *gocui.Gui, v *gocui.View) error {
	return gocui.ErrQuit
}

func cursorUpPodcast(g *gocui.Gui, v *gocui.View) error {
	if v != nil {
		x, y := v.Cursor()
		//if Y is 1 at the top, so don't move up again
		//TODO fix celing
		if y == 1 { //y=1 is the top because of the row that describes the colums
			return nil
		}
		yCursorOffset--
		if err := v.SetCursor(x, y-1); err != nil {
			return err
		}
	}
	return nil
}
func cursorDownPodcast(g *gocui.Gui, v *gocui.View) error {
	if v != nil {
		x, y := v.Cursor()
		if y == len(cachedPodcast) {
			return nil
		}
		yCursorOffset++
		if err := v.SetCursor(x, y+1); err != nil {
			return err
		}
	}
	return nil
}
<commit_msg>podcast episode list now kind of works, need to fix scrolling<commit_after>//this holds the tui functions and information
package main

import (
	"fmt"
	"github.com/jroimartin/gocui"
	"strings"
)

var (
	yCursorOffset   int      = 0
	selectedPodcast *Podcast = nil
	cachedPodcast   []PodcastEntry
)

func listLayout(g *gocui.Gui) error {
	maxX, maxY := g.Size()
	g.Cursor = true
	v, err := g.SetView("list", -1, -1, maxX+1, maxY-2)
	if err != nil {
		if err != gocui.ErrUnknownView {
			return err
		}
	}
	//now print subscribed
	err = printSubscribed(v)
	if err != nil {
		return err
	}
	//now print the player
	err = printPlayer(g)
	//now set current view to main view
	if err := g.SetCurrentView("list"); err != nil {
		return err
	}
	if err := v.SetCursor(0, 1+yCursorOffset); err != nil {
		return err
	}
	return err
}

func printSubscribed(v *gocui.View) error {
	//first clear
	v.Clear()
	//then set properties
	//TODO add in colors from settings
	setProperties(v)
	xMax, _ := v.Size()
	spacing := (xMax - 34) / 3 //43 chracters
	space := strings.Repeat("-", spacing)
	fmt.Fprintf(v, "Podcast Name %s Artist %s Description %s\n", space, space, space)
	for _, item := range globals.Config.Subscribed {
		strin := item.CollectionName + " - " + item.ArtistName + " - " + item.Description
		if len(item.Description+item.CollectionName+item.ArtistName)+6 < xMax {
			//do nothing
		} else { //else truncate string
			strin = strin[0:xMax]
		}
		fmt.Fprintf(v, "%s\n", strin)
	}
	return nil
}

//the audio player
func printPlayer(g *gocui.Gui) error {
	maxX, maxY := g.Size()
	v, err := g.SetView("player", -1, maxY-2, maxX+1, maxY)
	if err != nil {
		if err != gocui.ErrUnknownView {
			return err
		}
		err = nil
	}
	setProperties(v)
	fmt.Fprintf(v, "Play Something: [%s]", strings.Repeat("=", maxX-18))
	return nil
}

//cursor movement functions, should consolodate
//TODO reduce number
func cursorUpList(g *gocui.Gui, v *gocui.View) error {
	if v != nil {
		x, y := v.Cursor()
		//if Y is 1 at the top, so don't move up again
		if y == 1 {
			return nil
		}
		yCursorOffset--
		if err := v.SetCursor(x, y-1); err != nil {
			return err
		}
	}
	return nil
}

func cursorDownList(g *gocui.Gui, v *gocui.View) error {
	if v != nil {
		x, y := v.Cursor()
		//if y is equal to number subscribed+1 is at bottom
		if y == len(globals.Config.Subscribed) {
			return nil
		}
		yCursorOffset++
		if err := v.SetCursor(x, y+1); err != nil {
			return err
		}
	}
	return nil
}

//the second view a list podcast

func switchListPodcast(g *gocui.Gui, v *gocui.View) error {
	_, position := v.Cursor()                                  //get cursor position to select
	yCursorOffset = 0                                          //reset cursor
	selectedPodcast = &(globals.Config.Subscribed[position-1]) //minus 1 is needed
	cachedPodcast = nil                                        //now delete the cache from the last time
	g.SetLayout(listPodcast)
	return nil
}

func listPodcast(g *gocui.Gui) error {
	maxX, maxY := g.Size()
	//first 5 rows reserved for description
	v, err := g.SetView("podcast", -1, 5, maxX+1, maxY-1)
	if err != nil {
		if err != gocui.ErrUnknownView { //if not created yet cool we make it
			return err
		}
	}
	d, err := g.SetView("podcastDescription", -1, -1, maxX+1, 5)
	if err != nil {
		if err != gocui.ErrUnknownView { //if not created yet cool we make it
			return err
		}
	}
	//set current view to podcast
	if err := g.SetCurrentView("podcast"); err != nil {
		return err
	}
	if err := v.SetCursor(0, 0+yCursorOffset); err != nil {
		return err
	}
	//first print the podcast description
	err = printPodcastDescription(d)
	//first print the list
	err = printList(v)
	//now print the player
	err = printPlayer(g)
	return err
}

//this function will print the podcast information when it goes to a podcast
func printPodcastDescription(v *gocui.View) error {
	setProperties(v)
	//now actually print
	fmt.Fprintf(v, "BLEH! \n %v", *selectedPodcast)
	return nil
}
func printList(v *gocui.View) error {
	setProperties(v)
	var err error = nil
	//if nil then cache them
	if cachedPodcast == nil {
		cachedPodcast, err = parseRss(selectedPodcast.FeedURL)
	}
	if err != nil {
		fmt.Fprintln(v, "Cannot download podcast list, check your connection")
		return nil
	}
	//now actually print
	for i, thing := range cachedPodcast {
		fmt.Fprintf(v, "%d %s - %s\n", i+1, thing.Title, thing.Content)
	}
	return nil
}

func setProperties(v *gocui.View) {
	//First clear
	v.Clear()
	//set properties
	v.BgColor = gocui.ColorWhite
	v.FgColor = gocui.ColorBlack
	v.Wrap = false
	v.Frame = false
}
func switchToDownload(g *gocui.Gui, v *gocui.View) error {
	return nil
}

func playSelected(g *gocui.Gui, v *gocui.View) error {
	return nil
}
func quitGui(g *gocui.Gui, v *gocui.View) error {
	return gocui.ErrQuit
}

func cursorUpPodcast(g *gocui.Gui, v *gocui.View) error {
	if v != nil {
		x, y := v.Cursor()
		//if Y is 1 at the top, so don't move up again
		//TODO fix celing
		if y == 0 {
			return nil
		}
		yCursorOffset--
		if err := v.SetCursor(x, y-1); err != nil {
			return err
		}
	}
	return nil
}
func cursorDownPodcast(g *gocui.Gui, v *gocui.View) error {
	if v != nil {
		x, y := v.Cursor()
		if y == len(cachedPodcast)-1 {
			return nil
		}
		yCursorOffset++
		if err := v.SetCursor(x, y+1); err != nil {
			return err
		}
	}
	return nil
}
<|endoftext|>


<commit_before>package puddle_test

import (
	"context"
	"errors"
	"fmt"
	"sync"
	"testing"
	"time"

	"github.com/jackc/puddle"
	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/require"
)

type Counter struct {
	mutex sync.Mutex
	n     int
}

// Next increments the counter and returns the value
func (c *Counter) Next() int {
	c.mutex.Lock()
	c.n += 1
	n := c.n
	c.mutex.Unlock()
	return n
}

// Value returns the counter
func (c *Counter) Value() int {
	c.mutex.Lock()
	n := c.n
	c.mutex.Unlock()
	return n
}

func createCreateResourceFunc() (puddle.CreateFunc, *Counter) {
	var c Counter
	f := func(ctx context.Context) (interface{}, error) {
		return c.Next(), nil
	}
	return f, &c
}

func createCreateResourceFuncWithNotifierChan() (puddle.CreateFunc, *Counter, chan int) {
	ch := make(chan int)
	var c Counter
	f := func(ctx context.Context) (interface{}, error) {
		n := c.Next()

		// Because the tests will not read from ch until after the create function f returns.
		go func() { ch <- n }()

		return n, nil
	}
	return f, &c, ch
}

func createCloseResourceFuncWithNotifierChan() (puddle.CloseFunc, *Counter, chan int) {
	ch := make(chan int)
	var c Counter
	f := func(interface{}) {
		n := c.Next()

		// Because the tests will not read from ch until after the close function f returns.
		go func() { ch <- n }()
	}
	return f, &c, ch
}

func stubCloseRes(interface{}) {}

func waitForRead(ch chan int) bool {
	select {
	case <-ch:
		return true
	case <-time.NewTimer(time.Second).C:
		return false
	}
}

func TestPoolAcquireCreatesResourceWhenNoneAvailable(t *testing.T) {
	createFunc, _ := createCreateResourceFunc()
	pool := puddle.NewPool(createFunc, stubCloseRes)
	defer pool.Close()

	res, err := pool.Acquire(context.Background())
	require.NoError(t, err)
	assert.Equal(t, 1, res.Value())
	res.Release()
}

func TestPoolAcquireDoesNotCreatesResourceWhenItWouldExceedMaxSize(t *testing.T) {
	createFunc, createCounter := createCreateResourceFunc()
	pool := puddle.NewPool(createFunc, stubCloseRes)
	pool.SetMaxSize(1)

	wg := &sync.WaitGroup{}

	for i := 0; i < 100; i++ {
		wg.Add(1)
		go func() {
			for j := 0; j < 100; j++ {
				res, err := pool.Acquire(context.Background())
				assert.NoError(t, err)
				assert.Equal(t, 1, res.Value())
				res.Release()
			}
			wg.Done()
		}()
	}

	wg.Wait()

	assert.Equal(t, 1, createCounter.Value())
	assert.Equal(t, 1, pool.Size())
}

func TestPoolAcquireReturnsErrorFromFailedResourceCreate(t *testing.T) {
	errCreateFailed := errors.New("create failed")
	createFunc := func(ctx context.Context) (interface{}, error) {
		return nil, errCreateFailed
	}
	pool := puddle.NewPool(createFunc, stubCloseRes)

	res, err := pool.Acquire(context.Background())
	assert.Equal(t, errCreateFailed, err)
	assert.Nil(t, res)
}

func TestPoolAcquireReusesResources(t *testing.T) {
	createFunc, createCounter := createCreateResourceFunc()
	pool := puddle.NewPool(createFunc, stubCloseRes)

	res, err := pool.Acquire(context.Background())
	require.NoError(t, err)
	assert.Equal(t, 1, res.Value())

	res.Release()

	res, err = pool.Acquire(context.Background())
	require.NoError(t, err)
	assert.Equal(t, 1, res.Value())

	res.Release()

	assert.Equal(t, 1, createCounter.Value())
}

func TestPoolAcquireContextAlreadyCanceled(t *testing.T) {
	createFunc := func(ctx context.Context) (interface{}, error) {
		panic("should never be called")
	}
	pool := puddle.NewPool(createFunc, stubCloseRes)

	ctx, cancel := context.WithCancel(context.Background())
	cancel()
	res, err := pool.Acquire(ctx)
	assert.Equal(t, context.Canceled, err)
	assert.Nil(t, res)
}

func TestPoolAcquireContextCanceledDuringCreate(t *testing.T) {
	ctx, cancel := context.WithCancel(context.Background())
	time.AfterFunc(100*time.Millisecond, cancel)
	timeoutChan := time.After(1 * time.Second)

	var createCalls Counter
	createFunc := func(ctx context.Context) (interface{}, error) {
		select {
		case <-ctx.Done():
			return nil, ctx.Err()
		case <-timeoutChan:
		}
		return createCalls.Next(), nil
	}
	pool := puddle.NewPool(createFunc, stubCloseRes)

	res, err := pool.Acquire(ctx)
	assert.Equal(t, context.Canceled, err)
	assert.Nil(t, res)
}

func TestResourceReleaseClosesAndRemovesResourceIfOlderThanMaxDuration(t *testing.T) {
	createFunc, _ := createCreateResourceFunc()
	closeFunc, closeCalls, closeCallsChan := createCloseResourceFuncWithNotifierChan()

	pool := puddle.NewPool(createFunc, closeFunc)

	res, err := pool.Acquire(context.Background())
	require.NoError(t, err)

	assert.Equal(t, 1, pool.Size())

	pool.SetMaxResourceDuration(time.Nanosecond)
	time.Sleep(2 * time.Nanosecond)
	res.Release()

	waitForRead(closeCallsChan)
	assert.Equal(t, 0, pool.Size())
	assert.Equal(t, 1, closeCalls.Value())
}

func TestResourceReleaseClosesAndRemovesResourceWhenResourceCheckoutCountIsMaxResourceCheckouts(t *testing.T) {
	createFunc, _ := createCreateResourceFunc()
	closeFunc, closeCalls, closeCallsChan := createCloseResourceFuncWithNotifierChan()

	pool := puddle.NewPool(createFunc, closeFunc)
	pool.SetMaxResourceCheckouts(1)

	res, err := pool.Acquire(context.Background())
	require.NoError(t, err)

	res.Release()

	waitForRead(closeCallsChan)

	assert.Equal(t, 1, closeCalls.Value())
	assert.Equal(t, 0, pool.Size())
}

func TestPoolCloseClosesAllAvailableResources(t *testing.T) {
	createFunc, _ := createCreateResourceFunc()

	var closeCalls Counter
	closeFunc := func(interface{}) {
		closeCalls.Next()
	}

	p := puddle.NewPool(createFunc, closeFunc)

	resources := make([]*puddle.Resource, 4)
	for i := range resources {
		var err error
		resources[i], err = p.Acquire(context.Background())
		require.Nil(t, err)
	}

	for _, res := range resources {
		res.Release()
	}

	p.Close()

	assert.Equal(t, len(resources), closeCalls.Value())
}

func TestPoolReleaseClosesResourcePoolIsAlreadyClosed(t *testing.T) {
	createFunc, _ := createCreateResourceFunc()
	closeFunc, closeCalls, closeCallsChan := createCloseResourceFuncWithNotifierChan()

	p := puddle.NewPool(createFunc, closeFunc)

	resources := make([]*puddle.Resource, 4)
	for i := range resources {
		var err error
		resources[i], err = p.Acquire(context.Background())
		require.Nil(t, err)
	}

	p.Close()
	assert.Equal(t, 0, closeCalls.Value())

	for _, res := range resources {
		res.Release()
	}

	waitForRead(closeCallsChan)
	waitForRead(closeCallsChan)
	waitForRead(closeCallsChan)
	waitForRead(closeCallsChan)

	assert.Equal(t, len(resources), closeCalls.Value())
}

func TestResourceDestroyRemovesResourceFromPool(t *testing.T) {
	createFunc, _ := createCreateResourceFunc()
	pool := puddle.NewPool(createFunc, stubCloseRes)

	res, err := pool.Acquire(context.Background())
	require.NoError(t, err)
	assert.Equal(t, 1, res.Value())

	assert.Equal(t, 1, pool.Size())
	res.Destroy()
	assert.Equal(t, 0, pool.Size())
}

func TestPoolAcquireReturnsErrorWhenPoolIsClosed(t *testing.T) {
	createFunc, _ := createCreateResourceFunc()
	pool := puddle.NewPool(createFunc, stubCloseRes)
	pool.Close()

	res, err := pool.Acquire(context.Background())
	assert.Equal(t, puddle.ErrClosedPool, err)
	assert.Nil(t, res)
}

func BenchmarkPoolAcquireAndRelease(b *testing.B) {
	benchmarks := []struct {
		poolSize              int
		concurrentClientCount int
		loanDuration          time.Duration
	}{
		// Small pool
		{10, 1, 0},
		{10, 5, 0},
		{10, 10, 0},
		{10, 20, 0},
		{10, 1, 1 * time.Millisecond},
		{10, 5, 1 * time.Millisecond},
		{10, 10, 1 * time.Millisecond},
		{10, 20, 1 * time.Millisecond},

		// large pool
		{100, 1, 0},
		{100, 50, 0},
		{100, 100, 0},
		{100, 200, 0},
		{100, 1, 1 * time.Millisecond},
		{100, 50, 1 * time.Millisecond},
		{100, 100, 1 * time.Millisecond},
		{100, 200, 1 * time.Millisecond},

		// huge pool
		{1000, 1, 0},
		{1000, 500, 0},
		{1000, 1000, 0},
		{1000, 2000, 0},
		{1000, 1, 1 * time.Millisecond},
		{1000, 500, 1 * time.Millisecond},
		{1000, 1000, 1 * time.Millisecond},
		{1000, 2000, 1 * time.Millisecond},
	}

	for _, bm := range benchmarks {
		name := fmt.Sprintf("PoolSize=%d/ConcurrentClientCount=%d/LoanDuration=%v", bm.poolSize, bm.concurrentClientCount, bm.loanDuration)

		createFunc, _ := createCreateResourceFunc()
		pool := puddle.NewPool(createFunc, stubCloseRes)
		pool.SetMaxSize(bm.poolSize)

		acquireAndRelease := func() {
			res, err := pool.Acquire(context.Background())
			if err != nil {
				b.Fatal(err)
			}
			time.Sleep(bm.loanDuration)
			res.Release()
		}

		b.Run(name, func(b *testing.B) {
			doneChan := make(chan struct{})
			defer close(doneChan)
			for i := 0; i < bm.concurrentClientCount-1; i++ {
				go func() {
					for {
						select {
						case <-doneChan:
							return
						default:
						}

						acquireAndRelease()
					}
				}()
			}

			for i := 0; i < b.N; i++ {
				acquireAndRelease()
			}
		})
	}
}
<commit_msg>Better benchmarking<commit_after>package puddle_test

import (
	"context"
	"errors"
	"fmt"
	"sync"
	"testing"
	"time"

	"github.com/jackc/puddle"
	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/require"
)

type Counter struct {
	mutex sync.Mutex
	n     int
}

// Next increments the counter and returns the value
func (c *Counter) Next() int {
	c.mutex.Lock()
	c.n += 1
	n := c.n
	c.mutex.Unlock()
	return n
}

// Value returns the counter
func (c *Counter) Value() int {
	c.mutex.Lock()
	n := c.n
	c.mutex.Unlock()
	return n
}

func createCreateResourceFunc() (puddle.CreateFunc, *Counter) {
	var c Counter
	f := func(ctx context.Context) (interface{}, error) {
		return c.Next(), nil
	}
	return f, &c
}

func createCreateResourceFuncWithNotifierChan() (puddle.CreateFunc, *Counter, chan int) {
	ch := make(chan int)
	var c Counter
	f := func(ctx context.Context) (interface{}, error) {
		n := c.Next()

		// Because the tests will not read from ch until after the create function f returns.
		go func() { ch <- n }()

		return n, nil
	}
	return f, &c, ch
}

func createCloseResourceFuncWithNotifierChan() (puddle.CloseFunc, *Counter, chan int) {
	ch := make(chan int)
	var c Counter
	f := func(interface{}) {
		n := c.Next()

		// Because the tests will not read from ch until after the close function f returns.
		go func() { ch <- n }()
	}
	return f, &c, ch
}

func stubCloseRes(interface{}) {}

func waitForRead(ch chan int) bool {
	select {
	case <-ch:
		return true
	case <-time.NewTimer(time.Second).C:
		return false
	}
}

func TestPoolAcquireCreatesResourceWhenNoneAvailable(t *testing.T) {
	createFunc, _ := createCreateResourceFunc()
	pool := puddle.NewPool(createFunc, stubCloseRes)
	defer pool.Close()

	res, err := pool.Acquire(context.Background())
	require.NoError(t, err)
	assert.Equal(t, 1, res.Value())
	res.Release()
}

func TestPoolAcquireDoesNotCreatesResourceWhenItWouldExceedMaxSize(t *testing.T) {
	createFunc, createCounter := createCreateResourceFunc()
	pool := puddle.NewPool(createFunc, stubCloseRes)
	pool.SetMaxSize(1)

	wg := &sync.WaitGroup{}

	for i := 0; i < 100; i++ {
		wg.Add(1)
		go func() {
			for j := 0; j < 100; j++ {
				res, err := pool.Acquire(context.Background())
				assert.NoError(t, err)
				assert.Equal(t, 1, res.Value())
				res.Release()
			}
			wg.Done()
		}()
	}

	wg.Wait()

	assert.Equal(t, 1, createCounter.Value())
	assert.Equal(t, 1, pool.Size())
}

func TestPoolAcquireReturnsErrorFromFailedResourceCreate(t *testing.T) {
	errCreateFailed := errors.New("create failed")
	createFunc := func(ctx context.Context) (interface{}, error) {
		return nil, errCreateFailed
	}
	pool := puddle.NewPool(createFunc, stubCloseRes)

	res, err := pool.Acquire(context.Background())
	assert.Equal(t, errCreateFailed, err)
	assert.Nil(t, res)
}

func TestPoolAcquireReusesResources(t *testing.T) {
	createFunc, createCounter := createCreateResourceFunc()
	pool := puddle.NewPool(createFunc, stubCloseRes)

	res, err := pool.Acquire(context.Background())
	require.NoError(t, err)
	assert.Equal(t, 1, res.Value())

	res.Release()

	res, err = pool.Acquire(context.Background())
	require.NoError(t, err)
	assert.Equal(t, 1, res.Value())

	res.Release()

	assert.Equal(t, 1, createCounter.Value())
}

func TestPoolAcquireContextAlreadyCanceled(t *testing.T) {
	createFunc := func(ctx context.Context) (interface{}, error) {
		panic("should never be called")
	}
	pool := puddle.NewPool(createFunc, stubCloseRes)

	ctx, cancel := context.WithCancel(context.Background())
	cancel()
	res, err := pool.Acquire(ctx)
	assert.Equal(t, context.Canceled, err)
	assert.Nil(t, res)
}

func TestPoolAcquireContextCanceledDuringCreate(t *testing.T) {
	ctx, cancel := context.WithCancel(context.Background())
	time.AfterFunc(100*time.Millisecond, cancel)
	timeoutChan := time.After(1 * time.Second)

	var createCalls Counter
	createFunc := func(ctx context.Context) (interface{}, error) {
		select {
		case <-ctx.Done():
			return nil, ctx.Err()
		case <-timeoutChan:
		}
		return createCalls.Next(), nil
	}
	pool := puddle.NewPool(createFunc, stubCloseRes)

	res, err := pool.Acquire(ctx)
	assert.Equal(t, context.Canceled, err)
	assert.Nil(t, res)
}

func TestResourceReleaseClosesAndRemovesResourceIfOlderThanMaxDuration(t *testing.T) {
	createFunc, _ := createCreateResourceFunc()
	closeFunc, closeCalls, closeCallsChan := createCloseResourceFuncWithNotifierChan()

	pool := puddle.NewPool(createFunc, closeFunc)

	res, err := pool.Acquire(context.Background())
	require.NoError(t, err)

	assert.Equal(t, 1, pool.Size())

	pool.SetMaxResourceDuration(time.Nanosecond)
	time.Sleep(2 * time.Nanosecond)
	res.Release()

	waitForRead(closeCallsChan)
	assert.Equal(t, 0, pool.Size())
	assert.Equal(t, 1, closeCalls.Value())
}

func TestResourceReleaseClosesAndRemovesResourceWhenResourceCheckoutCountIsMaxResourceCheckouts(t *testing.T) {
	createFunc, _ := createCreateResourceFunc()
	closeFunc, closeCalls, closeCallsChan := createCloseResourceFuncWithNotifierChan()

	pool := puddle.NewPool(createFunc, closeFunc)
	pool.SetMaxResourceCheckouts(1)

	res, err := pool.Acquire(context.Background())
	require.NoError(t, err)

	res.Release()

	waitForRead(closeCallsChan)

	assert.Equal(t, 1, closeCalls.Value())
	assert.Equal(t, 0, pool.Size())
}

func TestPoolCloseClosesAllAvailableResources(t *testing.T) {
	createFunc, _ := createCreateResourceFunc()

	var closeCalls Counter
	closeFunc := func(interface{}) {
		closeCalls.Next()
	}

	p := puddle.NewPool(createFunc, closeFunc)

	resources := make([]*puddle.Resource, 4)
	for i := range resources {
		var err error
		resources[i], err = p.Acquire(context.Background())
		require.Nil(t, err)
	}

	for _, res := range resources {
		res.Release()
	}

	p.Close()

	assert.Equal(t, len(resources), closeCalls.Value())
}

func TestPoolReleaseClosesResourcePoolIsAlreadyClosed(t *testing.T) {
	createFunc, _ := createCreateResourceFunc()
	closeFunc, closeCalls, closeCallsChan := createCloseResourceFuncWithNotifierChan()

	p := puddle.NewPool(createFunc, closeFunc)

	resources := make([]*puddle.Resource, 4)
	for i := range resources {
		var err error
		resources[i], err = p.Acquire(context.Background())
		require.Nil(t, err)
	}

	p.Close()
	assert.Equal(t, 0, closeCalls.Value())

	for _, res := range resources {
		res.Release()
	}

	waitForRead(closeCallsChan)
	waitForRead(closeCallsChan)
	waitForRead(closeCallsChan)
	waitForRead(closeCallsChan)

	assert.Equal(t, len(resources), closeCalls.Value())
}

func TestResourceDestroyRemovesResourceFromPool(t *testing.T) {
	createFunc, _ := createCreateResourceFunc()
	pool := puddle.NewPool(createFunc, stubCloseRes)

	res, err := pool.Acquire(context.Background())
	require.NoError(t, err)
	assert.Equal(t, 1, res.Value())

	assert.Equal(t, 1, pool.Size())
	res.Destroy()
	assert.Equal(t, 0, pool.Size())
}

func TestPoolAcquireReturnsErrorWhenPoolIsClosed(t *testing.T) {
	createFunc, _ := createCreateResourceFunc()
	pool := puddle.NewPool(createFunc, stubCloseRes)
	pool.Close()

	res, err := pool.Acquire(context.Background())
	assert.Equal(t, puddle.ErrClosedPool, err)
	assert.Nil(t, res)
}

func BenchmarkPoolAcquireAndRelease(b *testing.B) {
	benchmarks := []struct {
		poolSize    int
		clientCount int
	}{
		{8, 1},
		{8, 4},
		{8, 8},
		{8, 16},
		{8, 32},
		{8, 64},
		{8, 64},
		{8, 128},
		{8, 256},
		{8, 512},
		{8, 1024},
		{8, 2048},

		{64, 1},
		{64, 4},
		{64, 8},
		{64, 16},
		{64, 32},
		{64, 64},
		{64, 64},
		{64, 128},
		{64, 256},
		{64, 512},
		{64, 1024},
		{64, 2048},

		{512, 1},
		{512, 4},
		{512, 8},
		{512, 16},
		{512, 32},
		{512, 64},
		{512, 64},
		{512, 128},
		{512, 256},
		{512, 512},
		{512, 1024},
		{512, 2048},
	}

	for _, bm := range benchmarks {
		name := fmt.Sprintf("PoolSize=%d/ClientCount=%d", bm.poolSize, bm.clientCount)

		b.Run(name, func(b *testing.B) {
			wg := &sync.WaitGroup{}

			createFunc, _ := createCreateResourceFunc()
			pool := puddle.NewPool(createFunc, stubCloseRes)
			pool.SetMaxSize(bm.poolSize)

			for i := 0; i < bm.clientCount; i++ {
				wg.Add(1)
				go func() {
					defer wg.Done()

					for j := 0; j < b.N; j++ {
						res, err := pool.Acquire(context.Background())
						if err != nil {
							b.Fatal(err)
						}
						res.Release()
					}
				}()
			}

			wg.Wait()
		})
	}
}
<|endoftext|>


<commit_before>package post

import (
	"database/sql"
	"fmt"
	"log"
	"math"

	"github.com/rwcarlsen/cyan/query"
	_ "github.com/rwcarlsen/go-sqlite/sqlite3"
)

// The number of sql commands to buffer before dumping to the output database.
const DumpFreq = 100000

var (
	preExecStmts = []string{
		"CREATE TABLE IF NOT EXISTS TimeSeriesPower (SimId BLOB,AgentId INTEGER,Time INTEGER, Value REAL);",
		"CREATE TABLE IF NOT EXISTS AgentExit (SimId BLOB,AgentId INTEGER,ExitTime INTEGER);",
		"CREATE TABLE IF NOT EXISTS Compositions (SimId BLOB,QualId INTEGER,NucId INTEGER, MassFrac REAL);",
		"CREATE TABLE IF NOT EXISTS Products (SimId BLOB,QualId INTEGER,Quality TEXT);",
		"CREATE TABLE IF NOT EXISTS Resources (SimId INTEGER,ResourceId INTEGER,ObjId INTEGER,Type TEXT,TimeCreated INTEGER,Quantity REAL,Units TEXT,QualId INTEGER,Parent1 INTEGER,Parent2 INTEGER);",
		"CREATE TABLE IF NOT EXISTS ResCreators (SimId INTEGER,ResourceId INTEGER,AgentId INTEGER);",
		"CREATE TABLE IF NOT EXISTS Agents (SimId BLOB,AgentId INTEGER,Kind TEXT,Spec TEXT,Prototype TEXT,ParentId INTEGER,Lifetime INTEGER,EnterTime INTEGER,ExitTime INTEGER);",
		"CREATE TABLE IF NOT EXISTS Inventories (SimId BLOB,ResourceId INTEGER,AgentId INTEGER,StartTime INTEGER,EndTime INTEGER,QualId INTEGER,Quantity REAL);",
		"CREATE TABLE IF NOT EXISTS TimeList (SimId BLOB, Time INTEGER);",
		"CREATE TABLE IF NOT EXISTS Transactions (SimId BLOB, TransactionId INTEGER, SenderId INTEGER, ReceiverId INTEGER, ResourceId INTEGER, Commodity TEXT, Time INTEGER);",
		query.Index("TimeSeriesPower", "SimId", "AgentId", "Time", "Value"),
		query.Index("TimeList", "Time"),
		query.Index("TimeList", "SimId", "Time"),
		query.Index("Resources", "SimId", "ResourceId", "QualId"),
		query.Index("Compositions", "SimId", "QualId", "NucId"),
		query.Index("Transactions", "SimId", "ResourceId"),
		query.Index("Transactions", "TransactionId"),
		query.Index("ResCreators", "SimId", "ResourceId"),
	}
	postExecStmts = []string{
		query.Index("Agents", "SimId", "Prototype"),
		query.Index("Agents", "SimId", "AgentId", "Prototype"),
		query.Index("Inventories", "SimId", "AgentId", "StartTime", "EndTime", "Quantity"),
		query.Index("Inventories", "SimId", "ResourceId", "StartTime"),
		query.Index("Inventories", "SimId", "StartTime", "EndTime", "ResourceId", "Quantity"),
		"ANALYZE;",
	}
	dumpSql    = "INSERT INTO Inventories VALUES (?,?,?,?,?,?,?);"
	resSqlHead = "SELECT ResourceId,TimeCreated,QualId,Quantity FROM "
	resSqlTail = " WHERE Parent1 = ? OR Parent2 = ?;"

	ownerSql = `SELECT tr.ReceiverId, tr.Time FROM Transactions AS tr
				  WHERE tr.ResourceId = ? AND tr.SimId = ?
				  ORDER BY tr.Time ASC;`
	rootsSql = `SELECT res.ResourceId,res.TimeCreated,rc.AgentId,res.QualId,Quantity FROM Resources AS res
				  INNER JOIN ResCreators AS rc ON res.ResourceId = rc.ResourceId
				  WHERE res.SimId = ? AND rc.SimId = ?;`
)

func Process(db *sql.DB) (simids [][]byte, err error) {
	err = Prepare(db)
	if err != nil {
		return nil, err
	}

	simids, err = GetSimIds(db)
	if err != nil {
		return nil, err
	}

	nprocessed := 0
	for _, id := range simids {
		ctx := NewContext(db, id)
		if err2 := ctx.WalkAll(); err2 != nil {
			if IsAlreadyPostErr(err2) {
			} else {
				err = err2
			}
		} else {
			nprocessed++
		}
	}
	if nprocessed > 0 {
		Finish(db)
	}
	return simids, nil
}

// Prepare creates necessary indexes and tables required for efficient
// calculation of cyclus simulation inventory information.  Should be called
// once before walking begins.
func Prepare(db *sql.DB) (err error) {
	for _, s := range preExecStmts {
		if _, err := db.Exec(s); err != nil {
			log.Println("    ", err)
		}
	}
	return nil
}

// Finish should be called for a cyclus database after all walkers have
// completed processing inventory data. It creates final indexes and other
// finishing tasks.
func Finish(db *sql.DB) (err error) {
	for _, s := range postExecStmts {
		if _, err := db.Exec(s); err != nil {
			return err
		}
	}
	return nil
}

type Node struct {
	ResId     int
	OwnerId   int
	StartTime int
	EndTime   int
	QualId    int
	Quantity  float64
}

// Context encapsulates the logic for building a fast, queryable inventories
// table for a specific simulation from raw cyclus output database.
type Context struct {
	*sql.DB
	// Simid is the cyclus simulation id targeted by this context.  Must be
	// set.
	Simid       []byte
	Log         *log.Logger
	mappednodes map[int32]struct{}
	tmpResTbl   string
	tmpResStmt  *sql.Stmt
	dumpStmt    *sql.Stmt
	ownerStmt   *sql.Stmt
	resCount    int
	nodes       []*Node
}

func NewContext(db *sql.DB, simid []byte) *Context {
	return &Context{
		DB:    db,
		Simid: simid,
		Log:   log.New(NullWriter{}, "", 0),
	}
}

type AlreadyPostErr []byte

func (s AlreadyPostErr) Error() string {
	return fmt.Sprintf("SimId %x is already post processed", []byte(s))
}

func IsAlreadyPostErr(err error) bool {
	_, ok := err.(AlreadyPostErr)
	return ok
}

func (c *Context) init() {
	// skip if the post processing already exists for this simid in the db
	dummy := 0
	err := c.QueryRow("SELECT AgentId FROM Agents WHERE SimId = ? LIMIT 1", c.Simid).Scan(&dummy)
	if err == nil {
		panic(AlreadyPostErr(c.Simid))
	} else if err != sql.ErrNoRows {
		panicif(err)
	}

	tx, err := c.Begin()
	panicif(err)

	// build Agents table
	sql := `INSERT INTO Agents
				SELECT n.SimId,n.AgentId,n.Kind,n.Spec,n.Prototype,n.ParentId,n.Lifetime,n.EnterTime,x.ExitTime
				FROM
					AgentEntry AS n
					LEFT JOIN AgentExit AS x ON n.AgentId = x.AgentId AND n.SimId = x.SimId
					WHERE n.SimId = ?;`
	_, err = tx.Exec(sql, c.Simid)
	panicif(err)

	c.nodes = make([]*Node, 0, 10000)
	c.mappednodes = map[int32]struct{}{}

	// build TimeList table
	sql = "SELECT Duration FROM Info WHERE SimId = ?;"
	rows, err := tx.Query(sql, c.Simid)
	panicif(err)
	defer rows.Close()
	for rows.Next() {
		var dur int
		panicif(rows.Scan(&dur))
		for i := 0; i < dur; i++ {
			_, err := tx.Exec("INSERT INTO TimeList VALUES (?, ?);", c.Simid, i)
			panicif(err)
		}
	}
	panicif(rows.Err())

	// create temp res table without simid
	c.Log.Println("Creating temporary resource table...")
	c.tmpResTbl = "tmp_restbl_" + fmt.Sprintf("%x", c.Simid)
	_, err = tx.Exec("DROP TABLE IF EXISTS " + c.tmpResTbl)
	panicif(err)

	sql = "CREATE TABLE " + c.tmpResTbl + " AS SELECT ResourceId,TimeCreated,Parent1,Parent2,QualId,Quantity FROM Resources WHERE SimId = ?;"
	_, err = tx.Exec(sql, c.Simid)
	panicif(err)

	c.Log.Println("Indexing temporary resource table...")
	_, err = tx.Exec(query.Index(c.tmpResTbl, "Parent1"))
	panicif(err)

	_, err = tx.Exec(query.Index(c.tmpResTbl, "Parent2"))
	panicif(err)

	tx.Commit()

	// create prepared statements
	c.tmpResStmt, err = c.Prepare(resSqlHead + c.tmpResTbl + resSqlTail)
	panicif(err)

	c.dumpStmt, err = c.Prepare(dumpSql)
	panicif(err)

	c.ownerStmt, err = c.Prepare(ownerSql)
	panicif(err)
}

// WalkAll constructs the inventories table in the cyclus database alongside
// other tables. Creates several indexes in the process.  Finish should be
// called on the database connection after all simulation id's have been
// walked.
func (c *Context) WalkAll() (err error) {
	defer func() {
		if r := recover(); r != nil {
			if er, ok := r.(error); ok {
				err = er
			} else {
				err = fmt.Errorf("%v", r)
			}
		}
	}()

	c.Log.Printf("--- Building inventories for simid %x ---\n", c.Simid)
	c.init()

	c.Log.Println("Retrieving root resource nodes...")
	roots := c.getRoots()

	c.Log.Printf("Found %v root nodes\n", len(roots))
	for i, n := range roots {
		c.Log.Printf("    Processing root %d...\n", i)
		c.walkDown(n)
	}

	c.Log.Println("Dropping temporary resource table...")
	_, err = c.Exec("DROP TABLE " + c.tmpResTbl)
	panicif(err)

	c.dumpNodes()

	return nil
}

func (c *Context) getRoots() (roots []*Node) {
	sql := "SELECT COUNT(*) FROM ResCreators WHERE SimId = ?"
	row := c.QueryRow(sql, c.Simid)

	n := 0
	err := row.Scan(&n)
	panicif(err)

	roots = make([]*Node, 0, n)
	rows, err := c.Query(rootsSql, c.Simid, c.Simid)
	panicif(err)
	defer rows.Close()
	for rows.Next() {
		node := &Node{EndTime: math.MaxInt32}
		err := rows.Scan(&node.ResId, &node.StartTime, &node.OwnerId, &node.QualId, &node.Quantity)
		panicif(err)

		roots = append(roots, node)
	}
	panicif(rows.Err())
	return roots
}

func (c *Context) walkDown(node *Node) {
	if _, ok := c.mappednodes[int32(node.ResId)]; ok {
		return
	}
	c.mappednodes[int32(node.ResId)] = struct{}{}

	// dump if necessary
	c.resCount++
	if c.resCount%DumpFreq == 0 {
		c.dumpNodes()
	}

	// find resource's children
	kids := make([]*Node, 0, 2)
	rows, err := c.tmpResStmt.Query(node.ResId, node.ResId)
	panicif(err)
	defer rows.Close()

	for rows.Next() {
		child := &Node{EndTime: math.MaxInt32}
		err := rows.Scan(&child.ResId, &child.StartTime, &child.QualId, &child.Quantity)
		panicif(err)
		node.EndTime = child.StartTime
		kids = append(kids, child)
	}
	panicif(rows.Err())
	rows.Close() // Close is idempotent - and this function is recursive.

	// find resources owner changes (that occurred before children)
	owners, times := c.getNewOwners(node.ResId)

	childOwner := node.OwnerId
	if len(owners) > 0 {
		node.EndTime = times[0]
		childOwner = owners[len(owners)-1]

		lastend := math.MaxInt32
		if len(kids) > 0 {
			lastend = kids[0].StartTime
		}
		times = append(times, lastend)
		for i := range owners {
			n := &Node{ResId: node.ResId,
				OwnerId:   owners[i],
				StartTime: times[i],
				EndTime:   times[i+1],
				QualId:    node.QualId,
				Quantity:  node.Quantity,
			}
			c.nodes = append(c.nodes, n)
		}
	}

	c.nodes = append(c.nodes, node)

	// walk down resource's children
	for _, child := range kids {
		child.OwnerId = childOwner
		c.walkDown(child)
	}
}

func (c *Context) getNewOwners(id int) (owners, times []int) {
	var owner, t int
	rows, err := c.ownerStmt.Query(id, c.Simid)
	panicif(err)
	defer rows.Close()
	for rows.Next() {
		err := rows.Scan(&owner, &t)
		panicif(err)

		if id == owner {
			continue
		}
		owners = append(owners, owner)
		times = append(times, t)
	}
	panicif(rows.Err())
	return owners, times
}

func (c *Context) dumpNodes() {
	c.Log.Printf("    Dumping inventories (%d resources done)...\n", c.resCount)
	tx, err := c.Begin()
	panicif(err)
	stmt := tx.Stmt(c.dumpStmt)

	for _, n := range c.nodes {
		if n.EndTime > n.StartTime {
			_, err = stmt.Exec(c.Simid, n.ResId, n.OwnerId, n.StartTime, n.EndTime, n.QualId, n.Quantity)
			panicif(err)
		}
	}

	err = tx.Commit()
	panicif(err)
	c.nodes = c.nodes[:0]
}
<commit_msg>remove sqlite import from post pkg<commit_after>package post

import (
	"database/sql"
	"fmt"
	"log"
	"math"

	"github.com/rwcarlsen/cyan/query"
)

// The number of sql commands to buffer before dumping to the output database.
const DumpFreq = 100000

var (
	preExecStmts = []string{
		"CREATE TABLE IF NOT EXISTS TimeSeriesPower (SimId BLOB,AgentId INTEGER,Time INTEGER, Value REAL);",
		"CREATE TABLE IF NOT EXISTS AgentExit (SimId BLOB,AgentId INTEGER,ExitTime INTEGER);",
		"CREATE TABLE IF NOT EXISTS Compositions (SimId BLOB,QualId INTEGER,NucId INTEGER, MassFrac REAL);",
		"CREATE TABLE IF NOT EXISTS Products (SimId BLOB,QualId INTEGER,Quality TEXT);",
		"CREATE TABLE IF NOT EXISTS Resources (SimId INTEGER,ResourceId INTEGER,ObjId INTEGER,Type TEXT,TimeCreated INTEGER,Quantity REAL,Units TEXT,QualId INTEGER,Parent1 INTEGER,Parent2 INTEGER);",
		"CREATE TABLE IF NOT EXISTS ResCreators (SimId INTEGER,ResourceId INTEGER,AgentId INTEGER);",
		"CREATE TABLE IF NOT EXISTS Agents (SimId BLOB,AgentId INTEGER,Kind TEXT,Spec TEXT,Prototype TEXT,ParentId INTEGER,Lifetime INTEGER,EnterTime INTEGER,ExitTime INTEGER);",
		"CREATE TABLE IF NOT EXISTS Inventories (SimId BLOB,ResourceId INTEGER,AgentId INTEGER,StartTime INTEGER,EndTime INTEGER,QualId INTEGER,Quantity REAL);",
		"CREATE TABLE IF NOT EXISTS TimeList (SimId BLOB, Time INTEGER);",
		"CREATE TABLE IF NOT EXISTS Transactions (SimId BLOB, TransactionId INTEGER, SenderId INTEGER, ReceiverId INTEGER, ResourceId INTEGER, Commodity TEXT, Time INTEGER);",
		query.Index("TimeSeriesPower", "SimId", "AgentId", "Time", "Value"),
		query.Index("TimeList", "Time"),
		query.Index("TimeList", "SimId", "Time"),
		query.Index("Resources", "SimId", "ResourceId", "QualId"),
		query.Index("Compositions", "SimId", "QualId", "NucId"),
		query.Index("Transactions", "SimId", "ResourceId"),
		query.Index("Transactions", "TransactionId"),
		query.Index("ResCreators", "SimId", "ResourceId"),
	}
	postExecStmts = []string{
		query.Index("Agents", "SimId", "Prototype"),
		query.Index("Agents", "SimId", "AgentId", "Prototype"),
		query.Index("Inventories", "SimId", "AgentId", "StartTime", "EndTime", "Quantity"),
		query.Index("Inventories", "SimId", "ResourceId", "StartTime"),
		query.Index("Inventories", "SimId", "StartTime", "EndTime", "ResourceId", "Quantity"),
		"ANALYZE;",
	}
	dumpSql    = "INSERT INTO Inventories VALUES (?,?,?,?,?,?,?);"
	resSqlHead = "SELECT ResourceId,TimeCreated,QualId,Quantity FROM "
	resSqlTail = " WHERE Parent1 = ? OR Parent2 = ?;"

	ownerSql = `SELECT tr.ReceiverId, tr.Time FROM Transactions AS tr
				  WHERE tr.ResourceId = ? AND tr.SimId = ?
				  ORDER BY tr.Time ASC;`
	rootsSql = `SELECT res.ResourceId,res.TimeCreated,rc.AgentId,res.QualId,Quantity FROM Resources AS res
				  INNER JOIN ResCreators AS rc ON res.ResourceId = rc.ResourceId
				  WHERE res.SimId = ? AND rc.SimId = ?;`
)

func Process(db *sql.DB) (simids [][]byte, err error) {
	err = Prepare(db)
	if err != nil {
		return nil, err
	}

	simids, err = GetSimIds(db)
	if err != nil {
		return nil, err
	}

	nprocessed := 0
	for _, id := range simids {
		ctx := NewContext(db, id)
		if err2 := ctx.WalkAll(); err2 != nil {
			if IsAlreadyPostErr(err2) {
			} else {
				err = err2
			}
		} else {
			nprocessed++
		}
	}
	if nprocessed > 0 {
		Finish(db)
	}
	return simids, nil
}

// Prepare creates necessary indexes and tables required for efficient
// calculation of cyclus simulation inventory information.  Should be called
// once before walking begins.
func Prepare(db *sql.DB) (err error) {
	for _, s := range preExecStmts {
		if _, err := db.Exec(s); err != nil {
			log.Println("    ", err)
		}
	}
	return nil
}

// Finish should be called for a cyclus database after all walkers have
// completed processing inventory data. It creates final indexes and other
// finishing tasks.
func Finish(db *sql.DB) (err error) {
	for _, s := range postExecStmts {
		if _, err := db.Exec(s); err != nil {
			return err
		}
	}
	return nil
}

type Node struct {
	ResId     int
	OwnerId   int
	StartTime int
	EndTime   int
	QualId    int
	Quantity  float64
}

// Context encapsulates the logic for building a fast, queryable inventories
// table for a specific simulation from raw cyclus output database.
type Context struct {
	*sql.DB
	// Simid is the cyclus simulation id targeted by this context.  Must be
	// set.
	Simid       []byte
	Log         *log.Logger
	mappednodes map[int32]struct{}
	tmpResTbl   string
	tmpResStmt  *sql.Stmt
	dumpStmt    *sql.Stmt
	ownerStmt   *sql.Stmt
	resCount    int
	nodes       []*Node
}

func NewContext(db *sql.DB, simid []byte) *Context {
	return &Context{
		DB:    db,
		Simid: simid,
		Log:   log.New(NullWriter{}, "", 0),
	}
}

type AlreadyPostErr []byte

func (s AlreadyPostErr) Error() string {
	return fmt.Sprintf("SimId %x is already post processed", []byte(s))
}

func IsAlreadyPostErr(err error) bool {
	_, ok := err.(AlreadyPostErr)
	return ok
}

func (c *Context) init() {
	// skip if the post processing already exists for this simid in the db
	dummy := 0
	err := c.QueryRow("SELECT AgentId FROM Agents WHERE SimId = ? LIMIT 1", c.Simid).Scan(&dummy)
	if err == nil {
		panic(AlreadyPostErr(c.Simid))
	} else if err != sql.ErrNoRows {
		panicif(err)
	}

	tx, err := c.Begin()
	panicif(err)

	// build Agents table
	sql := `INSERT INTO Agents
				SELECT n.SimId,n.AgentId,n.Kind,n.Spec,n.Prototype,n.ParentId,n.Lifetime,n.EnterTime,x.ExitTime
				FROM
					AgentEntry AS n
					LEFT JOIN AgentExit AS x ON n.AgentId = x.AgentId AND n.SimId = x.SimId
					WHERE n.SimId = ?;`
	_, err = tx.Exec(sql, c.Simid)
	panicif(err)

	c.nodes = make([]*Node, 0, 10000)
	c.mappednodes = map[int32]struct{}{}

	// build TimeList table
	sql = "SELECT Duration FROM Info WHERE SimId = ?;"
	rows, err := tx.Query(sql, c.Simid)
	panicif(err)
	defer rows.Close()
	for rows.Next() {
		var dur int
		panicif(rows.Scan(&dur))
		for i := 0; i < dur; i++ {
			_, err := tx.Exec("INSERT INTO TimeList VALUES (?, ?);", c.Simid, i)
			panicif(err)
		}
	}
	panicif(rows.Err())

	// create temp res table without simid
	c.Log.Println("Creating temporary resource table...")
	c.tmpResTbl = "tmp_restbl_" + fmt.Sprintf("%x", c.Simid)
	_, err = tx.Exec("DROP TABLE IF EXISTS " + c.tmpResTbl)
	panicif(err)

	sql = "CREATE TABLE " + c.tmpResTbl + " AS SELECT ResourceId,TimeCreated,Parent1,Parent2,QualId,Quantity FROM Resources WHERE SimId = ?;"
	_, err = tx.Exec(sql, c.Simid)
	panicif(err)

	c.Log.Println("Indexing temporary resource table...")
	_, err = tx.Exec(query.Index(c.tmpResTbl, "Parent1"))
	panicif(err)

	_, err = tx.Exec(query.Index(c.tmpResTbl, "Parent2"))
	panicif(err)

	tx.Commit()

	// create prepared statements
	c.tmpResStmt, err = c.Prepare(resSqlHead + c.tmpResTbl + resSqlTail)
	panicif(err)

	c.dumpStmt, err = c.Prepare(dumpSql)
	panicif(err)

	c.ownerStmt, err = c.Prepare(ownerSql)
	panicif(err)
}

// WalkAll constructs the inventories table in the cyclus database alongside
// other tables. Creates several indexes in the process.  Finish should be
// called on the database connection after all simulation id's have been
// walked.
func (c *Context) WalkAll() (err error) {
	defer func() {
		if r := recover(); r != nil {
			if er, ok := r.(error); ok {
				err = er
			} else {
				err = fmt.Errorf("%v", r)
			}
		}
	}()

	c.Log.Printf("--- Building inventories for simid %x ---\n", c.Simid)
	c.init()

	c.Log.Println("Retrieving root resource nodes...")
	roots := c.getRoots()

	c.Log.Printf("Found %v root nodes\n", len(roots))
	for i, n := range roots {
		c.Log.Printf("    Processing root %d...\n", i)
		c.walkDown(n)
	}

	c.Log.Println("Dropping temporary resource table...")
	_, err = c.Exec("DROP TABLE " + c.tmpResTbl)
	panicif(err)

	c.dumpNodes()

	return nil
}

func (c *Context) getRoots() (roots []*Node) {
	sql := "SELECT COUNT(*) FROM ResCreators WHERE SimId = ?"
	row := c.QueryRow(sql, c.Simid)

	n := 0
	err := row.Scan(&n)
	panicif(err)

	roots = make([]*Node, 0, n)
	rows, err := c.Query(rootsSql, c.Simid, c.Simid)
	panicif(err)
	defer rows.Close()
	for rows.Next() {
		node := &Node{EndTime: math.MaxInt32}
		err := rows.Scan(&node.ResId, &node.StartTime, &node.OwnerId, &node.QualId, &node.Quantity)
		panicif(err)

		roots = append(roots, node)
	}
	panicif(rows.Err())
	return roots
}

func (c *Context) walkDown(node *Node) {
	if _, ok := c.mappednodes[int32(node.ResId)]; ok {
		return
	}
	c.mappednodes[int32(node.ResId)] = struct{}{}

	// dump if necessary
	c.resCount++
	if c.resCount%DumpFreq == 0 {
		c.dumpNodes()
	}

	// find resource's children
	kids := make([]*Node, 0, 2)
	rows, err := c.tmpResStmt.Query(node.ResId, node.ResId)
	panicif(err)
	defer rows.Close()

	for rows.Next() {
		child := &Node{EndTime: math.MaxInt32}
		err := rows.Scan(&child.ResId, &child.StartTime, &child.QualId, &child.Quantity)
		panicif(err)
		node.EndTime = child.StartTime
		kids = append(kids, child)
	}
	panicif(rows.Err())
	rows.Close() // Close is idempotent - and this function is recursive.

	// find resources owner changes (that occurred before children)
	owners, times := c.getNewOwners(node.ResId)

	childOwner := node.OwnerId
	if len(owners) > 0 {
		node.EndTime = times[0]
		childOwner = owners[len(owners)-1]

		lastend := math.MaxInt32
		if len(kids) > 0 {
			lastend = kids[0].StartTime
		}
		times = append(times, lastend)
		for i := range owners {
			n := &Node{ResId: node.ResId,
				OwnerId:   owners[i],
				StartTime: times[i],
				EndTime:   times[i+1],
				QualId:    node.QualId,
				Quantity:  node.Quantity,
			}
			c.nodes = append(c.nodes, n)
		}
	}

	c.nodes = append(c.nodes, node)

	// walk down resource's children
	for _, child := range kids {
		child.OwnerId = childOwner
		c.walkDown(child)
	}
}

func (c *Context) getNewOwners(id int) (owners, times []int) {
	var owner, t int
	rows, err := c.ownerStmt.Query(id, c.Simid)
	panicif(err)
	defer rows.Close()
	for rows.Next() {
		err := rows.Scan(&owner, &t)
		panicif(err)

		if id == owner {
			continue
		}
		owners = append(owners, owner)
		times = append(times, t)
	}
	panicif(rows.Err())
	return owners, times
}

func (c *Context) dumpNodes() {
	c.Log.Printf("    Dumping inventories (%d resources done)...\n", c.resCount)
	tx, err := c.Begin()
	panicif(err)
	stmt := tx.Stmt(c.dumpStmt)

	for _, n := range c.nodes {
		if n.EndTime > n.StartTime {
			_, err = stmt.Exec(c.Simid, n.ResId, n.OwnerId, n.StartTime, n.EndTime, n.QualId, n.Quantity)
			panicif(err)
		}
	}

	err = tx.Commit()
	panicif(err)
	c.nodes = c.nodes[:0]
}
<|endoftext|>