<commit_before>package main

import (
	"encoding/json"
	"flag"
	"fmt"
	"io/ioutil"
	"os"
	"os/signal"
	"sync"
	"time"

	"github.com/golang/glog"
	honeybadger "github.com/honeybadger-io/honeybadger-go"
	"github.com/joho/godotenv"
	"github.com/pkg/errors"
	"github.com/robfig/cron"
	"gopkg.in/yaml.v2"
	"k8s.io/client-go/1.4/kubernetes"
	"k8s.io/client-go/1.4/pkg/api"
	"k8s.io/client-go/1.4/pkg/apis/batch/v1"
	"k8s.io/client-go/1.4/pkg/fields"
	"k8s.io/client-go/1.4/rest"
	"k8s.io/client-go/1.4/tools/clientcmd"
)

var (
	kubeClient   *kubernetes.Clientset
	manager      Manager
	configPath   string
	configDir    string
	scheduleName string
)

func init() {
	flag.StringVar(&configDir, "config", ".", "path to schedule yaml")
	flag.StringVar(&scheduleName, "schedule-name", "schedule.yml", "name of schedule config file")
	flag.StringVar(&configPath, "kubeconfig", "", "absolute path to kubernetes credentials dir")

	if configPath == "" {
		configPath = os.Getenv("KUBE_CONFIG_PATH")
	}

	if os.Getenv("SCHEDULE_NAME") != "" {
		scheduleName = os.Getenv("SCHEDULE_NAME")
	}

	configureHoneybadger()
}

func main() {
	flag.Parse()
	godotenv.Load()
	glog.Info("Kube Scheduler")

	var err error
	kubeClient, err = createKubernetesClient()
	if err != nil {
		honeybadger.Notify(
			"Scheduler could not create kubernetes client",
			honeybadger.Context{"error": err}, honeybadger.Fingerprint{time.Now().String()},
		)
		glog.Fatalf("Could not create kubernetes client: %s", err)
	}

	manager, err = createScheduleManager(filePath(scheduleName))
	if err != nil {
		honeybadger.Notify(
			"Scheduler could not create job manager",
			honeybadger.Context{"error": err}, honeybadger.Fingerprint{time.Now().String()},
		)
		glog.Fatalf("Could not create manager: %s", err)
	}

	cronManager := cron.New()
	for name, job := range manager.jobList {
		glog.Infof("Adding job %s (%s) with schedule %s", name, job.Description, job.Cron)
		go func(name string, job Job) {
			cronManager.AddFunc(job.Cron, func() {
				glog.Infof("Running %s...", name)
				if err := job.Run(); err != nil {
					glog.Warningf("Unable to create & run job: %s", err)
					honeybadger.Notify(fmt.Sprintf("Unable to schedule %s", name), honeybadger.Context{"error": err})
				}
				glog.V(1).Infof("Finished job %s", name)
			})
		}(name, job)
	}
	cronManager.Start()
	defer cronManager.Stop()

	sigs := make(chan os.Signal, 1)
	signal.Notify(sigs, os.Interrupt, os.Kill)
	s := <-sigs
	glog.Infof("Got signal: %s", s)
}

func createKubernetesClient() (*kubernetes.Clientset, error) {
	var (
		kubeConfig *rest.Config
		err        error
	)

	// If no config path is given assume we are in the cluster
	if configPath == "" {
		kubeConfig, err = rest.InClusterConfig()
	} else {
		kubeConfig, err = clientcmd.BuildConfigFromFlags("", configPath)
	}

	if err != nil {
		return nil, errors.Wrap(err, "Failed to connect to kubernetes")
	}

	return kubernetes.NewForConfig(kubeConfig)
}

func createScheduleManager(schedulePath string) (Manager, error) {
	scheduleData, err := ioutil.ReadFile(schedulePath)
	if err != nil {
		return Manager{}, errors.Wrap(err, "Unable to read schedule yaml")
	}

	var config JobList
	if err := yaml.Unmarshal(scheduleData, &config); err != nil {
		return Manager{}, errors.Wrap(err, "Unable to unmarshal schedule yaml")
	}

	return Manager{
		jobList: config,
		jobLock: make(map[string]string),
		mutex:   &sync.Mutex{},
	}, nil
}

type JobList map[string]Job

type Manager struct {
	jobList JobList
	jobLock map[string]string
	mutex   *sync.Mutex
}

func (m *Manager) ReadFromJobLock(template string) (string, bool) {
	m.mutex.Lock()
	defer m.mutex.Unlock()
	v, ok := m.jobLock[template]
	return v, ok
}

func (m *Manager) WriteToJobLock(template, state string) {
	m.mutex.Lock()
	defer m.mutex.Unlock()
	m.jobLock[template] = state
}

func (m *Manager) DeleteFromJobLock(template string) {
	m.mutex.Lock()
	defer m.mutex.Unlock()
	delete(m.jobLock, template)
}

func (m *Manager) NameFromJob(j Job) (string, error) {
	for k, v := range m.jobList {
		if v.Equal(j) {
			return k, nil
		}
	}

	return "", fmt.Errorf("Job is not in job list")
}

type Job struct {
	Cron        string
	Template    string
	Description string
	Args        []string
	Namespace   string
}

func (j Job) Equal(job Job) bool {
	return j.Cron == job.Cron &&
		j.Template == job.Template &&
		j.Description == job.Description &&
		j.Namespace == job.Namespace
}

func (j Job) Run() error {
	jobName, err := manager.NameFromJob(j)
	if err != nil {
		return errors.Wrap(err, "Unable to get name from job")
	}

	if _, ok := manager.ReadFromJobLock(jobName); ok {
		glog.Warningf("Unable to start %s becuase it is already running", jobName)
		return nil
	}

	manager.WriteToJobLock(jobName, "started")
	defer manager.DeleteFromJobLock(jobName)

	if err := createTaskJob(jobName, j); err != nil {
		return errors.Wrap(err, fmt.Sprintf("Failed to created job %s", jobName))
	}

	return nil
}

func createTaskJob(jobName string, j Job) error {
	jobData, err := ioutil.ReadFile(filePath(j.Template))
	if err != nil {
		return errors.Wrap(err, "Error reading job template")
	}

	job := v1.Job{}
	if err = json.Unmarshal(jobData, &job); err != nil {
		return errors.Wrap(err, "Error parsing task pod")
	}

	glog.V(2).Infof("For %s found args: %v", jobName, j.Args)
	glog.V(2).Infof("For %s found namespace: %s", jobName, j.Namespace)

	job.Spec.Template.Spec.Containers[0].Args = j.Args
	job.ObjectMeta.Namespace = j.Namespace

	jobsClient := kubeClient.Batch().Jobs(j.Namespace)
	newJob, err := jobsClient.Create(&job)
	if err != nil {
		return errors.Wrap(err, "Error creating task job")
	}

	glog.V(2).Infof("Created kubernetes job %s", newJob.Name)

	events, err := jobsClient.Watch(api.ListOptions{
		FieldSelector:   fields.OneTermEqualSelector("metadata.name", newJob.Name),
		Watch:           true,
		ResourceVersion: newJob.ResourceVersion,
	})
	if err != nil {
		return errors.Wrap(err, "Error creating job watcher")
	}
	defer events.Stop()

	glog.V(2).Infof("Watching kubernetes job %s for status events", newJob.Name)

	var jobErr error
	for event := range events.ResultChan() {
		job := event.Object.(*v1.Job)
		if len(job.Status.Conditions) > 0 {
			if job.Status.Conditions[0].Type == v1.JobComplete {
				break
			}

			if job.Status.Conditions[0].Type == v1.JobFailed {
				jobErr = fmt.Errorf("Error creating job task")
				break
			}
		}
	}

	if err = jobsClient.Delete(newJob.Name, &api.DeleteOptions{}); err != nil {
		return errors.Wrap(err, "Error deleting job.")
	}

	glog.V(2).Infof("Deleted kubernetes job %s", newJob.Name)

	return jobErr
}

func filePath(filename string) string {
	return fmt.Sprintf("%s/%s", configDir, filename)
}

func configureHoneybadger() {
	honeybadger.Configure(honeybadger.Configuration{APIKey: os.Getenv("HONEYBADGER_API_KEY"), Env: os.Getenv("NAMESPACE")})
}
<commit_msg>added bypass for empty schedules<commit_after>package main

import (
	"encoding/json"
	"flag"
	"fmt"
	"io/ioutil"
	"os"
	"os/signal"
	"sync"
	"time"

	"github.com/golang/glog"
	honeybadger "github.com/honeybadger-io/honeybadger-go"
	"github.com/joho/godotenv"
	"github.com/pkg/errors"
	"github.com/robfig/cron"
	"gopkg.in/yaml.v2"
	"k8s.io/client-go/1.4/kubernetes"
	"k8s.io/client-go/1.4/pkg/api"
	"k8s.io/client-go/1.4/pkg/apis/batch/v1"
	"k8s.io/client-go/1.4/pkg/fields"
	"k8s.io/client-go/1.4/rest"
	"k8s.io/client-go/1.4/tools/clientcmd"
)

var (
	kubeClient   *kubernetes.Clientset
	manager      Manager
	configPath   string
	configDir    string
	scheduleName string
)

func init() {
	flag.StringVar(&configDir, "config", ".", "path to schedule yaml")
	flag.StringVar(&scheduleName, "schedule-name", "schedule.yml", "name of schedule config file")
	flag.StringVar(&configPath, "kubeconfig", "", "absolute path to kubernetes credentials dir")

	if configPath == "" {
		configPath = os.Getenv("KUBE_CONFIG_PATH")
	}

	if os.Getenv("SCHEDULE_NAME") != "" {
		scheduleName = os.Getenv("SCHEDULE_NAME")
	}

	configureHoneybadger()
}

func main() {
	flag.Parse()
	godotenv.Load()
	glog.Info("Kube Scheduler")

	var err error
	kubeClient, err = createKubernetesClient()
	if err != nil {
		honeybadger.Notify(
			"Scheduler could not create kubernetes client",
			honeybadger.Context{"error": err}, honeybadger.Fingerprint{time.Now().String()},
		)
		glog.Fatalf("Could not create kubernetes client: %s", err)
	}

	manager, err = createScheduleManager(filePath(scheduleName))
	if err != nil {
		honeybadger.Notify(
			"Scheduler could not create job manager",
			honeybadger.Context{"error": err}, honeybadger.Fingerprint{time.Now().String()},
		)
		glog.Fatalf("Could not create manager: %s", err)
	}

	cronManager := cron.New()
	// loop over the jobs and add to the cronManager
	for name, job := range manager.jobList {
		if job.Cron == "" {
			glog.Infof("Ignoring job %s (%s) without schedule", name, job.Description)
			continue // Skip adding empty schedules to the cronManager
		}
		glog.Infof("Adding job %s (%s) with schedule %s", name, job.Description, job.Cron)
		go func(name string, job Job) {
			cronManager.AddFunc(job.Cron, func() {
				glog.Infof("Running %s...", name)
				if err := job.Run(); err != nil {
					glog.Warningf("Unable to create & run job: %s", err)
					honeybadger.Notify(fmt.Sprintf("Unable to schedule %s", name), honeybadger.Context{"error": err})
				}
				glog.V(1).Infof("Finished job %s", name)
			})
		}(name, job)
	}
	cronManager.Start()
	defer cronManager.Stop()

	sigs := make(chan os.Signal, 1)
	signal.Notify(sigs, os.Interrupt, os.Kill)
	s := <-sigs
	glog.Infof("Got signal: %s", s)
}

func createKubernetesClient() (*kubernetes.Clientset, error) {
	var (
		kubeConfig *rest.Config
		err        error
	)

	// If no config path is given assume we are in the cluster
	if configPath == "" {
		kubeConfig, err = rest.InClusterConfig()
	} else {
		kubeConfig, err = clientcmd.BuildConfigFromFlags("", configPath)
	}

	if err != nil {
		return nil, errors.Wrap(err, "Failed to connect to kubernetes")
	}

	return kubernetes.NewForConfig(kubeConfig)
}

func createScheduleManager(schedulePath string) (Manager, error) {
	scheduleData, err := ioutil.ReadFile(schedulePath)
	if err != nil {
		return Manager{}, errors.Wrap(err, "Unable to read schedule yaml")
	}

	var config JobList
	if err := yaml.Unmarshal(scheduleData, &config); err != nil {
		return Manager{}, errors.Wrap(err, "Unable to unmarshal schedule yaml")
	}

	return Manager{
		jobList: config,
		jobLock: make(map[string]string),
		mutex:   &sync.Mutex{},
	}, nil
}

type JobList map[string]Job

type Manager struct {
	jobList JobList
	jobLock map[string]string
	mutex   *sync.Mutex
}

func (m *Manager) ReadFromJobLock(template string) (string, bool) {
	m.mutex.Lock()
	defer m.mutex.Unlock()
	v, ok := m.jobLock[template]
	return v, ok
}

func (m *Manager) WriteToJobLock(template, state string) {
	m.mutex.Lock()
	defer m.mutex.Unlock()
	m.jobLock[template] = state
}

func (m *Manager) DeleteFromJobLock(template string) {
	m.mutex.Lock()
	defer m.mutex.Unlock()
	delete(m.jobLock, template)
}

func (m *Manager) NameFromJob(j Job) (string, error) {
	for k, v := range m.jobList {
		if v.Equal(j) {
			return k, nil
		}
	}

	return "", fmt.Errorf("Job is not in job list")
}

type Job struct {
	Cron        string
	Template    string
	Description string
	Args        []string
	Namespace   string
}

func (j Job) Equal(job Job) bool {
	return j.Cron == job.Cron &&
		j.Template == job.Template &&
		j.Description == job.Description &&
		j.Namespace == job.Namespace
}

func (j Job) Run() error {
	jobName, err := manager.NameFromJob(j)
	if err != nil {
		return errors.Wrap(err, "Unable to get name from job")
	}

	if _, ok := manager.ReadFromJobLock(jobName); ok {
		glog.Warningf("Unable to start %s becuase it is already running", jobName)
		return nil
	}

	manager.WriteToJobLock(jobName, "started")
	defer manager.DeleteFromJobLock(jobName)

	if err := createTaskJob(jobName, j); err != nil {
		return errors.Wrap(err, fmt.Sprintf("Failed to created job %s", jobName))
	}

	return nil
}

func createTaskJob(jobName string, j Job) error {
	jobData, err := ioutil.ReadFile(filePath(j.Template))
	if err != nil {
		return errors.Wrap(err, "Error reading job template")
	}

	job := v1.Job{}
	if err = json.Unmarshal(jobData, &job); err != nil {
		return errors.Wrap(err, "Error parsing task pod")
	}

	glog.V(2).Infof("For %s found args: %v", jobName, j.Args)
	glog.V(2).Infof("For %s found namespace: %s", jobName, j.Namespace)

	job.Spec.Template.Spec.Containers[0].Args = j.Args
	job.ObjectMeta.Namespace = j.Namespace

	jobsClient := kubeClient.Batch().Jobs(j.Namespace)
	newJob, err := jobsClient.Create(&job)
	if err != nil {
		return errors.Wrap(err, "Error creating task job")
	}

	glog.V(2).Infof("Created kubernetes job %s", newJob.Name)

	events, err := jobsClient.Watch(api.ListOptions{
		FieldSelector:   fields.OneTermEqualSelector("metadata.name", newJob.Name),
		Watch:           true,
		ResourceVersion: newJob.ResourceVersion,
	})
	if err != nil {
		return errors.Wrap(err, "Error creating job watcher")
	}
	defer events.Stop()

	glog.V(2).Infof("Watching kubernetes job %s for status events", newJob.Name)

	var jobErr error
	for event := range events.ResultChan() {
		job := event.Object.(*v1.Job)
		if len(job.Status.Conditions) > 0 {
			if job.Status.Conditions[0].Type == v1.JobComplete {
				break
			}

			if job.Status.Conditions[0].Type == v1.JobFailed {
				jobErr = fmt.Errorf("Error creating job task")
				break
			}
		}
	}

	if err = jobsClient.Delete(newJob.Name, &api.DeleteOptions{}); err != nil {
		return errors.Wrap(err, "Error deleting job.")
	}

	glog.V(2).Infof("Deleted kubernetes job %s", newJob.Name)

	return jobErr
}

func filePath(filename string) string {
	return fmt.Sprintf("%s/%s", configDir, filename)
}

func configureHoneybadger() {
	honeybadger.Configure(honeybadger.Configuration{APIKey: os.Getenv("HONEYBADGER_API_KEY"), Env: os.Getenv("NAMESPACE")})
}
<|endoftext|>


<commit_before>package main

import (
	"crypto/sha256"
	"encoding/binary"
	"flag"
	"fmt"
	"io"
	"io/ioutil"
	"math/rand"
	"net"
	"os"
	"os/user"
	"path"
	"reflect"
	"strconv"
	"strings"
	"time"

	"github.com/mitnk/goutils/encrypt"
	"github.com/orcaman/concurrent-map"
)

var VERSION = "1.6.3"
var countConnected = 0
var KEY = getKey()
var DEBUG = false

var Servers = cmap.New()

func main() {
	host := flag.String("host", "0.0.0.0", "host")
	port := flag.String("port", "12345", "port")
	_debug := flag.Bool("v", false, "verbose")
	flag.Usage = func() {
		fmt.Printf("Usage of lightsocks v%s:\n", VERSION)
		fmt.Printf("lightsocks [flags]\n")
		flag.PrintDefaults()
		os.Exit(0)
	}
	flag.Parse()
	remote, err := net.Listen("tcp", *host+":"+*port)
	if err != nil {
		fmt.Printf("net listen: %v\n", err)
		os.Exit(2)
	}
	defer remote.Close()
	DEBUG = *_debug

	info("lightsocks v%s", VERSION)
	info("listen on %s:%s", *host, *port)

	go printServersInfo()
	for {
		local, err := remote.Accept()
		if err != nil {
			info("error when accept: %v", err)
			continue
		}
		go handleLocal(local)
	}
}

func printServersInfo() {
	for {
		select {
		case <-time.After(600 * time.Second):
			ts_now := time.Now().Unix()
			keys := Servers.Keys()
			info("[REPORT] We have %d servers connected", len(keys))
			for i, key := range keys {
				if tmp, ok := Servers.Get(key); ok {
					bytes := int64(0)
					ts_span := int64(0)
					if tmp, ok := tmp.(cmap.ConcurrentMap).Get("bytes"); ok {
						bytes = tmp.(int64)
					}
					if tmp, ok := tmp.(cmap.ConcurrentMap).Get("ts"); ok {
						ts_span = ts_now - tmp.(int64)
					}

					str_bytes := ""
					if bytes > 1024*1024*1024 {
						str_bytes += fmt.Sprintf("%.2fG", float64(bytes)/(1024*1024*1024))
					} else if bytes > 1024*1024 {
						str_bytes += fmt.Sprintf("%.2fM", float64(bytes)/(1024*1024))
					} else {
						str_bytes += fmt.Sprintf("%.2fK", float64(bytes)/1024)
					}

					str_span := ""
					if ts_span > 3600 {
						str_span += fmt.Sprintf("%dh", ts_span/3600)
					}
					if ts_span > 60 {
						str_span += fmt.Sprintf("%dm", (ts_span%3600)/60)
					}
					str_span += fmt.Sprintf("%ds", ts_span%60)
					info("[REPORT] [%d][%s] %s: %s", i, str_span, key, str_bytes)
				}
			}
		}
	}
}

func handleLocal(local net.Conn) {
	countConnected += 1
	defer func() {
		local.Close()
		countConnected -= 1
		debug("closed local")
	}()

	debug("local connected: %v", local.RemoteAddr())
	buffer := make([]byte, 1)
	_, err := io.ReadFull(local, buffer)
	if err != nil {
		info("cannot read first byte from local")
		return
	}
	buffer = make([]byte, buffer[0])
	_, err = io.ReadFull(local, buffer)
	if err != nil {
		info("cannot read validation data from local")
		return
	}
	dataCheck, err := encrypt.Decrypt(buffer, KEY)
	if err != nil {
		info("invalid local: %v", err)
		return
	}
	if !reflect.DeepEqual(KEY[8:16], dataCheck) {
		info("invalid local: checker types not eq")
		return
	}
	buffer = make([]byte, 1)
	_, err = io.ReadFull(local, buffer)
	if err != nil {
		info("cannot read size from local")
		return
	}

	buffer = make([]byte, buffer[0])
	_, err = io.ReadFull(local, buffer)
	if err != nil {
		info("cannot read host from local")
		return
	}
	host, err := encrypt.Decrypt(buffer, KEY)
	if err != nil {
		info("ERROR: cannot decrypt host")
		return
	}

	buffer = make([]byte, 2)
	_, err = io.ReadFull(local, buffer)
	if err != nil {
		info("cannot read port from local")
		return
	}
	port := binary.BigEndian.Uint16(buffer)

	url := net.JoinHostPort(string(host), strconv.Itoa(int(port)))
	if strings.Contains(url, "[[") && strings.Contains(url, "]]") {
		// not known yet why, but the url could be something like this:
		// dial tcp: missing port in address [[::ffff:220.249.243.126]]:80
		// we just fix it here.
		url = strings.Replace(url, "[[", "[", 1)
		url = strings.Replace(url, "]]", "]", 1)
	}
	server, err := net.DialTimeout("tcp", url, time.Second*60)
	if err != nil {
		info("ERROR: cannot dial to server %s: %v", url, err)
		return
	}
	info("connected to server: %s", url)
	initServers(url, 0)

	defer func() {
		server.Close()
		deleteServers(url)
		debug("closed server")
	}()

	ch_local := make(chan []byte)
	ch_server := make(chan DataInfo)
	go readDataFromLocal(ch_local, local)
	go readDataFromServer(ch_server, server, url)

	shouldStop := false
	for {
		if shouldStop {
			break
		}

		select {
		case data, ok := <-ch_local:
			if !ok {
				shouldStop = true
				break
			}
			server.Write(data)
		case di, ok := <-ch_server:
			if !ok {
				shouldStop = true
				break
			}
			buffer = encrypt.Encrypt(di.data[:di.size], KEY)
			b := make([]byte, 2)
			binary.BigEndian.PutUint16(b, uint16(len(buffer)))
			local.Write(b)
			local.Write(buffer)
		case <-time.After(3600 * time.Second):
			debug("timeout on %s", url)
			shouldStop = true
			break
		}
	}
}

func readDataFromServer(ch chan DataInfo, conn net.Conn, url string) {
	debug("enter readDataFromServer")
	defer func() {
		debug("leave readDataFromServer")
	}()
	for {
		data := make([]byte, 7000+rand.Intn(2000))
		n, err := conn.Read(data)
		if err != nil {
			break
		}
		incrServers(url, int64(n))
		debug("data from server:\n%s", data[:n])
		ch <- DataInfo{data, n}
	}
	close(ch)
}

func readDataFromLocal(ch chan []byte, conn net.Conn) {
	debug("enter readDataFromLocal")
	defer func() {
		debug("leave readDataFromLocal")
	}()
	for {
		buffer := make([]byte, 2)
		_, err := io.ReadFull(conn, buffer)
		if err != nil {
			break
		}
		size := binary.BigEndian.Uint16(buffer)
		buffer = make([]byte, size)
		_, err = io.ReadFull(conn, buffer)
		if err != nil {
			break
		}
		data, err := encrypt.Decrypt(buffer, KEY)
		if err != nil {
			fmt.Printf("ERROR: cannot decrypt data from local.")
			break
		}
		debug("data from local:\n%s", data)
		ch <- data
	}
	close(ch)
}

func getKey() []byte {
	usr, err := user.Current()
	if err != nil {
		fmt.Printf("user current: %v\n", err)
		os.Exit(2)
	}
	fileKey := path.Join(usr.HomeDir, ".lightsockskey")
	data, err := ioutil.ReadFile(fileKey)
	if err != nil {
		fmt.Printf("Failed to load key file: %v\n", err)
		os.Exit(1)
	}
	s := strings.TrimSpace(string(data))
	sum := sha256.Sum256([]byte(s))
	return sum[:]
}

func initServers(key string, bytes int64) {
	m := cmap.New()
	now := time.Now()
	m.Set("ts", now.Unix())
	m.Set("bytes", bytes)
	Servers.Set(key, m)
}

func incrServers(key string, n int64) {
	if m, ok := Servers.Get(key); ok {
		if tmp, ok := m.(cmap.ConcurrentMap).Get("bytes"); ok {
			m.(cmap.ConcurrentMap).Set("bytes", tmp.(int64)+n)
		}
	} else {
		initServers(key, n)
	}
}

func deleteServers(key string) {
	Servers.Remove(key)
}

func info(format string, a ...interface{}) {
	ts := time.Now().Format("2006-01-02 15:04:05")
	prefix := fmt.Sprintf("[%s][%d] ", ts, countConnected)
	fmt.Printf(prefix+format+"\n", a...)
}

func debug(format string, a ...interface{}) {
	if DEBUG {
		info(format, a...)
	}
}

type DataInfo struct {
	data []byte
	size int
}
<commit_msg>use goixy key instead of its own<commit_after>package main

import (
	"crypto/sha256"
	"encoding/binary"
	"encoding/json"
	"flag"
	"fmt"
	"io"
	"io/ioutil"
	"math/rand"
	"net"
	"os"
	"os/user"
	"path"
	"reflect"
	"strconv"
	"strings"
	"time"

	"github.com/mitnk/goutils/encrypt"
	"github.com/orcaman/concurrent-map"
)

var VERSION = "1.7.0"
var countConnected = 0
var KEY = getKey()
var DEBUG = false

type GoixyConfig struct {
	Key        string
}
var GC GoixyConfig = GoixyConfig{}

var Servers = cmap.New()

func main() {
	host := flag.String("host", "0.0.0.0", "host")
	port := flag.String("port", "12345", "port")
	_debug := flag.Bool("v", false, "verbose")
	flag.Usage = func() {
		fmt.Printf("Usage of lightsocks v%s:\n", VERSION)
		fmt.Printf("lightsocks [flags]\n")
		flag.PrintDefaults()
		os.Exit(0)
	}
	flag.Parse()
	remote, err := net.Listen("tcp", *host+":"+*port)
	if err != nil {
		fmt.Printf("net listen: %v\n", err)
		os.Exit(2)
	}
	defer remote.Close()
	DEBUG = *_debug

	info("lightsocks v%s", VERSION)
	info("listen on %s:%s", *host, *port)

	go printServersInfo()
	for {
		local, err := remote.Accept()
		if err != nil {
			info("error when accept: %v", err)
			continue
		}
		go handleLocal(local)
	}
}

func printServersInfo() {
	for {
		select {
		case <-time.After(600 * time.Second):
			ts_now := time.Now().Unix()
			keys := Servers.Keys()
			info("[REPORT] We have %d servers connected", len(keys))
			for i, key := range keys {
				if tmp, ok := Servers.Get(key); ok {
					bytes := int64(0)
					ts_span := int64(0)
					if tmp, ok := tmp.(cmap.ConcurrentMap).Get("bytes"); ok {
						bytes = tmp.(int64)
					}
					if tmp, ok := tmp.(cmap.ConcurrentMap).Get("ts"); ok {
						ts_span = ts_now - tmp.(int64)
					}

					str_bytes := ""
					if bytes > 1024*1024*1024 {
						str_bytes += fmt.Sprintf("%.2fG", float64(bytes)/(1024*1024*1024))
					} else if bytes > 1024*1024 {
						str_bytes += fmt.Sprintf("%.2fM", float64(bytes)/(1024*1024))
					} else {
						str_bytes += fmt.Sprintf("%.2fK", float64(bytes)/1024)
					}

					str_span := ""
					if ts_span > 3600 {
						str_span += fmt.Sprintf("%dh", ts_span/3600)
					}
					if ts_span > 60 {
						str_span += fmt.Sprintf("%dm", (ts_span%3600)/60)
					}
					str_span += fmt.Sprintf("%ds", ts_span%60)
					info("[REPORT] [%d][%s] %s: %s", i, str_span, key, str_bytes)
				}
			}
		}
	}
}

func handleLocal(local net.Conn) {
	countConnected += 1
	defer func() {
		local.Close()
		countConnected -= 1
		debug("closed local")
	}()

	debug("local connected: %v", local.RemoteAddr())
	buffer := make([]byte, 1)
	_, err := io.ReadFull(local, buffer)
	if err != nil {
		info("cannot read first byte from local")
		return
	}
	buffer = make([]byte, buffer[0])
	_, err = io.ReadFull(local, buffer)
	if err != nil {
		info("cannot read validation data from local")
		return
	}
	dataCheck, err := encrypt.Decrypt(buffer, KEY)
	if err != nil {
		info("invalid local: %v", err)
		return
	}
	if !reflect.DeepEqual(KEY[8:16], dataCheck) {
		info("invalid local: checker types not eq")
		return
	}
	buffer = make([]byte, 1)
	_, err = io.ReadFull(local, buffer)
	if err != nil {
		info("cannot read size from local")
		return
	}

	buffer = make([]byte, buffer[0])
	_, err = io.ReadFull(local, buffer)
	if err != nil {
		info("cannot read host from local")
		return
	}
	host, err := encrypt.Decrypt(buffer, KEY)
	if err != nil {
		info("ERROR: cannot decrypt host")
		return
	}

	buffer = make([]byte, 2)
	_, err = io.ReadFull(local, buffer)
	if err != nil {
		info("cannot read port from local")
		return
	}
	port := binary.BigEndian.Uint16(buffer)

	url := net.JoinHostPort(string(host), strconv.Itoa(int(port)))
	if strings.Contains(url, "[[") && strings.Contains(url, "]]") {
		// not known yet why, but the url could be something like this:
		// dial tcp: missing port in address [[::ffff:220.249.243.126]]:80
		// we just fix it here.
		url = strings.Replace(url, "[[", "[", 1)
		url = strings.Replace(url, "]]", "]", 1)
	}
	server, err := net.DialTimeout("tcp", url, time.Second*60)
	if err != nil {
		info("ERROR: cannot dial to server %s: %v", url, err)
		return
	}
	info("connected to server: %s", url)
	initServers(url, 0)

	defer func() {
		server.Close()
		deleteServers(url)
		debug("closed server")
	}()

	ch_local := make(chan []byte)
	ch_server := make(chan DataInfo)
	go readDataFromLocal(ch_local, local)
	go readDataFromServer(ch_server, server, url)

	shouldStop := false
	for {
		if shouldStop {
			break
		}

		select {
		case data, ok := <-ch_local:
			if !ok {
				shouldStop = true
				break
			}
			server.Write(data)
		case di, ok := <-ch_server:
			if !ok {
				shouldStop = true
				break
			}
			buffer = encrypt.Encrypt(di.data[:di.size], KEY)
			b := make([]byte, 2)
			binary.BigEndian.PutUint16(b, uint16(len(buffer)))
			local.Write(b)
			local.Write(buffer)
		case <-time.After(3600 * time.Second):
			debug("timeout on %s", url)
			shouldStop = true
			break
		}
	}
}

func readDataFromServer(ch chan DataInfo, conn net.Conn, url string) {
	debug("enter readDataFromServer")
	defer func() {
		debug("leave readDataFromServer")
	}()
	for {
		data := make([]byte, 7000+rand.Intn(2000))
		n, err := conn.Read(data)
		if err != nil {
			break
		}
		incrServers(url, int64(n))
		debug("data from server:\n%s", data[:n])
		ch <- DataInfo{data, n}
	}
	close(ch)
}

func readDataFromLocal(ch chan []byte, conn net.Conn) {
	debug("enter readDataFromLocal")
	defer func() {
		debug("leave readDataFromLocal")
	}()
	for {
		buffer := make([]byte, 2)
		_, err := io.ReadFull(conn, buffer)
		if err != nil {
			break
		}
		size := binary.BigEndian.Uint16(buffer)
		buffer = make([]byte, size)
		_, err = io.ReadFull(conn, buffer)
		if err != nil {
			break
		}
		data, err := encrypt.Decrypt(buffer, KEY)
		if err != nil {
			fmt.Printf("ERROR: cannot decrypt data from local.")
			break
		}
		debug("data from local:\n%s", data)
		ch <- data
	}
	close(ch)
}

func getKey() []byte {
	b := getGoixyConfig()
	if b == nil {
		fmt.Printf("Goixy Config not found")
		os.Exit(2)
	}
	err := json.Unmarshal(b, &GC)
	if err != nil {
		fmt.Printf("invalid json in Goixy Config: %v\n", err)
		os.Exit(2)
	}

	s := strings.TrimSpace(GC.Key)
	sum := sha256.Sum256([]byte(s))
	return sum[:]
}

func initServers(key string, bytes int64) {
	m := cmap.New()
	now := time.Now()
	m.Set("ts", now.Unix())
	m.Set("bytes", bytes)
	Servers.Set(key, m)
}

func incrServers(key string, n int64) {
	if m, ok := Servers.Get(key); ok {
		if tmp, ok := m.(cmap.ConcurrentMap).Get("bytes"); ok {
			m.(cmap.ConcurrentMap).Set("bytes", tmp.(int64)+n)
		}
	} else {
		initServers(key, n)
	}
}

func deleteServers(key string) {
	Servers.Remove(key)
}

func getGoixyConfig() []byte {
	usr, err := user.Current()
	if err != nil {
		fmt.Printf("user current: %v\n", err)
		os.Exit(2)
	}
	fileConfig := path.Join(usr.HomeDir, ".goixy/config.json")
	if _, err := os.Stat(fileConfig); os.IsNotExist(err) {
		fmt.Printf("config file is missing: %v\n", fileConfig)
		os.Exit(2)
	}

	data, err := ioutil.ReadFile(fileConfig)
	if err != nil {
		fmt.Printf("failed to load direct-servers file: %v\n", err)
		os.Exit(1)
	}
	return data
}

func info(format string, a ...interface{}) {
	ts := time.Now().Format("2006-01-02 15:04:05")
	prefix := fmt.Sprintf("[%s][%d] ", ts, countConnected)
	fmt.Printf(prefix+format+"\n", a...)
}

func debug(format string, a ...interface{}) {
	if DEBUG {
		info(format, a...)
	}
}

type DataInfo struct {
	data []byte
	size int
}
<|endoftext|>


<commit_before>package linux

import (
	"io/ioutil"
	"strconv"
	"strings"
	"time"
)

type Stat struct {
	CPUStatAll      CPUStat   `json:"cpu_all"`
	CPUStats        []CPUStat `json:"cpus"`
	Interrupts      uint64    `json:"intr"`
	ContextSwitches uint64    `json:"ctxt"`
	BootTime        time.Time `json:"btime"`
	Processes       uint64    `json:"processes"`
	ProcsRunning    uint64    `json:"procs_running"`
	ProcsBlocked    uint64    `json:"procs_blocked"`
}

type CPUStat struct {
	Id        string `json:"id"`
	User      uint64 `json:"user"`
	Nice      uint64 `json:"nice"`
	System    uint64 `json:"system"`
	Idle      uint64 `json:"idle"`
	IOWait    uint64 `json:"iowait"`
	IRQ       uint64 `json:"irq"`
	SoftIRQ   uint64 `json:"softirq"`
	Steal     uint64 `json:"steal"`
	Guest     uint64 `json:"guest"`
	GuestNice uint64 `json:"guest_nice"`
}

func createCPUStat(fields []string) *CPUStat {
	s := CPUStat{}
	s.Id = fields[0]

	for i := 1; i < len(fields); i++ {
		v, _ := strconv.ParseUint(fields[i], 10, 64)
		switch i {
		case 1:
			s.User = v
		case 2:
			s.Nice = v
		case 3:
			s.System = v
		case 4:
			s.Idle = v
		case 5:
			s.IOWait = v
		case 6:
			s.IRQ = v
		case 7:
			s.SoftIRQ = v
		case 8:
			s.Steal = v
		case 9:
			s.Guest = v
		case 10:
			s.GuestNice = v
		}
	}
	return &s
}

func ReadStat(path string) (*Stat, error) {
	b, err := ioutil.ReadFile(path)
	if err != nil {
		return nil, err
	}
	content := string(b)
	lines := strings.Split(content, "\n")

	var stat Stat = Stat{}

	for i, line := range lines {
		fields := strings.Fields(line)
		if len(fields) == 0 {
			continue
		}
		if fields[0][:3] == "cpu" {
			if cpuStat := createCPUStat(fields); cpuStat != nil {
				if i == 0 {
					stat.CPUStatAll = *cpuStat
				} else {
					stat.CPUStats = append(stat.CPUStats, *cpuStat)
				}
			}
		} else if fields[0] == "intr" {
			stat.Interrupts, _ = strconv.ParseUint(fields[1], 10, 64)
		} else if fields[0] == "ctxt" {
			stat.ContextSwitches, _ = strconv.ParseUint(fields[1], 10, 64)
		} else if fields[0] == "btime" {
			seconds, _ := strconv.ParseInt(fields[1], 10, 64)
			stat.BootTime = time.Unix(seconds, 0)
		} else if fields[0] == "processes" {
			stat.Processes, _ = strconv.ParseUint(fields[1], 10, 64)
		} else if fields[0] == "procs_running" {
			stat.ProcsRunning, _ = strconv.ParseUint(fields[1], 10, 64)
		} else if fields[0] == "procs_blocked" {
			stat.ProcsBlocked, _ = strconv.ParseUint(fields[1], 10, 64)
		}
	}
	return &stat, nil
}
<commit_msg>Adding ReadStatFromBytes, refactor<commit_after>package linux

import (
	"bufio"
	"bytes"
	"io/ioutil"
	"strconv"
	"strings"
	"time"
)

type (
	Stat struct {
		CPUStatAll      CPUStat   `json:"cpu_all"`
		CPUStats        []CPUStat `json:"cpus"`
		Interrupts      uint64    `json:"intr"`
		ContextSwitches uint64    `json:"ctxt"`
		BootTime        time.Time `json:"btime"`
		Processes       uint64    `json:"processes"`
		ProcsRunning    uint64    `json:"procs_running"`
		ProcsBlocked    uint64    `json:"procs_blocked"`
	}

	CPUStat struct {
		Id        string `json:"id"`
		User      uint64 `json:"user"`
		Nice      uint64 `json:"nice"`
		System    uint64 `json:"system"`
		Idle      uint64 `json:"idle"`
		IOWait    uint64 `json:"iowait"`
		IRQ       uint64 `json:"irq"`
		SoftIRQ   uint64 `json:"softirq"`
		Steal     uint64 `json:"steal"`
		Guest     uint64 `json:"guest"`
		GuestNice uint64 `json:"guest_nice"`
	}
)

func createCPUStat(fields []string) *CPUStat {

	s := &CPUStat{}

	if len(fields) == 10 {
		s.Id = fields[0]
		s.User, _ = strconv.ParseUint(fields[1], 10, 64)
		s.Nice, _ = strconv.ParseUint(fields[2], 10, 64)
		s.System, _ = strconv.ParseUint(fields[3], 10, 64)
		s.Idle, _ = strconv.ParseUint(fields[4], 10, 64)
		s.IOWait, _ = strconv.ParseUint(fields[5], 10, 64)
		s.IRQ, _ = strconv.ParseUint(fields[6], 10, 64)
		s.SoftIRQ, _ = strconv.ParseUint(fields[7], 10, 64)
		s.Steal, _ = strconv.ParseUint(fields[8], 10, 64)
		s.Guest, _ = strconv.ParseUint(fields[9], 10, 64)
		s.GuestNice, _ = strconv.ParseUint(fields[10], 10, 64)
	}

	return s

}

func ReadStatFromBytes(content []byte) (*Stat, error) {

	stat := &Stat{}

	for s := bufio.NewScanner(bytes.NewReader(content)); s.Scan(); {

		fields := strings.Fields(s.Text())
		if len(fields) == 0 {
			continue
		}

		switch fields[0] {
		case "intr":
			stat.Interrupts, _ = strconv.ParseUint(fields[1], 10, 64)
		case "ctxt":
			stat.ContextSwitches, _ = strconv.ParseUint(fields[1], 10, 64)
		case "btime":
			seconds, _ := strconv.ParseInt(fields[1], 10, 64)
			stat.BootTime = time.Unix(seconds, 0)
		case "processes":
			stat.Processes, _ = strconv.ParseUint(fields[1], 10, 64)
		case "procs_running":
			stat.ProcsRunning, _ = strconv.ParseUint(fields[1], 10, 64)
		case "procs_blocked":
			stat.ProcsBlocked, _ = strconv.ParseUint(fields[1], 10, 64)
		case "cpu":
			if cpuStat := createCPUStat(fields); cpuStat != nil {
				stat.CPUStatAll = *cpuStat
			}
		default:
			if strings.HasPrefix(fields[0], "cpu") {
				if cpuStat := createCPUStat(fields); cpuStat != nil {
					stat.CPUStats = append(stat.CPUStats, *cpuStat)
				}
			}
		}

	}

	return stat, nil

}

func ReadStat(path string) (*Stat, error) {

	b, err := ioutil.ReadFile(path)
	if err != nil {
		return nil, err
	}

	return ReadStatFromBytes(b)

}
<|endoftext|>


<commit_before>// Package es implements an Elasticsearch batch handler. Currently this implementation
// assumes the index format of "index-YY-MM-DD".
package es

import (
	"io"
	stdlog "log"
	"sync"
	"time"

	"github.com/tj/go-elastic/batch"

	"github.com/apex/log"
)

// TODO(tj): allow index configuration
// TODO(tj): allow dumping logs to stderr on timeout
// TODO(tj): allow custom format that does not include .fields etc

// index for the current time.
func index() string {
	return time.Now().Format("logs-06-01-02")
}

// Elasticsearch interface.
type Elasticsearch interface {
	Bulk(io.Reader) error
}

// Config for handler.
type Config struct {
	BufferSize int           // BufferSize is the number of logs to buffer before flush (default: 100)
	Client     Elasticsearch // Client for ES
}

// defaults applies defaults to the config.
func (c *Config) defaults() {
	if c.BufferSize == 0 {
		c.BufferSize = 100
	}
}

// Handler implementation.
type Handler struct {
	*Config

	mu    sync.Mutex
	batch *batch.Batch
}

// New handler with BufferSize
func New(config *Config) *Handler {
	config.defaults()
	return &Handler{
		Config: config,
	}
}

// HandleLog implements log.Handler.
func (h *Handler) HandleLog(e *log.Entry) error {
	h.mu.Lock()
	defer h.mu.Unlock()

	if h.batch == nil {
		h.batch = &batch.Batch{
			Elastic: h.Client,
			Index:   index(),
			Type:    "log",
		}
	}

	h.batch.Add(e)

	if h.batch.Size() >= h.BufferSize {
		h.flush(h.batch)
		h.batch = nil
	}

	return nil
}

// flush the given `batch` asynchronously.
func (h *Handler) flush(batch *batch.Batch) {
	size := batch.Size()
	start := time.Now()
	stdlog.Printf("log/elastic: flushing %d logs", size)

	if err := batch.Flush(); err != nil {
		stdlog.Printf("log/elastic: failed to flush %d logs: %s", size, err)
	}

	stdlog.Printf("log/elastic: flushed %d logs in %s", size, time.Since(start))
}
<commit_msg>fix typo in es package docs<commit_after>// Package es implements an Elasticsearch batch handler. Currently this implementation
// assumes the index format of "logs-YY-MM-DD".
package es

import (
	"io"
	stdlog "log"
	"sync"
	"time"

	"github.com/tj/go-elastic/batch"

	"github.com/apex/log"
)

// TODO(tj): allow index configuration
// TODO(tj): allow dumping logs to stderr on timeout
// TODO(tj): allow custom format that does not include .fields etc

// index for the current time.
func index() string {
	return time.Now().Format("logs-06-01-02")
}

// Elasticsearch interface.
type Elasticsearch interface {
	Bulk(io.Reader) error
}

// Config for handler.
type Config struct {
	BufferSize int           // BufferSize is the number of logs to buffer before flush (default: 100)
	Client     Elasticsearch // Client for ES
}

// defaults applies defaults to the config.
func (c *Config) defaults() {
	if c.BufferSize == 0 {
		c.BufferSize = 100
	}
}

// Handler implementation.
type Handler struct {
	*Config

	mu    sync.Mutex
	batch *batch.Batch
}

// New handler with BufferSize
func New(config *Config) *Handler {
	config.defaults()
	return &Handler{
		Config: config,
	}
}

// HandleLog implements log.Handler.
func (h *Handler) HandleLog(e *log.Entry) error {
	h.mu.Lock()
	defer h.mu.Unlock()

	if h.batch == nil {
		h.batch = &batch.Batch{
			Elastic: h.Client,
			Index:   index(),
			Type:    "log",
		}
	}

	h.batch.Add(e)

	if h.batch.Size() >= h.BufferSize {
		h.flush(h.batch)
		h.batch = nil
	}

	return nil
}

// flush the given `batch` asynchronously.
func (h *Handler) flush(batch *batch.Batch) {
	size := batch.Size()
	start := time.Now()
	stdlog.Printf("log/elastic: flushing %d logs", size)

	if err := batch.Flush(); err != nil {
		stdlog.Printf("log/elastic: failed to flush %d logs: %s", size, err)
	}

	stdlog.Printf("log/elastic: flushed %d logs in %s", size, time.Since(start))
}
<|endoftext|>


<commit_before>package vlc

// #cgo LDFLAGS: -lvlc
// #include <vlc/vlc.h>
// #include <stdlib.h>
import "C"
import (
	"errors"
)

// PlaybackMode defines playback modes for a media list.
type PlaybackMode uint

// Playback modes.
const (
	Default PlaybackMode = iota
	Loop
	Repeat
)

// ListPlayer is an enhanced media player used to play media lists.
type ListPlayer struct {
	player *C.libvlc_media_list_player_t
	list   *MediaList
}

// NewListPlayer creates an instance of a multi-media player.
func NewListPlayer() (*ListPlayer, error) {
	if instance == nil {
		return nil, errors.New("Module must be initialized first")
	}

	if player := C.libvlc_media_list_player_new(instance); player != nil {
		return &ListPlayer{player: player}, nil
	}

	return nil, getError()
}

// Release destroys the media player instance.
func (lp *ListPlayer) Release() error {
	if lp.player == nil {
		return nil
	}

	C.libvlc_media_list_player_release(lp.player)
	lp.player = nil

	return getError()
}

// Play plays the current media list.
func (lp *ListPlayer) Play() error {
	if lp.player == nil {
		return errors.New("A list player must be initialized first")
	}
	if lp.IsPlaying() {
		return nil
	}

	C.libvlc_media_list_player_play(lp.player)
	return getError()
}

// PlayNext plays the next media in the current media list.
func (lp *ListPlayer) PlayNext() error {
	if lp.player == nil {
		return errors.New("A list player must be initialized first")
	}

	if C.libvlc_media_list_player_next(lp.player) < 0 {
		return getError()
	}

	return nil
}

// PlayPrevious plays the previous media in the current media list.
func (lp *ListPlayer) PlayPrevious() error {
	if lp.player == nil {
		return errors.New("A list player must be initialized first")
	}

	if C.libvlc_media_list_player_previous(lp.player) < 0 {
		return getError()
	}

	return nil
}

// PlayAtIndex plays the media at the specified index from the
// current media list.
func (lp ListPlayer) PlayAtIndex(index uint) error {
	if lp.player == nil {
		return errors.New("A list player must be initialized first")
	}

	idx := C.int(index)
	if C.libvlc_media_list_player_play_item_at_index(lp.player, idx) < 0 {
		return getError()
	}

	return nil
}

// IsPlaying returns a boolean value specifying if the player is currently
// playing.
func (lp *ListPlayer) IsPlaying() bool {
	if lp.player == nil {
		return false
	}

	return C.libvlc_media_list_player_is_playing(lp.player) != 0
}

// Stop cancels the currently playing media list, if there is one.
func (lp *ListPlayer) Stop() error {
	if lp.player == nil {
		return errors.New("A list player must be initialized first")
	}
	if !lp.IsPlaying() {
		return nil
	}

	C.libvlc_media_list_player_stop(lp.player)
	return getError()
}

// TogglePause pauses/resumes the player.
// Calling this method has no effect if there is no media.
func (lp *ListPlayer) TogglePause() error {
	if lp.player == nil {
		return errors.New("A list player must be initialized first")
	}

	C.libvlc_media_list_player_pause(lp.player)
	return getError()
}

// SetPlaybackMode sets the player playback mode for the media list.
// By default, it plays the media list once and then stops.
func (lp *ListPlayer) SetPlaybackMode(mode PlaybackMode) error {
	if lp.player == nil {
		return errors.New("A list player must be initialized first")
	}

	m := C.libvlc_playback_mode_t(mode)
	C.libvlc_media_list_player_set_playback_mode(lp.player, m)
	return getError()
}

// MediaState returns the state of the current media.
func (lp *ListPlayer) MediaState() (MediaState, error) {
	if lp.player == nil {
		return 0, errors.New("A list player must be initialized first")
	}

	state := int(C.libvlc_media_list_player_get_state(lp.player))
	return MediaState(state), getError()
}

// MediaList returns the current media list of the player, if one exists
func (lp *ListPlayer) MediaList() *MediaList {
	return lp.list
}

// SetMediaList sets the media list to be played.
func (lp *ListPlayer) SetMediaList(ml *MediaList) error {
	if lp.player == nil {
		return errors.New("A list player must be initialized first")
	}
	if ml.list == nil {
		return errors.New("A media list must be initialized first")
	}

	lp.list = ml
	C.libvlc_media_list_player_set_media_list(lp.player, ml.list)

	return getError()
}

// GetMediaPlayer returns the media player of the media list player instance
func (lp *ListPlayer) GetMediaPlayer() (*Player, error) {
	if lp.player == nil {
		return nil, errors.New("A list player must be initialized first")
	}

	if player := C.libvlc_media_list_player_get_media_player(lp.player); player != nil {
		return &Player{player: player}, nil
	}

	return nil, getError()
}
<commit_msg>Do not claim ownership of the player reference obtained from the list player<commit_after>package vlc

// #cgo LDFLAGS: -lvlc
// #include <vlc/vlc.h>
// #include <stdlib.h>
import "C"
import (
	"errors"
)

// PlaybackMode defines playback modes for a media list.
type PlaybackMode uint

// Playback modes.
const (
	Default PlaybackMode = iota
	Loop
	Repeat
)

// ListPlayer is an enhanced media player used to play media lists.
type ListPlayer struct {
	player *C.libvlc_media_list_player_t
	list   *MediaList
}

// NewListPlayer creates an instance of a multi-media player.
func NewListPlayer() (*ListPlayer, error) {
	if instance == nil {
		return nil, errors.New("Module must be initialized first")
	}

	if player := C.libvlc_media_list_player_new(instance); player != nil {
		return &ListPlayer{player: player}, nil
	}

	return nil, getError()
}

// Release destroys the media player instance.
func (lp *ListPlayer) Release() error {
	if lp.player == nil {
		return nil
	}

	C.libvlc_media_list_player_release(lp.player)
	lp.player = nil

	return getError()
}

// MediaPlayer returns the underlying Player instance of the ListPlayer.
func (lp *ListPlayer) MediaPlayer() (*Player, error) {
	if lp.player == nil {
		return nil, errors.New("A list player must be initialized first")
	}

	player := C.libvlc_media_list_player_get_media_player(lp.player)
	if player == nil {
		return nil, getError()
	}

	// This call will not release the player. Instead, it will decrement the
	// reference count increased by libvlc_media_list_player_get_media_player.
	C.libvlc_media_player_release(player)

	return &Player{player: player}, nil
}

// Play plays the current media list.
func (lp *ListPlayer) Play() error {
	if lp.player == nil {
		return errors.New("A list player must be initialized first")
	}
	if lp.IsPlaying() {
		return nil
	}

	C.libvlc_media_list_player_play(lp.player)
	return getError()
}

// PlayNext plays the next media in the current media list.
func (lp *ListPlayer) PlayNext() error {
	if lp.player == nil {
		return errors.New("A list player must be initialized first")
	}

	if C.libvlc_media_list_player_next(lp.player) < 0 {
		return getError()
	}

	return nil
}

// PlayPrevious plays the previous media in the current media list.
func (lp *ListPlayer) PlayPrevious() error {
	if lp.player == nil {
		return errors.New("A list player must be initialized first")
	}

	if C.libvlc_media_list_player_previous(lp.player) < 0 {
		return getError()
	}

	return nil
}

// PlayAtIndex plays the media at the specified index from the
// current media list.
func (lp ListPlayer) PlayAtIndex(index uint) error {
	if lp.player == nil {
		return errors.New("A list player must be initialized first")
	}

	idx := C.int(index)
	if C.libvlc_media_list_player_play_item_at_index(lp.player, idx) < 0 {
		return getError()
	}

	return nil
}

// IsPlaying returns a boolean value specifying if the player is currently
// playing.
func (lp *ListPlayer) IsPlaying() bool {
	if lp.player == nil {
		return false
	}

	return C.libvlc_media_list_player_is_playing(lp.player) != 0
}

// Stop cancels the currently playing media list, if there is one.
func (lp *ListPlayer) Stop() error {
	if lp.player == nil {
		return errors.New("A list player must be initialized first")
	}
	if !lp.IsPlaying() {
		return nil
	}

	C.libvlc_media_list_player_stop(lp.player)
	return getError()
}

// TogglePause pauses/resumes the player.
// Calling this method has no effect if there is no media.
func (lp *ListPlayer) TogglePause() error {
	if lp.player == nil {
		return errors.New("A list player must be initialized first")
	}

	C.libvlc_media_list_player_pause(lp.player)
	return getError()
}

// SetPlaybackMode sets the player playback mode for the media list.
// By default, it plays the media list once and then stops.
func (lp *ListPlayer) SetPlaybackMode(mode PlaybackMode) error {
	if lp.player == nil {
		return errors.New("A list player must be initialized first")
	}

	m := C.libvlc_playback_mode_t(mode)
	C.libvlc_media_list_player_set_playback_mode(lp.player, m)
	return getError()
}

// MediaState returns the state of the current media.
func (lp *ListPlayer) MediaState() (MediaState, error) {
	if lp.player == nil {
		return 0, errors.New("A list player must be initialized first")
	}

	state := int(C.libvlc_media_list_player_get_state(lp.player))
	return MediaState(state), getError()
}

// MediaList returns the current media list of the player, if one exists
func (lp *ListPlayer) MediaList() *MediaList {
	return lp.list
}

// SetMediaList sets the media list to be played.
func (lp *ListPlayer) SetMediaList(ml *MediaList) error {
	if lp.player == nil {
		return errors.New("A list player must be initialized first")
	}
	if ml.list == nil {
		return errors.New("A media list must be initialized first")
	}

	lp.list = ml
	C.libvlc_media_list_player_set_media_list(lp.player, ml.list)

	return getError()
}
<|endoftext|>


<commit_before>// Copyright 2015 The Numgrad Authors. All rights reserved.
// See the LICENSE file for rights to use this source code.

// Package tipe defines data structures representing Numengrad types.
//
// Go took the usual spelling of type.
package tipe

import (
	"bytes"
	"fmt"
)

type Type interface {
	Sexp() string
	tipe()
}

type Field struct {
	Name string
	Type Type
}

type Func struct {
	In  []*Field
	Out []*Field
}

type Struct struct {
	Fields []*Field
}

type BasicKind int

const (
	Invalid BasicKind = iota
	Bool
	Byte
	Int64
	Float32
	Float64
	Integer
	Float
	String
)

type Basic struct {
	Kind BasicKind
	Name string
}

type Unresolved struct {
	Name interface{} // string or *expr.Selector
}

var (
	_ = Type((*Func)(nil))
	_ = Type((*Struct)(nil))
	_ = Type((*Unresolved)(nil))
)

func (t Func) tipe()       {}
func (t Struct) tipe()     {}
func (t Unresolved) tipe() {}

func (e *Func) Sexp() string {
	return fmt.Sprintf("(functype (in %s) (out %s))", fieldsStr(e.In), fieldsStr(e.Out))
}
func (e *Struct) Sexp() string {
	return fmt.Sprintf("(struct )", fieldsStr(e.Fields))
}
func (e *Unresolved) Sexp() string {
	switch n := e.Name.(type) {
	case string:
		return n
	case interface {
		Sexp() string
	}:
		return "(type " + n.Sexp() + ")"
	default:
		return fmt.Sprintf("unknown:%s", e)
	}
}

func fieldsStr(fields []*Field) string {
	buf := new(bytes.Buffer)
	for i, f := range fields {
		if i > 0 {
			buf.WriteRune(' ')
		}
		fmt.Fprintf(buf, "(%s %s)", f.Name, f.Type.Sexp())
	}
	return buf.String()
}
<commit_msg>lang/tipe: simplify basic types<commit_after>// Copyright 2015 The Numgrad Authors. All rights reserved.
// See the LICENSE file for rights to use this source code.

// Package tipe defines data structures representing Numengrad types.
//
// Go took the usual spelling of type.
package tipe

import (
	"bytes"
	"fmt"
)

type Type interface {
	Sexp() string
	tipe()
}

type Field struct {
	Name string
	Type Type
}

type Func struct {
	In  []*Field
	Out []*Field
}

type Struct struct {
	Fields []*Field
}

type Basic string

const (
	Invalid Basic = "invalid"
	Bool    Basic = "bool"
	Integer Basic = "integer"
	Float   Basic = "float"
	Complex Basic = "complex"
	String  Basic = "string"

	Int64   Basic = "int64"
	Float32 Basic = "float32"
	Float64 Basic = "float64"

	UntypedBool    Basic = "untyped bool"
	UntypedInteger Basic = "untyped integer"
	UntypedFloat   Basic = "untyped float"
	UntypedComplex Basic = "untyped complex"
)

type Named struct {
	Name string // not an identifier, only for debugging
	// TODO: move Ref to a Checker map?
	Ref        interface{} // a *typecheck.Obj after type checking
	Underlying Type
	// TODO: Methods []*Obj
}

type Unresolved struct {
	Name interface{} // string or *expr.Selector
}

var (
	_ = Type(Basic(""))
	_ = Type((*Func)(nil))
	_ = Type((*Struct)(nil))
	_ = Type((*Unresolved)(nil))
)

func (t Basic) tipe()       {}
func (t *Func) tipe()       {}
func (t *Struct) tipe()     {}
func (t *Named) tipe()      {}
func (t *Unresolved) tipe() {}

func (e Basic) Sexp() string { return fmt.Sprintf("(basictype %s)", string(e)) }
func (e *Func) Sexp() string {
	return fmt.Sprintf("(functype (in %s) (out %s))", fieldsStr(e.In), fieldsStr(e.Out))
}
func (e *Struct) Sexp() string {
	return fmt.Sprintf("(structtype %s)", fieldsStr(e.Fields))
}
func (e *Named) Sexp() string {
	u := "nilunderlying"
	if e.Underlying != nil {
		u = e.Underlying.Sexp()
	}
	return fmt.Sprintf("(namedtype %s %s)", e.Name, u)
}
func (e *Unresolved) Sexp() string {
	switch n := e.Name.(type) {
	case string:
		return n
	case interface {
		Sexp() string
	}:
		return "(unresolvedtype " + n.Sexp() + ")"
	default:
		return fmt.Sprintf("unknown:%s", e)
	}
}

func fieldsStr(fields []*Field) string {
	buf := new(bytes.Buffer)
	for i, f := range fields {
		if i > 0 {
			buf.WriteRune(' ')
		}
		fmt.Fprintf(buf, "(%s %s)", f.Name, f.Type.Sexp())
	}
	return buf.String()
}
<|endoftext|>


<commit_before>/*
 *
 *     Copyright 2016 Christophe Hesters
 *
 *     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.
 *
 */
package statistics

type Stopwatch interface {
    ElapsedMillis() int64
    ElapsedNanos() int64
}

type Distribution interface {
    GetSampleCount() int64
    GetMinimum() int64
    GetMaximum() int64
    GetSampleAverage() float64
    GetSampleVariance() float64
    GetSampleStdDeviation() float64
}

type Snapshot struct {
    TimestampTaken int64
    Durations      map[string]Distribution `json:"durations"`
    Counters       map[string]int64        `json:"counters"`
    Samples        map[string]Distribution `json:"samples"`
}

type Store interface {
    StartStopwatch() Stopwatch

    // Finds a duration and returns (true, duration distribution) if found else (false, distribution with all fields set to 0)
    FindDuration(key string) (bool, Distribution)

    // Finds a counter and returns (true, counter value) if found, else (false, 0)
    FindCounter(key string) (bool, int64)

    // Finds a sample and returns (true, sample distribution) if found else (false, distribution with all fields set to 0)
    FindSample(key string) (bool, Distribution)

    // Records the elapsed time of the stopwatch and adds that to the distribution identified by key.
    // Returns the recorded millis
    RecordElapsedTime(key string, stopwatch Stopwatch) int64

    // Records duration of the subject function and adds that to the distribution identified by key.
    // Returns the recorded millis
    MeasureFunc(key string, subject func()) int64

    // Records duration of the subject function and adds that to the distribution identified by key.
    // Returns the recorded millis and the returned value of the subject function
    MeasureFuncWithReturn(key string, subject func() interface{}) (int64, interface{})

    // Increments the counter identified with key by 1. If the counter does not yet exist, it will be created
    // with initial value of 1
    IncrementCounter(key string)

    // Decrements the counter identified with key by 1. If the counter does not yet exist, it will be created
    // with initial value of -1
    DecrementCounter(key string)

    // Adds value to the counter identified with key with key, if the counter does not yet exist, it will be created
    // and initialized to value. Value can be negative.
    AddToCounter(key string, value int)

    // Adds a value to the sample distribution identified by key. If the distribution does not yet exist, value will be it's initial value.
    AddSample(key string, value int64)

    // Clears all durations, counters and samples
    Reset()

    // Creates a snapshot containing all currently registered durations, counters and samples
    Snapshot() Snapshot

    // Creates a snapshot and then calls Reset()
    SnapshotAndReset() Snapshot
}
<commit_msg>Updated the API so that everything is an interface and the implementation can specify the concrete types<commit_after>/*
 *
 *     Copyright 2016 Christophe Hesters
 *
 *     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.
 *
 */
package statistics

type Stopwatch interface {
	ElapsedMillis() int64
	ElapsedNanos() int64
}

type Distribution interface {
	GetSampleCount() int64
	GetMinimum() int64
	GetMaximum() int64
	GetSampleAverage() float64
	GetSampleVariance() float64
	GetSampleStdDeviation() float64
}

type Snapshot interface {
	GetTimestampTaken() int64
	GetDurations() map[string]Distribution
	GetCounters() map[string]int64
	GetSamples() map[string]Distribution
}

type Store interface {
	StartStopwatch() Stopwatch

	// Finds a duration and returns (true, duration distribution) if found else (false, distribution with all fields set to 0)
	FindDuration(key string) (bool, Distribution)

	// Finds a counter and returns (true, counter value) if found, else (false, 0)
	FindCounter(key string) (bool, int64)

	// Finds a sample and returns (true, sample distribution) if found else (false, distribution with all fields set to 0)
	FindSample(key string) (bool, Distribution)

	// Records the elapsed time of the stopwatch and adds that to the distribution identified by key.
	// Returns the recorded millis
	RecordElapsedTime(key string, stopwatch Stopwatch) int64

	// Records duration of the subject function and adds that to the distribution identified by key.
	// Returns the recorded millis
	MeasureFunc(key string, subject func()) int64

	// Records duration of the subject function and adds that to the distribution identified by key.
	// Returns the recorded millis and the returned value of the subject function
	MeasureFuncWithReturn(key string, subject func() interface{}) (int64, interface{})

	// Increments the counter identified with key by 1. If the counter does not yet exist, it will be created
	// with initial value of 1
	IncrementCounter(key string)

	// Decrements the counter identified with key by 1. If the counter does not yet exist, it will be created
	// with initial value of -1
	DecrementCounter(key string)

	// Adds value to the counter identified with key with key, if the counter does not yet exist, it will be created
	// and initialized to value. Value can be negative.
	AddToCounter(key string, value int)

	// Adds a value to the sample distribution identified by key. If the distribution does not yet exist, value will be it's initial value.
	AddSample(key string, value int64)

	// Clears all durations, counters and samples
	Reset()

	// Creates a snapshot containing all currently registered durations, counters and samples
	Snapshot() Snapshot

	// Creates a snapshot and then calls Reset()
	SnapshotAndReset() Snapshot
}
<|endoftext|>


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

package storage

import (
	"errors"
	"fmt"
	"math"
	"math/rand"
	"os"
	"strconv"
	"sync"
	"sync/atomic"
	"time"

	"github.com/m3db/m3db/client"
	"github.com/m3db/m3db/clock"
	"github.com/m3db/m3db/context"
	"github.com/m3db/m3db/retention"
	"github.com/m3db/m3db/storage/block"
	"github.com/m3db/m3db/storage/repair"
	"github.com/m3db/m3db/ts"
	"github.com/m3db/m3x/errors"
	"github.com/m3db/m3x/log"
	"github.com/m3db/m3x/time"

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

var (
	errNoRepairOptions  = errors.New("no repair options")
	errRepairInProgress = errors.New("repair already in progress")
)

type recordFn func(namespace ts.ID, shard databaseShard, diffRes repair.MetadataComparisonResult)

type shardRepairer struct {
	opts      Options
	rpopts    repair.Options
	rtopts    retention.Options
	client    client.AdminClient
	recordFn  recordFn
	logger    xlog.Logger
	scope     tally.Scope
	nowFn     clock.NowFn
	blockSize time.Duration
}

func newShardRepairer(opts Options, rpopts repair.Options) (databaseShardRepairer, error) {
	hostname, err := os.Hostname()
	if err != nil {
		return nil, err
	}

	iopts := opts.InstrumentOptions()
	scope := iopts.MetricsScope().SubScope("database.repair").Tagged(map[string]string{"host": hostname})
	rtopts := opts.RetentionOptions()

	r := shardRepairer{
		opts:      opts,
		rpopts:    rpopts,
		rtopts:    rtopts,
		client:    rpopts.AdminClient(),
		logger:    iopts.Logger(),
		scope:     scope,
		nowFn:     opts.ClockOptions().NowFn(),
		blockSize: rtopts.BlockSize(),
	}
	r.recordFn = r.recordDifferences

	return r, nil
}

func (r shardRepairer) Options() repair.Options {
	return r.rpopts
}

func (r shardRepairer) Repair(
	ctx context.Context,
	namespace ts.ID,
	tr xtime.Range,
	shard databaseShard,
) (repair.MetadataComparisonResult, error) {
	session, err := r.client.DefaultAdminSession()
	if err != nil {
		return repair.MetadataComparisonResult{}, err
	}

	var (
		start    = tr.Start
		end      = tr.End
		origin   = session.Origin()
		replicas = session.Replicas()
	)

	metadata := repair.NewReplicaMetadataComparer(replicas, r.rpopts)
	ctx.RegisterFinalizer(metadata)

	// Add local metadata
	localMetadata, _ := shard.FetchBlocksMetadata(ctx, start, end, math.MaxInt64, 0, true, true)
	localIter := block.NewFilteredBlocksMetadataIter(localMetadata)
	metadata.AddLocalMetadata(origin, localIter)
	localMetadata.Close()

	// Add peer metadata
	peerIter, err := session.FetchBlocksMetadataFromPeers(namespace, shard.ID(), start, end)
	if err != nil {
		return repair.MetadataComparisonResult{}, err
	}
	if err := metadata.AddPeerMetadata(peerIter); err != nil {
		return repair.MetadataComparisonResult{}, err
	}

	metadataRes := metadata.Compare()

	r.recordFn(namespace, shard, metadataRes)

	return metadataRes, nil
}

func (r shardRepairer) recordDifferences(
	namespace ts.ID,
	shard databaseShard,
	diffRes repair.MetadataComparisonResult,
) {
	var (
		shardScope = r.scope.Tagged(map[string]string{
			"namespace": namespace.String(),
			"shard":     strconv.Itoa(int(shard.ID())),
		})
		totalScope        = shardScope.Tagged(map[string]string{"resultType": "total"})
		sizeDiffScope     = shardScope.Tagged(map[string]string{"resultType": "sizeDiff"})
		checksumDiffScope = shardScope.Tagged(map[string]string{"resultType": "checksumDiff"})
	)

	// Record total number of series and total number of blocks
	totalScope.Counter("series").Inc(diffRes.NumSeries)
	totalScope.Counter("blocks").Inc(diffRes.NumBlocks)

	// Record size differences
	sizeDiffScope.Counter("series").Inc(diffRes.SizeDifferences.NumSeries())
	sizeDiffScope.Counter("blocks").Inc(diffRes.SizeDifferences.NumBlocks())

	// Record checksum differences
	checksumDiffScope.Counter("series").Inc(diffRes.ChecksumDifferences.NumSeries())
	checksumDiffScope.Counter("blocks").Inc(diffRes.ChecksumDifferences.NumBlocks())
}

type repairFn func() error

type sleepFn func(d time.Duration)

type repairStatus int

const (
	repairNotStarted repairStatus = iota
	repairSuccess
	repairFailed
)

type repairState struct {
	Status      repairStatus
	NumFailures int
}

type dbRepairer struct {
	sync.Mutex

	database      database
	ropts         repair.Options
	rtopts        retention.Options
	shardRepairer databaseShardRepairer
	repairStates  map[time.Time]repairState

	repairFn            repairFn
	sleepFn             sleepFn
	nowFn               clock.NowFn
	logger              xlog.Logger
	repairInterval      time.Duration
	repairTimeOffset    time.Duration
	repairTimeJitter    time.Duration
	repairCheckInterval time.Duration
	repairMaxRetries    int
	closed              bool
	running             int32
}

func newDatabaseRepairer(database database) (databaseRepairer, error) {
	opts := database.Options()
	nowFn := opts.ClockOptions().NowFn()
	ropts := opts.RepairOptions()
	if ropts == nil {
		return nil, errNoRepairOptions
	}
	if err := ropts.Validate(); err != nil {
		return nil, err
	}

	shardRepairer, err := newShardRepairer(opts, ropts)
	if err != nil {
		return nil, err
	}

	var jitter time.Duration
	if repairJitter := ropts.RepairTimeJitter(); repairJitter > 0 {
		src := rand.NewSource(nowFn().UnixNano())
		jitter = time.Duration(float64(repairJitter) * (float64(src.Int63()) / float64(math.MaxInt64)))
	}

	r := &dbRepairer{
		database:            database,
		ropts:               ropts,
		rtopts:              opts.RetentionOptions(),
		shardRepairer:       shardRepairer,
		repairStates:        make(map[time.Time]repairState),
		sleepFn:             time.Sleep,
		nowFn:               nowFn,
		logger:              opts.InstrumentOptions().Logger(),
		repairInterval:      ropts.RepairInterval(),
		repairTimeOffset:    ropts.RepairTimeOffset(),
		repairTimeJitter:    jitter,
		repairCheckInterval: ropts.RepairCheckInterval(),
		repairMaxRetries:    ropts.RepairMaxRetries(),
	}
	r.repairFn = r.Repair

	return r, nil
}

func (r *dbRepairer) run() {
	var curIntervalStart time.Time

	for {
		r.Lock()
		closed := r.closed
		r.Unlock()

		if closed {
			break
		}

		r.sleepFn(r.repairCheckInterval)

		now := r.nowFn()
		intervalStart := now.Truncate(r.repairInterval)

		// If we haven't reached the offset yet, skip
		target := intervalStart.Add(r.repairTimeOffset + r.repairTimeJitter)
		if now.Before(target) {
			continue
		}

		// If we are in the same interval, we must have already repaired, skip
		if intervalStart == curIntervalStart {
			continue
		}

		curIntervalStart = intervalStart
		if err := r.repairFn(); err != nil {
			r.logger.Errorf("error repairing database: %v", err)
		}
	}
}

func (r *dbRepairer) repairTimeRanges() xtime.Ranges {
	var (
		now       = r.nowFn()
		blockSize = r.rtopts.BlockSize()
		start     = now.Add(-r.rtopts.RetentionPeriod()).Truncate(blockSize)
		end       = now.Add(-r.rtopts.BufferPast()).Truncate(blockSize)
	)

	targetRanges := xtime.NewRanges().AddRange(xtime.Range{Start: start, End: end})
	for t := range r.repairStates {
		if !r.needsRepair(t) {
			targetRanges = targetRanges.RemoveRange(xtime.Range{Start: t, End: t.Add(blockSize)})
		}
	}

	return targetRanges
}

func (r *dbRepairer) needsRepair(t time.Time) bool {
	repairState, exists := r.repairStates[t]
	if !exists {
		return true
	}
	return repairState.Status == repairFailed && repairState.NumFailures < r.repairMaxRetries
}

func (r *dbRepairer) Start() {
	if r.repairInterval <= 0 {
		return
	}

	go r.run()
}

func (r *dbRepairer) Stop() {
	r.Lock()
	r.closed = true
	r.Unlock()
}

func (r *dbRepairer) Repair() error {
	// Don't attempt a repair if the database is not bootstrapped yet
	if !r.database.IsBootstrapped() {
		return nil
	}

	if !atomic.CompareAndSwapInt32(&r.running, 0, 1) {
		return errRepairInProgress
	}

	defer func() {
		atomic.StoreInt32(&r.running, 0)
	}()

	multiErr := xerrors.NewMultiError()
	blockSize := r.rtopts.BlockSize()
	iter := r.repairTimeRanges().Iter()
	for iter.Next() {
		tr := iter.Value()
		err := r.repairWithTimeRange(tr)
		for t := tr.Start; t.Before(tr.End); t = t.Add(blockSize) {
			repairState := r.repairStates[t]
			if err == nil {
				repairState.Status = repairSuccess
			} else {
				repairState.Status = repairFailed
				repairState.NumFailures++
			}
			r.repairStates[t] = repairState
		}
		multiErr = multiErr.Add(err)
	}

	return multiErr.FinalError()
}

func (r *dbRepairer) repairWithTimeRange(tr xtime.Range) error {
	multiErr := xerrors.NewMultiError()
	namespaces := r.database.getOwnedNamespaces()
	for _, n := range namespaces {
		if err := n.Repair(r.shardRepairer, tr); err != nil {
			detailedErr := fmt.Errorf("namespace %s failed to repair time range %v: %v", n.ID().String(), tr, err)
			multiErr = multiErr.Add(detailedErr)
		}
	}
	return multiErr.FinalError()
}

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

package storage

import (
	"errors"
	"fmt"
	"math"
	"math/rand"
	"os"
	"strconv"
	"sync"
	"sync/atomic"
	"time"

	"github.com/m3db/m3db/client"
	"github.com/m3db/m3db/clock"
	"github.com/m3db/m3db/context"
	"github.com/m3db/m3db/retention"
	"github.com/m3db/m3db/storage/block"
	"github.com/m3db/m3db/storage/repair"
	"github.com/m3db/m3db/ts"
	"github.com/m3db/m3x/errors"
	"github.com/m3db/m3x/log"
	"github.com/m3db/m3x/time"

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

var (
	errNoRepairOptions  = errors.New("no repair options")
	errRepairInProgress = errors.New("repair already in progress")
)

type recordFn func(namespace ts.ID, shard databaseShard, diffRes repair.MetadataComparisonResult)

type shardRepairer struct {
	opts      Options
	rpopts    repair.Options
	rtopts    retention.Options
	client    client.AdminClient
	recordFn  recordFn
	logger    xlog.Logger
	scope     tally.Scope
	nowFn     clock.NowFn
	blockSize time.Duration
}

func newShardRepairer(opts Options, rpopts repair.Options) (databaseShardRepairer, error) {
	hostname, err := os.Hostname()
	if err != nil {
		return nil, err
	}

	iopts := opts.InstrumentOptions()
	scope := iopts.MetricsScope().SubScope("database.repair").Tagged(map[string]string{"host": hostname})
	rtopts := opts.RetentionOptions()

	r := shardRepairer{
		opts:      opts,
		rpopts:    rpopts,
		rtopts:    rtopts,
		client:    rpopts.AdminClient(),
		logger:    iopts.Logger(),
		scope:     scope,
		nowFn:     opts.ClockOptions().NowFn(),
		blockSize: rtopts.BlockSize(),
	}
	r.recordFn = r.recordDifferences

	return r, nil
}

func (r shardRepairer) Options() repair.Options {
	return r.rpopts
}

func (r shardRepairer) Repair(
	ctx context.Context,
	namespace ts.ID,
	tr xtime.Range,
	shard databaseShard,
) (repair.MetadataComparisonResult, error) {
	session, err := r.client.DefaultAdminSession()
	if err != nil {
		return repair.MetadataComparisonResult{}, err
	}

	var (
		start    = tr.Start
		end      = tr.End
		origin   = session.Origin()
		replicas = session.Replicas()
	)

	metadata := repair.NewReplicaMetadataComparer(replicas, r.rpopts)
	ctx.RegisterFinalizer(metadata)

	// Add local metadata
	localMetadata, _ := shard.FetchBlocksMetadata(ctx, start, end, math.MaxInt64, 0, true, true)
	ctx.RegisterFinalizer(context.FinalizerFn(localMetadata.Close))

	localIter := block.NewFilteredBlocksMetadataIter(localMetadata)
	metadata.AddLocalMetadata(origin, localIter)

	// Add peer metadata
	peerIter, err := session.FetchBlocksMetadataFromPeers(namespace, shard.ID(), start, end)
	if err != nil {
		return repair.MetadataComparisonResult{}, err
	}
	if err := metadata.AddPeerMetadata(peerIter); err != nil {
		return repair.MetadataComparisonResult{}, err
	}

	metadataRes := metadata.Compare()

	r.recordFn(namespace, shard, metadataRes)

	return metadataRes, nil
}

func (r shardRepairer) recordDifferences(
	namespace ts.ID,
	shard databaseShard,
	diffRes repair.MetadataComparisonResult,
) {
	var (
		shardScope = r.scope.Tagged(map[string]string{
			"namespace": namespace.String(),
			"shard":     strconv.Itoa(int(shard.ID())),
		})
		totalScope        = shardScope.Tagged(map[string]string{"resultType": "total"})
		sizeDiffScope     = shardScope.Tagged(map[string]string{"resultType": "sizeDiff"})
		checksumDiffScope = shardScope.Tagged(map[string]string{"resultType": "checksumDiff"})
	)

	// Record total number of series and total number of blocks
	totalScope.Counter("series").Inc(diffRes.NumSeries)
	totalScope.Counter("blocks").Inc(diffRes.NumBlocks)

	// Record size differences
	sizeDiffScope.Counter("series").Inc(diffRes.SizeDifferences.NumSeries())
	sizeDiffScope.Counter("blocks").Inc(diffRes.SizeDifferences.NumBlocks())

	// Record checksum differences
	checksumDiffScope.Counter("series").Inc(diffRes.ChecksumDifferences.NumSeries())
	checksumDiffScope.Counter("blocks").Inc(diffRes.ChecksumDifferences.NumBlocks())
}

type repairFn func() error

type sleepFn func(d time.Duration)

type repairStatus int

const (
	repairNotStarted repairStatus = iota
	repairSuccess
	repairFailed
)

type repairState struct {
	Status      repairStatus
	NumFailures int
}

type dbRepairer struct {
	sync.Mutex

	database      database
	ropts         repair.Options
	rtopts        retention.Options
	shardRepairer databaseShardRepairer
	repairStates  map[time.Time]repairState

	repairFn            repairFn
	sleepFn             sleepFn
	nowFn               clock.NowFn
	logger              xlog.Logger
	repairInterval      time.Duration
	repairTimeOffset    time.Duration
	repairTimeJitter    time.Duration
	repairCheckInterval time.Duration
	repairMaxRetries    int
	closed              bool
	running             int32
}

func newDatabaseRepairer(database database) (databaseRepairer, error) {
	opts := database.Options()
	nowFn := opts.ClockOptions().NowFn()
	ropts := opts.RepairOptions()
	if ropts == nil {
		return nil, errNoRepairOptions
	}
	if err := ropts.Validate(); err != nil {
		return nil, err
	}

	shardRepairer, err := newShardRepairer(opts, ropts)
	if err != nil {
		return nil, err
	}

	var jitter time.Duration
	if repairJitter := ropts.RepairTimeJitter(); repairJitter > 0 {
		src := rand.NewSource(nowFn().UnixNano())
		jitter = time.Duration(float64(repairJitter) * (float64(src.Int63()) / float64(math.MaxInt64)))
	}

	r := &dbRepairer{
		database:            database,
		ropts:               ropts,
		rtopts:              opts.RetentionOptions(),
		shardRepairer:       shardRepairer,
		repairStates:        make(map[time.Time]repairState),
		sleepFn:             time.Sleep,
		nowFn:               nowFn,
		logger:              opts.InstrumentOptions().Logger(),
		repairInterval:      ropts.RepairInterval(),
		repairTimeOffset:    ropts.RepairTimeOffset(),
		repairTimeJitter:    jitter,
		repairCheckInterval: ropts.RepairCheckInterval(),
		repairMaxRetries:    ropts.RepairMaxRetries(),
	}
	r.repairFn = r.Repair

	return r, nil
}

func (r *dbRepairer) run() {
	var curIntervalStart time.Time

	for {
		r.Lock()
		closed := r.closed
		r.Unlock()

		if closed {
			break
		}

		r.sleepFn(r.repairCheckInterval)

		now := r.nowFn()
		intervalStart := now.Truncate(r.repairInterval)

		// If we haven't reached the offset yet, skip
		target := intervalStart.Add(r.repairTimeOffset + r.repairTimeJitter)
		if now.Before(target) {
			continue
		}

		// If we are in the same interval, we must have already repaired, skip
		if intervalStart == curIntervalStart {
			continue
		}

		curIntervalStart = intervalStart
		if err := r.repairFn(); err != nil {
			r.logger.Errorf("error repairing database: %v", err)
		}
	}
}

func (r *dbRepairer) repairTimeRanges() xtime.Ranges {
	var (
		now       = r.nowFn()
		blockSize = r.rtopts.BlockSize()
		start     = now.Add(-r.rtopts.RetentionPeriod()).Truncate(blockSize)
		end       = now.Add(-r.rtopts.BufferPast()).Truncate(blockSize)
	)

	targetRanges := xtime.NewRanges().AddRange(xtime.Range{Start: start, End: end})
	for t := range r.repairStates {
		if !r.needsRepair(t) {
			targetRanges = targetRanges.RemoveRange(xtime.Range{Start: t, End: t.Add(blockSize)})
		}
	}

	return targetRanges
}

func (r *dbRepairer) needsRepair(t time.Time) bool {
	repairState, exists := r.repairStates[t]
	if !exists {
		return true
	}
	return repairState.Status == repairFailed && repairState.NumFailures < r.repairMaxRetries
}

func (r *dbRepairer) Start() {
	if r.repairInterval <= 0 {
		return
	}

	go r.run()
}

func (r *dbRepairer) Stop() {
	r.Lock()
	r.closed = true
	r.Unlock()
}

func (r *dbRepairer) Repair() error {
	// Don't attempt a repair if the database is not bootstrapped yet
	if !r.database.IsBootstrapped() {
		return nil
	}

	if !atomic.CompareAndSwapInt32(&r.running, 0, 1) {
		return errRepairInProgress
	}

	defer func() {
		atomic.StoreInt32(&r.running, 0)
	}()

	multiErr := xerrors.NewMultiError()
	blockSize := r.rtopts.BlockSize()
	iter := r.repairTimeRanges().Iter()
	for iter.Next() {
		tr := iter.Value()
		err := r.repairWithTimeRange(tr)
		for t := tr.Start; t.Before(tr.End); t = t.Add(blockSize) {
			repairState := r.repairStates[t]
			if err == nil {
				repairState.Status = repairSuccess
			} else {
				repairState.Status = repairFailed
				repairState.NumFailures++
			}
			r.repairStates[t] = repairState
		}
		multiErr = multiErr.Add(err)
	}

	return multiErr.FinalError()
}

func (r *dbRepairer) repairWithTimeRange(tr xtime.Range) error {
	multiErr := xerrors.NewMultiError()
	namespaces := r.database.getOwnedNamespaces()
	for _, n := range namespaces {
		if err := n.Repair(r.shardRepairer, tr); err != nil {
			detailedErr := fmt.Errorf("namespace %s failed to repair time range %v: %v", n.ID().String(), tr, err)
			multiErr = multiErr.Add(detailedErr)
		}
	}
	return multiErr.FinalError()
}

func (r *dbRepairer) IsRepairing() bool {
	return atomic.LoadInt32(&r.running) == 1
}
<|endoftext|>


<commit_before>package storage

import (
	"bytes"
	"database/sql"
	"path"
	"strings"

	"log"
	"regexp"

	"github.com/dinedal/textql/inputs"
	"github.com/dinedal/textql/sqlparser"

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

// SQLite3Storage represents a TextQL compatible SQL backend based on in-memory SQLite3
type SQLite3Storage struct {
	options        *SQLite3Options
	db             *sql.DB
	connID         int
	firstTableName string
}

// SQLite3Options are options passed into SQLite3 connection as needed.
type SQLite3Options struct{}

var (
	sqlite3conn          = []*sqlite3.SQLiteConn{}
	allWhiteSpace        = regexp.MustCompile("^\\s+$")
	tableNameCheckRegEx  = regexp.MustCompile(`.*\[.*\].*`)
	columnNameCheckRegEx = regexp.MustCompile(`.*\[.*\].*`)
)

type entrypoint struct {
	lib  string
	proc string
}

var libNames = []entrypoint{
	{"libgo-sqlite3-extension-functions.so", "sqlite3_extension_init"},
	{"libgo-sqlite3-extension-functions.dylib", "sqlite3_extension_init"},
}

func init() {
	sql.Register("sqlite3_textql",
		&sqlite3.SQLiteDriver{
			ConnectHook: func(conn *sqlite3.SQLiteConn) error {
				for _, v := range libNames {
					conn.LoadExtension(v.lib, v.proc)
				}
				sqlite3conn = append(sqlite3conn, conn)
				return conn.RegisterFunc("regexp", regExp, true)
			},
		})
}

// NewSQLite3StorageWithDefaults returns a SQLite3Storage with the default options.
func NewSQLite3StorageWithDefaults() *SQLite3Storage {
	return NewSQLite3Storage(&SQLite3Options{})
}

// NewSQLite3Storage returns a SQLite3Storage with the SQLite3Options provided applied.
func NewSQLite3Storage(opts *SQLite3Options) *SQLite3Storage {
	sqlite3Storage := &SQLite3Storage{
		options:        opts,
		firstTableName: "",
	}

	sqlite3Storage.open()
	return sqlite3Storage
}

func (sqlite3Storage *SQLite3Storage) open() {
	db, err := sql.Open("sqlite3_textql", ":memory:")

	if err != nil {
		log.Fatalln(err)
	}

	err = db.Ping()

	if err != nil {
		log.Fatalln(err)
	}

	sqlite3Storage.connID = len(sqlite3conn) - 1
	sqlite3Storage.db = db
}

// LoadInput reads the entire Input provided into a table named after the Input name.
// The name is cooreced into a valid SQLite3 table name prior to use.
func (sqlite3Storage *SQLite3Storage) LoadInput(input inputs.Input) {
	tableName := strings.Replace(input.Name(), path.Ext(input.Name()), "", -1)
	sqlite3Storage.createTable(tableName, input.Header(), false)

	tx, txErr := sqlite3Storage.db.Begin()

	if txErr != nil {
		log.Fatalln(txErr)
	}

	stmt := sqlite3Storage.createLoadStmt(tableName, len(input.Header()), tx)

	row := input.ReadRecord()
	for {
		if row == nil {
			break
		}
		sqlite3Storage.loadRow(tableName, len(input.Header()), row, tx, stmt, true)
		row = input.ReadRecord()
	}
	stmt.Close()
	tx.Commit()

	if sqlite3Storage.firstTableName == "" {
		sqlite3Storage.firstTableName = tableName
	}
}

func (sqlite3Storage *SQLite3Storage) createTable(tableName string, columnNames []string, verbose bool) error {
	var buffer bytes.Buffer

	if tableNameCheckRegEx.FindString(tableName) != "" {
		log.Fatalln("Invalid table name", tableName)
	}

	buffer.WriteString("CREATE TABLE IF NOT EXISTS [" + (tableName) + "] (")

	for i, col := range columnNames {
		if columnNameCheckRegEx.FindString(col) != "" {
			log.Fatalln("Invalid table name", col)
		}

		buffer.WriteString("[" + col + "] NUMERIC")

		if i != len(columnNames)-1 {
			buffer.WriteString(", ")
		}
	}

	buffer.WriteString(");")

	_, err := sqlite3Storage.db.Exec(buffer.String())

	if err != nil {
		log.Fatalln(err)
	}

	if verbose {
		log.Println(buffer.String())
	}

	return err
}

func (sqlite3Storage *SQLite3Storage) createLoadStmt(tableName string, colCount int, db *sql.Tx) *sql.Stmt {
	if colCount == 0 {
		log.Fatalln("Nothing to build insert with!")
	}
	var buffer bytes.Buffer

	buffer.WriteString("INSERT INTO [" + (tableName) + "] VALUES (")
	// Don't write the comma for the last column
	for i := 1; i <= colCount; i++ {
		buffer.WriteString("nullif(?,'')")
		if i != colCount {
			buffer.WriteString(", ")
		}
	}

	buffer.WriteString(");")

	stmt, err := db.Prepare(buffer.String())

	if err != nil {
		log.Fatalln(err)
	}
	return stmt
}

func (sqlite3Storage *SQLite3Storage) loadRow(tableName string, colCount int, values []string, db *sql.Tx, stmt *sql.Stmt, verbose bool) error {
	if len(values) == 0 || colCount == 0 {
		return nil
	}

	var vals []interface{}

	for i := 0; i < colCount; i++ {
		if allWhiteSpace.MatchString(values[i]) {
			vals = append(vals, "")
		} else {
			vals = append(vals, values[i])
		}
	}

	_, err := stmt.Exec(vals...)

	if err != nil && verbose {
		log.Printf("Bad row: %v\n", err)
	}

	return err
}

// ExecuteSQLString maps the sqlQuery provided from short hand TextQL to SQL, then
// applies the query to the sqlite3 in memory database, and lastly returns the sql.Rows
// that resulted from the executing query.
func (sqlite3Storage *SQLite3Storage) ExecuteSQLString(sqlQuery string) (*sql.Rows, error) {
	var result *sql.Rows
	var err error

	if strings.Trim(sqlQuery, " ") != "" {
		implictFromSQL := sqlparser.Magicify(sqlQuery, sqlite3Storage.firstTableName)
		result, err = sqlite3Storage.db.Query(implictFromSQL)
		if err != nil {
			return nil, err
		}
	}

	return result, nil
}

// Exec maps the sqlQuery provided from short hand TextQL to SQL, then
// applies the query to the sqlite3 in memory database, and lastly returns the sql.Result
// that resulted from the executing query.
func (sqlite3Storage *SQLite3Storage) Exec(sqlQuery string) (sql.Result, error) {
	var result sql.Result
	var err error

	if strings.Trim(sqlQuery, " ") != "" {
		implictFromSQL := sqlparser.Magicify(sqlQuery, sqlite3Storage.firstTableName)
		result, err = sqlite3Storage.db.Exec(implictFromSQL)
		if err != nil {
			return nil, err
		}
	}

	return result, nil
}

// SaveTo saves the current in memory database to the path provided as a string.
func (sqlite3Storage *SQLite3Storage) SaveTo(path string) error {
	backupDb, openErr := sql.Open("sqlite3_textql", path)
	if openErr != nil {
		return openErr
	}

	backupPingErr := backupDb.Ping()
	if backupPingErr != nil {
		return backupPingErr
	}
	backupConnID := len(sqlite3conn) - 1

	backup, backupStartErr := sqlite3conn[backupConnID].Backup("main", sqlite3conn[sqlite3Storage.connID], "main")
	if backupStartErr != nil {
		return backupStartErr
	}

	_, backupPerformError := backup.Step(-1)
	if backupPerformError != nil {
		return backupPerformError
	}

	backupFinishError := backup.Finish()
	if backupFinishError != nil {
		return backupFinishError
	}

	backupCloseError := backupDb.Close()
	if backupCloseError != nil {
		return backupCloseError
	}

	return nil
}

// Close will close the current database
func (sqlite3Storage *SQLite3Storage) Close() {
	sqlite3Storage.db.Close()
}
<commit_msg>Ensure windows can also load `libgo-sqlite3-extension-functions`<commit_after>package storage

import (
	"bytes"
	"database/sql"
	"path"
	"strings"

	"log"
	"regexp"

	"github.com/dinedal/textql/inputs"
	"github.com/dinedal/textql/sqlparser"

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

// SQLite3Storage represents a TextQL compatible SQL backend based on in-memory SQLite3
type SQLite3Storage struct {
	options        *SQLite3Options
	db             *sql.DB
	connID         int
	firstTableName string
}

// SQLite3Options are options passed into SQLite3 connection as needed.
type SQLite3Options struct{}

var (
	sqlite3conn          = []*sqlite3.SQLiteConn{}
	allWhiteSpace        = regexp.MustCompile("^\\s+$")
	tableNameCheckRegEx  = regexp.MustCompile(`.*\[.*\].*`)
	columnNameCheckRegEx = regexp.MustCompile(`.*\[.*\].*`)
)

type entrypoint struct {
	lib  string
	proc string
}

var libNames = []entrypoint{
	{"libgo-sqlite3-extension-functions.so", "sqlite3_extension_init"},
	{"libgo-sqlite3-extension-functions.dylib", "sqlite3_extension_init"},
	{"libgo-sqlite3-extension-functions.dll", "sqlite3_extension_init"},
}

func init() {
	sql.Register("sqlite3_textql",
		&sqlite3.SQLiteDriver{
			ConnectHook: func(conn *sqlite3.SQLiteConn) error {
				for _, v := range libNames {
					conn.LoadExtension(v.lib, v.proc)
				}
				sqlite3conn = append(sqlite3conn, conn)
				return conn.RegisterFunc("regexp", regExp, true)
			},
		})
}

// NewSQLite3StorageWithDefaults returns a SQLite3Storage with the default options.
func NewSQLite3StorageWithDefaults() *SQLite3Storage {
	return NewSQLite3Storage(&SQLite3Options{})
}

// NewSQLite3Storage returns a SQLite3Storage with the SQLite3Options provided applied.
func NewSQLite3Storage(opts *SQLite3Options) *SQLite3Storage {
	sqlite3Storage := &SQLite3Storage{
		options:        opts,
		firstTableName: "",
	}

	sqlite3Storage.open()
	return sqlite3Storage
}

func (sqlite3Storage *SQLite3Storage) open() {
	db, err := sql.Open("sqlite3_textql", ":memory:")

	if err != nil {
		log.Fatalln(err)
	}

	err = db.Ping()

	if err != nil {
		log.Fatalln(err)
	}

	sqlite3Storage.connID = len(sqlite3conn) - 1
	sqlite3Storage.db = db
}

// LoadInput reads the entire Input provided into a table named after the Input name.
// The name is cooreced into a valid SQLite3 table name prior to use.
func (sqlite3Storage *SQLite3Storage) LoadInput(input inputs.Input) {
	tableName := strings.Replace(input.Name(), path.Ext(input.Name()), "", -1)
	sqlite3Storage.createTable(tableName, input.Header(), false)

	tx, txErr := sqlite3Storage.db.Begin()

	if txErr != nil {
		log.Fatalln(txErr)
	}

	stmt := sqlite3Storage.createLoadStmt(tableName, len(input.Header()), tx)

	row := input.ReadRecord()
	for {
		if row == nil {
			break
		}
		sqlite3Storage.loadRow(tableName, len(input.Header()), row, tx, stmt, true)
		row = input.ReadRecord()
	}
	stmt.Close()
	tx.Commit()

	if sqlite3Storage.firstTableName == "" {
		sqlite3Storage.firstTableName = tableName
	}
}

func (sqlite3Storage *SQLite3Storage) createTable(tableName string, columnNames []string, verbose bool) error {
	var buffer bytes.Buffer

	if tableNameCheckRegEx.FindString(tableName) != "" {
		log.Fatalln("Invalid table name", tableName)
	}

	buffer.WriteString("CREATE TABLE IF NOT EXISTS [" + (tableName) + "] (")

	for i, col := range columnNames {
		if columnNameCheckRegEx.FindString(col) != "" {
			log.Fatalln("Invalid table name", col)
		}

		buffer.WriteString("[" + col + "] NUMERIC")

		if i != len(columnNames)-1 {
			buffer.WriteString(", ")
		}
	}

	buffer.WriteString(");")

	_, err := sqlite3Storage.db.Exec(buffer.String())

	if err != nil {
		log.Fatalln(err)
	}

	if verbose {
		log.Println(buffer.String())
	}

	return err
}

func (sqlite3Storage *SQLite3Storage) createLoadStmt(tableName string, colCount int, db *sql.Tx) *sql.Stmt {
	if colCount == 0 {
		log.Fatalln("Nothing to build insert with!")
	}
	var buffer bytes.Buffer

	buffer.WriteString("INSERT INTO [" + (tableName) + "] VALUES (")
	// Don't write the comma for the last column
	for i := 1; i <= colCount; i++ {
		buffer.WriteString("nullif(?,'')")
		if i != colCount {
			buffer.WriteString(", ")
		}
	}

	buffer.WriteString(");")

	stmt, err := db.Prepare(buffer.String())

	if err != nil {
		log.Fatalln(err)
	}
	return stmt
}

func (sqlite3Storage *SQLite3Storage) loadRow(tableName string, colCount int, values []string, db *sql.Tx, stmt *sql.Stmt, verbose bool) error {
	if len(values) == 0 || colCount == 0 {
		return nil
	}

	var vals []interface{}

	for i := 0; i < colCount; i++ {
		if allWhiteSpace.MatchString(values[i]) {
			vals = append(vals, "")
		} else {
			vals = append(vals, values[i])
		}
	}

	_, err := stmt.Exec(vals...)

	if err != nil && verbose {
		log.Printf("Bad row: %v\n", err)
	}

	return err
}

// ExecuteSQLString maps the sqlQuery provided from short hand TextQL to SQL, then
// applies the query to the sqlite3 in memory database, and lastly returns the sql.Rows
// that resulted from the executing query.
func (sqlite3Storage *SQLite3Storage) ExecuteSQLString(sqlQuery string) (*sql.Rows, error) {
	var result *sql.Rows
	var err error

	if strings.Trim(sqlQuery, " ") != "" {
		implictFromSQL := sqlparser.Magicify(sqlQuery, sqlite3Storage.firstTableName)
		result, err = sqlite3Storage.db.Query(implictFromSQL)
		if err != nil {
			return nil, err
		}
	}

	return result, nil
}

// Exec maps the sqlQuery provided from short hand TextQL to SQL, then
// applies the query to the sqlite3 in memory database, and lastly returns the sql.Result
// that resulted from the executing query.
func (sqlite3Storage *SQLite3Storage) Exec(sqlQuery string) (sql.Result, error) {
	var result sql.Result
	var err error

	if strings.Trim(sqlQuery, " ") != "" {
		implictFromSQL := sqlparser.Magicify(sqlQuery, sqlite3Storage.firstTableName)
		result, err = sqlite3Storage.db.Exec(implictFromSQL)
		if err != nil {
			return nil, err
		}
	}

	return result, nil
}

// SaveTo saves the current in memory database to the path provided as a string.
func (sqlite3Storage *SQLite3Storage) SaveTo(path string) error {
	backupDb, openErr := sql.Open("sqlite3_textql", path)
	if openErr != nil {
		return openErr
	}

	backupPingErr := backupDb.Ping()
	if backupPingErr != nil {
		return backupPingErr
	}
	backupConnID := len(sqlite3conn) - 1

	backup, backupStartErr := sqlite3conn[backupConnID].Backup("main", sqlite3conn[sqlite3Storage.connID], "main")
	if backupStartErr != nil {
		return backupStartErr
	}

	_, backupPerformError := backup.Step(-1)
	if backupPerformError != nil {
		return backupPerformError
	}

	backupFinishError := backup.Finish()
	if backupFinishError != nil {
		return backupFinishError
	}

	backupCloseError := backupDb.Close()
	if backupCloseError != nil {
		return backupCloseError
	}

	return nil
}

// Close will close the current database
func (sqlite3Storage *SQLite3Storage) Close() {
	sqlite3Storage.db.Close()
}
<|endoftext|>