<commit_before>package mplinux

import (
	"bufio"
	"bytes"
	"fmt"
	"io"
	"os"
	"os/exec"
	"regexp"
	"strconv"
	"strings"

	mp "github.com/mackerelio/go-mackerel-plugin-helper"
	"github.com/urfave/cli"
)

const (
	pathVmstat    = "/proc/vmstat"
	pathDiskstats = "/proc/diskstats"
	pathStat      = "/proc/stat"
)

// metric value structure
// note: all metrics are add dynamic at collect*().
var graphdef = map[string]mp.Graphs{}

// LinuxPlugin mackerel plugin for linux
type LinuxPlugin struct {
	Tempfile string
	Typemap  map[string]bool
}

// GraphDefinition interface for mackerelplugin
func (c LinuxPlugin) GraphDefinition() map[string]mp.Graphs {
	var err error

	p := make(map[string]interface{})

	if c.Typemap["all"] || c.Typemap["swap"] {
		err = collectProcVmstat(pathVmstat, &p)
		if err != nil {
			return nil
		}
	}

	if c.Typemap["all"] || c.Typemap["netstat"] {
		err = collectNetworkStat(&p)
		if err != nil {
			return nil
		}
	}

	if c.Typemap["all"] || c.Typemap["diskstats"] {
		err = collectProcDiskstats(pathDiskstats, &p)
		if err != nil {
			return nil
		}
	}

	if c.Typemap["all"] || c.Typemap["proc_stat"] {
		err = collectProcStat(pathStat, &p)
		if err != nil {
			return nil
		}
	}

	if c.Typemap["all"] || c.Typemap["users"] {
		err = collectWho(&p)
		if err != nil {
			return nil
		}
	}

	return graphdef
}

// main function
func doMain(c *cli.Context) error {
	var linux LinuxPlugin

	typemap := map[string]bool{}
	types := c.StringSlice("type")
	// If no `type` is specified, fetch all metrics
	if len(types) == 0 {
		typemap["all"] = true
	} else {
		for _, t := range types {
			typemap[t] = true
		}
	}
	linux.Typemap = typemap
	helper := mp.NewMackerelPlugin(linux)
	helper.Tempfile = c.String("tempfile")

	helper.Run()
	return nil
}

// FetchMetrics interface for mackerelplugin
func (c LinuxPlugin) FetchMetrics() (map[string]interface{}, error) {
	var err error

	p := make(map[string]interface{})

	if c.Typemap["all"] || c.Typemap["swap"] {
		err = collectProcVmstat(pathVmstat, &p)
		if err != nil {
			return nil, err
		}
	}

	if c.Typemap["all"] || c.Typemap["netstat"] {
		err = collectNetworkStat(&p)
		if err != nil {
			return nil, err
		}
	}

	if c.Typemap["all"] || c.Typemap["diskstats"] {
		err = collectProcDiskstats(pathDiskstats, &p)
		if err != nil {
			return nil, err
		}
	}

	if c.Typemap["all"] || c.Typemap["proc_stat"] {
		err = collectProcStat(pathStat, &p)
		if err != nil {
			return nil, err
		}
	}

	if c.Typemap["all"] || c.Typemap["users"] {
		err = collectWho(&p)
		if err != nil {
			return nil, err
		}
	}

	return p, nil
}

// collect who
func collectWho(p *map[string]interface{}) error {
	var err error
	var data string

	graphdef["linux.users"] = mp.Graphs{
		Label: "Linux Users",
		Unit:  "integer",
		Metrics: []mp.Metrics{
			{Name: "users", Label: "Users", Diff: false},
		},
	}

	data, err = getWho()
	if err != nil {
		return err
	}
	err = parseWho(data, p)
	if err != nil {
		return err
	}

	return nil
}

// parsing metrics from /proc/stat
func parseWho(str string, p *map[string]interface{}) error {
	str = strings.TrimSpace(str)
	if str == "" {
		(*p)["users"] = 0
		return nil
	}
	line := strings.Split(str, "\n")
	(*p)["users"] = float64(len(line))

	return nil
}

// Getting who
func getWho() (string, error) {
	cmd := exec.Command("who")
	var out bytes.Buffer
	cmd.Stdout = &out
	err := cmd.Run()
	if err != nil {
		return "", err
	}
	return out.String(), nil
}

// collect /proc/stat
func collectProcStat(path string, p *map[string]interface{}) error {
	var err error
	var data string

	graphdef["linux.interrupts"] = mp.Graphs{
		Label: "Linux Interrupts",
		Unit:  "integer",
		Metrics: []mp.Metrics{
			{Name: "interrupts", Label: "Interrupts", Diff: true},
		},
	}
	graphdef["linux.context_switches"] = mp.Graphs{
		Label: "Linux Context Switches",
		Unit:  "integer",
		Metrics: []mp.Metrics{
			{Name: "context_switches", Label: "Context Switches", Diff: true},
		},
	}
	graphdef["linux.forks"] = mp.Graphs{
		Label: "Linux Forks",
		Unit:  "integer",
		Metrics: []mp.Metrics{
			{Name: "forks", Label: "Forks", Diff: true},
		},
	}

	data, err = getProc(path)
	if err != nil {
		return err
	}
	err = parseProcStat(data, p)
	if err != nil {
		return err
	}

	return nil
}

// parsing metrics from /proc/stat
func parseProcStat(str string, p *map[string]interface{}) error {
	for _, line := range strings.Split(str, "\n") {
		record := strings.Fields(line)
		if len(record) < 2 {
			continue
		}
		name := record[0]
		value, errParse := atof(record[1])
		if errParse != nil {
			return errParse
		}

		if name == "intr" {
			(*p)["interrupts"] = value
		} else if name == "ctxt" {
			(*p)["context_switches"] = value
		} else if name == "processes" {
			(*p)["forks"] = value
		}
	}

	return nil
}

// collect /proc/diskstats
func collectProcDiskstats(path string, p *map[string]interface{}) error {
	file, err := os.Open(path)
	if err != nil {
		return err
	}
	defer file.Close()
	err = parseProcDiskstats(file, p)
	if err != nil {
		return err
	}

	return nil
}

// parsing metrics from diskstats
func parseProcDiskstats(r io.Reader, p *map[string]interface{}) error {
	var elapsedData []mp.Metrics
	var rwtimeData []mp.Metrics

	scanner := bufio.NewScanner(r)

	for scanner.Scan() {
		line := scanner.Text()
		// See also: https://www.kernel.org/doc/Documentation/ABI/testing/procfs-diskstats
		record := strings.Fields(line)
		if len(record) < 14 {
			continue
		}
		device := record[2]
		matched, err := regexp.MatchString("[0-9]$", device)
		if matched || err != nil {
			continue
		}

		(*p)[fmt.Sprintf("iotime_%s", device)], _ = atof(record[12])
		(*p)[fmt.Sprintf("iotime_weighted_%s", device)], _ = atof(record[13])
		elapsedData = append(elapsedData, mp.Metrics{Name: fmt.Sprintf("iotime_%s", device), Label: fmt.Sprintf("%s IO Time", device), Diff: true})
		elapsedData = append(elapsedData, mp.Metrics{Name: fmt.Sprintf("iotime_weighted_%s", device), Label: fmt.Sprintf("%s IO Time Weighted", device), Diff: true})

		(*p)[fmt.Sprintf("tsreading_%s", device)], _ = atof(record[6])
		(*p)[fmt.Sprintf("tswriting_%s", device)], _ = atof(record[10])
		rwtimeData = append(rwtimeData, mp.Metrics{Name: fmt.Sprintf("tsreading_%s", device), Label: fmt.Sprintf("%s Read", device), Diff: true})
		rwtimeData = append(rwtimeData, mp.Metrics{Name: fmt.Sprintf("tswriting_%s", device), Label: fmt.Sprintf("%s Write", device), Diff: true})
	}

	graphdef["linux.disk.elapsed"] = mp.Graphs{
		Label:   "Disk Elapsed IO Time",
		Unit:    "integer",
		Metrics: elapsedData,
	}

	graphdef["linux.disk.rwtime"] = mp.Graphs{
		Label:   "Disk Read/Write Time",
		Unit:    "integer",
		Metrics: rwtimeData,
	}

	return nil
}

// collect ss
func collectNetworkStat(p *map[string]interface{}) error {
	var err error
	var data string

	graphdef["linux.ss"] = mp.Graphs{
		Label: "Linux Network Connection States",
		Unit:  "integer",
		Metrics: []mp.Metrics{
			{Name: "ESTAB", Label: "Established", Diff: false, Stacked: true},
			{Name: "SYN-SENT", Label: "Syn Sent", Diff: false, Stacked: true},
			{Name: "SYN-RECV", Label: "Syn Received", Diff: false, Stacked: true},
			{Name: "FIN-WAIT-1", Label: "Fin Wait 1", Diff: false, Stacked: true},
			{Name: "FIN-WAIT-2", Label: "Fin Wait 2", Diff: false, Stacked: true},
			{Name: "TIME-WAIT", Label: "Time Wait", Diff: false, Stacked: true},
			{Name: "UNCONN", Label: "Close", Diff: false, Stacked: true},
			{Name: "CLOSE-WAIT", Label: "Close Wait", Diff: false, Stacked: true},
			{Name: "LAST-ACK", Label: "Last Ack", Diff: false, Stacked: true},
			{Name: "LISTEN", Label: "Listen", Diff: false, Stacked: true},
			{Name: "CLOSING", Label: "Closing", Diff: false, Stacked: true},
			{Name: "UNKNOWN", Label: "Unknown", Diff: false, Stacked: true},
		},
	}
	data, err = getSs()
	if err != nil {
		return err
	}
	err = parseSs(data, p)
	if err != nil {
		return err
	}

	return nil
}

// parsing metrics from ss
func parseSs(str string, p *map[string]interface{}) error {
	status := 0
	for i, line := range strings.Split(str, "\n") {
		record := strings.Fields(line)
		if len(record) < 5 {
			continue
		}
		if i == 0 {
			if record[0] == "State" {
				// for RHEL6
				status = 0
			} else if record[1] == "State" {
				// for RHEL7
				status = 1
			}
		}
		v, _ := (*p)[record[status]].(float64)
		(*p)[record[status]] = v + 1
	}

	return nil
}

// Getting ss
func getSs() (string, error) {
	cmd := exec.Command("ss", "-na")
	var out bytes.Buffer
	cmd.Stdout = &out
	err := cmd.Run()
	if err != nil {
		return "", err
	}
	return out.String(), nil
}

// collect /proc/vmstat
func collectProcVmstat(path string, p *map[string]interface{}) error {
	graphdef["linux.swap"] = mp.Graphs{
		Label: "Linux Swap Usage",
		Unit:  "integer",
		Metrics: []mp.Metrics{
			{Name: "pswpin", Label: "Swap In", Diff: true},
			{Name: "pswpout", Label: "Swap Out", Diff: true},
		},
	}

	file, err := os.Open(path)
	if err != nil {
		return err
	}
	defer file.Close()
	return parseProcVmstat(file, p)
}

// parsing metrics from /proc/vmstat
func parseProcVmstat(r io.Reader, p *map[string]interface{}) error {
	scanner := bufio.NewScanner(r)

	for scanner.Scan() {
		line := scanner.Text()
		record := strings.Fields(line)
		if len(record) != 2 {
			continue
		}
		var errParse error
		(*p)[record[0]], errParse = atof(record[1])
		if errParse != nil {
			return errParse
		}
	}

	return nil
}

// Getting /proc/*
func getProc(path string) (string, error) {
	cmd := exec.Command("cat", path)
	var out bytes.Buffer
	cmd.Stdout = &out
	err := cmd.Run()
	if err != nil {
		return "", err
	}
	return out.String(), nil
}

// atof
func atof(str string) (float64, error) {
	return strconv.ParseFloat(strings.Trim(str, " "), 64)
}

// Do the plugin
func Do() {
	app := cli.NewApp()
	app.Name = "mackerel-plugin-linux"
	app.Version = version
	app.Usage = "Get metrics from Linux."
	app.Author = "Yuichiro Saito"
	app.Email = "saito@heartbeats.jp"
	app.Flags = flags
	app.Action = doMain

	app.Run(os.Args)
}
<commit_msg>simplify<commit_after>package mplinux

import (
	"bufio"
	"bytes"
	"fmt"
	"io"
	"os"
	"os/exec"
	"regexp"
	"strconv"
	"strings"

	mp "github.com/mackerelio/go-mackerel-plugin-helper"
	"github.com/urfave/cli"
)

const (
	pathVmstat    = "/proc/vmstat"
	pathDiskstats = "/proc/diskstats"
	pathStat      = "/proc/stat"
)

// metric value structure
// note: all metrics are add dynamic at collect*().
var graphdef = map[string]mp.Graphs{}

// LinuxPlugin mackerel plugin for linux
type LinuxPlugin struct {
	Tempfile string
	Typemap  map[string]bool
}

// GraphDefinition interface for mackerelplugin
func (c LinuxPlugin) GraphDefinition() map[string]mp.Graphs {
	var err error

	p := make(map[string]interface{})

	if c.Typemap["all"] || c.Typemap["swap"] {
		err = collectProcVmstat(pathVmstat, &p)
		if err != nil {
			return nil
		}
	}

	if c.Typemap["all"] || c.Typemap["netstat"] {
		err = collectNetworkStat(&p)
		if err != nil {
			return nil
		}
	}

	if c.Typemap["all"] || c.Typemap["diskstats"] {
		err = collectProcDiskstats(pathDiskstats, &p)
		if err != nil {
			return nil
		}
	}

	if c.Typemap["all"] || c.Typemap["proc_stat"] {
		err = collectProcStat(pathStat, &p)
		if err != nil {
			return nil
		}
	}

	if c.Typemap["all"] || c.Typemap["users"] {
		err = collectWho(&p)
		if err != nil {
			return nil
		}
	}

	return graphdef
}

// main function
func doMain(c *cli.Context) error {
	var linux LinuxPlugin

	typemap := map[string]bool{}
	types := c.StringSlice("type")
	// If no `type` is specified, fetch all metrics
	if len(types) == 0 {
		typemap["all"] = true
	} else {
		for _, t := range types {
			typemap[t] = true
		}
	}
	linux.Typemap = typemap
	helper := mp.NewMackerelPlugin(linux)
	helper.Tempfile = c.String("tempfile")

	helper.Run()
	return nil
}

// FetchMetrics interface for mackerelplugin
func (c LinuxPlugin) FetchMetrics() (map[string]interface{}, error) {
	var err error

	p := make(map[string]interface{})

	if c.Typemap["all"] || c.Typemap["swap"] {
		err = collectProcVmstat(pathVmstat, &p)
		if err != nil {
			return nil, err
		}
	}

	if c.Typemap["all"] || c.Typemap["netstat"] {
		err = collectNetworkStat(&p)
		if err != nil {
			return nil, err
		}
	}

	if c.Typemap["all"] || c.Typemap["diskstats"] {
		err = collectProcDiskstats(pathDiskstats, &p)
		if err != nil {
			return nil, err
		}
	}

	if c.Typemap["all"] || c.Typemap["proc_stat"] {
		err = collectProcStat(pathStat, &p)
		if err != nil {
			return nil, err
		}
	}

	if c.Typemap["all"] || c.Typemap["users"] {
		err = collectWho(&p)
		if err != nil {
			return nil, err
		}
	}

	return p, nil
}

// collect who
func collectWho(p *map[string]interface{}) error {
	var err error
	var data string

	graphdef["linux.users"] = mp.Graphs{
		Label: "Linux Users",
		Unit:  "integer",
		Metrics: []mp.Metrics{
			{Name: "users", Label: "Users", Diff: false},
		},
	}

	data, err = getWho()
	if err != nil {
		return err
	}
	err = parseWho(data, p)
	if err != nil {
		return err
	}

	return nil
}

// parsing metrics from /proc/stat
func parseWho(str string, p *map[string]interface{}) error {
	str = strings.TrimSpace(str)
	if str == "" {
		(*p)["users"] = 0
		return nil
	}
	line := strings.Split(str, "\n")
	(*p)["users"] = float64(len(line))

	return nil
}

// Getting who
func getWho() (string, error) {
	cmd := exec.Command("who")
	var out bytes.Buffer
	cmd.Stdout = &out
	err := cmd.Run()
	if err != nil {
		return "", err
	}
	return out.String(), nil
}

// collect /proc/stat
func collectProcStat(path string, p *map[string]interface{}) error {
	var err error
	var data string

	graphdef["linux.interrupts"] = mp.Graphs{
		Label: "Linux Interrupts",
		Unit:  "integer",
		Metrics: []mp.Metrics{
			{Name: "interrupts", Label: "Interrupts", Diff: true},
		},
	}
	graphdef["linux.context_switches"] = mp.Graphs{
		Label: "Linux Context Switches",
		Unit:  "integer",
		Metrics: []mp.Metrics{
			{Name: "context_switches", Label: "Context Switches", Diff: true},
		},
	}
	graphdef["linux.forks"] = mp.Graphs{
		Label: "Linux Forks",
		Unit:  "integer",
		Metrics: []mp.Metrics{
			{Name: "forks", Label: "Forks", Diff: true},
		},
	}

	data, err = getProc(path)
	if err != nil {
		return err
	}
	err = parseProcStat(data, p)
	if err != nil {
		return err
	}

	return nil
}

// parsing metrics from /proc/stat
func parseProcStat(str string, p *map[string]interface{}) error {
	for _, line := range strings.Split(str, "\n") {
		record := strings.Fields(line)
		if len(record) < 2 {
			continue
		}
		name := record[0]
		value, errParse := atof(record[1])
		if errParse != nil {
			return errParse
		}

		if name == "intr" {
			(*p)["interrupts"] = value
		} else if name == "ctxt" {
			(*p)["context_switches"] = value
		} else if name == "processes" {
			(*p)["forks"] = value
		}
	}

	return nil
}

// collect /proc/diskstats
func collectProcDiskstats(path string, p *map[string]interface{}) error {
	file, err := os.Open(path)
	if err != nil {
		return err
	}
	defer file.Close()
	return parseProcDiskstats(file, p)
}

// parsing metrics from diskstats
func parseProcDiskstats(r io.Reader, p *map[string]interface{}) error {
	var elapsedData []mp.Metrics
	var rwtimeData []mp.Metrics

	scanner := bufio.NewScanner(r)

	for scanner.Scan() {
		line := scanner.Text()
		// See also: https://www.kernel.org/doc/Documentation/ABI/testing/procfs-diskstats
		record := strings.Fields(line)
		if len(record) < 14 {
			continue
		}
		device := record[2]
		matched, err := regexp.MatchString("[0-9]$", device)
		if matched || err != nil {
			continue
		}

		(*p)[fmt.Sprintf("iotime_%s", device)], _ = atof(record[12])
		(*p)[fmt.Sprintf("iotime_weighted_%s", device)], _ = atof(record[13])
		elapsedData = append(elapsedData, mp.Metrics{Name: fmt.Sprintf("iotime_%s", device), Label: fmt.Sprintf("%s IO Time", device), Diff: true})
		elapsedData = append(elapsedData, mp.Metrics{Name: fmt.Sprintf("iotime_weighted_%s", device), Label: fmt.Sprintf("%s IO Time Weighted", device), Diff: true})

		(*p)[fmt.Sprintf("tsreading_%s", device)], _ = atof(record[6])
		(*p)[fmt.Sprintf("tswriting_%s", device)], _ = atof(record[10])
		rwtimeData = append(rwtimeData, mp.Metrics{Name: fmt.Sprintf("tsreading_%s", device), Label: fmt.Sprintf("%s Read", device), Diff: true})
		rwtimeData = append(rwtimeData, mp.Metrics{Name: fmt.Sprintf("tswriting_%s", device), Label: fmt.Sprintf("%s Write", device), Diff: true})
	}

	graphdef["linux.disk.elapsed"] = mp.Graphs{
		Label:   "Disk Elapsed IO Time",
		Unit:    "integer",
		Metrics: elapsedData,
	}

	graphdef["linux.disk.rwtime"] = mp.Graphs{
		Label:   "Disk Read/Write Time",
		Unit:    "integer",
		Metrics: rwtimeData,
	}

	return nil
}

// collect ss
func collectNetworkStat(p *map[string]interface{}) error {
	var err error
	var data string

	graphdef["linux.ss"] = mp.Graphs{
		Label: "Linux Network Connection States",
		Unit:  "integer",
		Metrics: []mp.Metrics{
			{Name: "ESTAB", Label: "Established", Diff: false, Stacked: true},
			{Name: "SYN-SENT", Label: "Syn Sent", Diff: false, Stacked: true},
			{Name: "SYN-RECV", Label: "Syn Received", Diff: false, Stacked: true},
			{Name: "FIN-WAIT-1", Label: "Fin Wait 1", Diff: false, Stacked: true},
			{Name: "FIN-WAIT-2", Label: "Fin Wait 2", Diff: false, Stacked: true},
			{Name: "TIME-WAIT", Label: "Time Wait", Diff: false, Stacked: true},
			{Name: "UNCONN", Label: "Close", Diff: false, Stacked: true},
			{Name: "CLOSE-WAIT", Label: "Close Wait", Diff: false, Stacked: true},
			{Name: "LAST-ACK", Label: "Last Ack", Diff: false, Stacked: true},
			{Name: "LISTEN", Label: "Listen", Diff: false, Stacked: true},
			{Name: "CLOSING", Label: "Closing", Diff: false, Stacked: true},
			{Name: "UNKNOWN", Label: "Unknown", Diff: false, Stacked: true},
		},
	}
	data, err = getSs()
	if err != nil {
		return err
	}
	err = parseSs(data, p)
	if err != nil {
		return err
	}

	return nil
}

// parsing metrics from ss
func parseSs(str string, p *map[string]interface{}) error {
	status := 0
	for i, line := range strings.Split(str, "\n") {
		record := strings.Fields(line)
		if len(record) < 5 {
			continue
		}
		if i == 0 {
			if record[0] == "State" {
				// for RHEL6
				status = 0
			} else if record[1] == "State" {
				// for RHEL7
				status = 1
			}
		}
		v, _ := (*p)[record[status]].(float64)
		(*p)[record[status]] = v + 1
	}

	return nil
}

// Getting ss
func getSs() (string, error) {
	cmd := exec.Command("ss", "-na")
	var out bytes.Buffer
	cmd.Stdout = &out
	err := cmd.Run()
	if err != nil {
		return "", err
	}
	return out.String(), nil
}

// collect /proc/vmstat
func collectProcVmstat(path string, p *map[string]interface{}) error {
	graphdef["linux.swap"] = mp.Graphs{
		Label: "Linux Swap Usage",
		Unit:  "integer",
		Metrics: []mp.Metrics{
			{Name: "pswpin", Label: "Swap In", Diff: true},
			{Name: "pswpout", Label: "Swap Out", Diff: true},
		},
	}

	file, err := os.Open(path)
	if err != nil {
		return err
	}
	defer file.Close()
	return parseProcVmstat(file, p)
}

// parsing metrics from /proc/vmstat
func parseProcVmstat(r io.Reader, p *map[string]interface{}) error {
	scanner := bufio.NewScanner(r)

	for scanner.Scan() {
		line := scanner.Text()
		record := strings.Fields(line)
		if len(record) != 2 {
			continue
		}
		var errParse error
		(*p)[record[0]], errParse = atof(record[1])
		if errParse != nil {
			return errParse
		}
	}

	return nil
}

// Getting /proc/*
func getProc(path string) (string, error) {
	cmd := exec.Command("cat", path)
	var out bytes.Buffer
	cmd.Stdout = &out
	err := cmd.Run()
	if err != nil {
		return "", err
	}
	return out.String(), nil
}

// atof
func atof(str string) (float64, error) {
	return strconv.ParseFloat(strings.Trim(str, " "), 64)
}

// Do the plugin
func Do() {
	app := cli.NewApp()
	app.Name = "mackerel-plugin-linux"
	app.Version = version
	app.Usage = "Get metrics from Linux."
	app.Author = "Yuichiro Saito"
	app.Email = "saito@heartbeats.jp"
	app.Flags = flags
	app.Action = doMain

	app.Run(os.Args)
}
<|endoftext|>


<commit_before>/*
 * Copyright (c) 2014 Hewlett-Packard Development Company, L.P.
 *
 * Licensed to the Apache Software Foundation (ASF) under one
 * or more contributor license agreements.  See the NOTICE file
 * distributed with this work for additional information
 * regarding copyright ownership.  The ASF licenses this file
 * to you 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 main

import (
	"flag"
	"os"
	"os/signal"
	"syscall"

	log "code.google.com/p/log4go"
	"github.com/ekarlso/gomdns/config"
	"github.com/ekarlso/gomdns/db"
	"github.com/ekarlso/gomdns/server"
	"github.com/ekarlso/gomdns/stats"
)

var (
	connection string
	nsBind     string
	nsPort     int
	apiBind    string
	apiPort    int

	tsig string
)

func main() {
	fileName := flag.String("config", "config.sample.toml", "Config file")
	flag.StringVar(&connection, "connection", "designate:designate@tcp(localhost:3306)/designate", "Connection string to use for Database")
	flag.StringVar(&nsBind, "nameserver_bind", "", "Addr to listen at")
	flag.IntVar(&nsPort, "nameserver_port", 5053, "Addr to listen at")
	flag.StringVar(&apiBind, "api_bind", "", "Addr to listen at")
	flag.IntVar(&apiPort, "api_port", 5080, "Addr to listen at")
	flag.StringVar(&tsig, "tsig", "", "use MD5 hmac tsig: keyname:base64")

	flag.Usage = func() {
		flag.PrintDefaults()
	}
	flag.Parse()

	cfg, err := config.LoadConfiguration(*fileName)

	if err != nil {
		return
	}

	cfg.NameServerBind = nsBind
	cfg.NameServerPort = nsPort
	cfg.ApiServerBind = apiBind
	cfg.ApiServerPort = apiPort
	cfg.StorageDSN = connection

	log.Info("Database is at connection %s", cfg.StorageDSN)

	db.Setup(cfg.StorageDSN)
	// Setup db access
	if db.CheckDB(cfg.StorageDSN) != true {
		log.Warn("Error verifying database connectivity, see above for errors")
		os.Exit(1)
	}

	stats.Setup()

	srv, err := server.NewServer(cfg)
	srv.ListenAndServe()

	sig := make(chan os.Signal)
	signal.Notify(sig, syscall.SIGINT, syscall.SIGTERM)

forever:
	for {
		select {
		case s := <-sig:
			log.Info("Signal (%d) received, stopping\n", s)
			break forever
		}
	}
}
<commit_msg>Move setup call<commit_after>/*
 * Copyright (c) 2014 Hewlett-Packard Development Company, L.P.
 *
 * Licensed to the Apache Software Foundation (ASF) under one
 * or more contributor license agreements.  See the NOTICE file
 * distributed with this work for additional information
 * regarding copyright ownership.  The ASF licenses this file
 * to you 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 main

import (
	"flag"
	"os"
	"os/signal"
	"syscall"

	log "code.google.com/p/log4go"
	"github.com/ekarlso/gomdns/config"
	"github.com/ekarlso/gomdns/db"
	"github.com/ekarlso/gomdns/server"
	"github.com/ekarlso/gomdns/stats"
)

var (
	connection string
	nsBind     string
	nsPort     int
	apiBind    string
	apiPort    int

	tsig string
)

func main() {
	fileName := flag.String("config", "config.sample.toml", "Config file")
	flag.StringVar(&connection, "connection", "designate:designate@tcp(localhost:3306)/designate", "Connection string to use for Database")
	flag.StringVar(&nsBind, "nameserver_bind", "", "Addr to listen at")
	flag.IntVar(&nsPort, "nameserver_port", 5053, "Addr to listen at")
	flag.StringVar(&apiBind, "api_bind", "", "Addr to listen at")
	flag.IntVar(&apiPort, "api_port", 5080, "Addr to listen at")
	flag.StringVar(&tsig, "tsig", "", "use MD5 hmac tsig: keyname:base64")

	flag.Usage = func() {
		flag.PrintDefaults()
	}
	flag.Parse()

	cfg, err := config.LoadConfiguration(*fileName)

	if err != nil {
		return
	}

	cfg.NameServerBind = nsBind
	cfg.NameServerPort = nsPort
	cfg.ApiServerBind = apiBind
	cfg.ApiServerPort = apiPort
	cfg.StorageDSN = connection

	log.Info("Database is at connection %s", cfg.StorageDSN)

	stats.Setup(cfg)

	db.Setup(cfg.StorageDSN)
	// Setup db access
	if db.CheckDB(cfg.StorageDSN) != true {
		log.Warn("Error verifying database connectivity, see above for errors")
		os.Exit(1)
	}

	srv, err := server.NewServer(cfg)
	srv.ListenAndServe()

	sig := make(chan os.Signal)
	signal.Notify(sig, syscall.SIGINT, syscall.SIGTERM)

forever:
	for {
		select {
		case s := <-sig:
			log.Info("Signal (%d) received, stopping\n", s)
			break forever
		}
	}
}
<|endoftext|>


<commit_before>package task

import (
	"fmt"
	"neon/build"
	"neon/util"
)

func init() {
	build.TaskMap["for"] = build.TaskDescriptor{
		Constructor: For,
		Help: `For loop

Arguments:
- for: the name of the variable to set at each loop iteration.
- in: the list of values or expression that generates this list.
- do: the block of steps to execute at each loop iteration.

Examples:
# create empty files
- for: file
  in:  ["foo", "bar"]
  do:
  - touch: "#{file}"
# print first 10 integers
- for: i
  in: range(10)
  do:
  - print: "#{i}"`,
	}
}

func For(target *build.Target, args util.Object) (build.Task, error) {
	fields := []string{"for", "in", "do"}
	if err := CheckFields(args, fields, fields); err != nil {
		return nil, err
	}
	variable, err := args.GetString("for")
	if err != nil {
		return nil, fmt.Errorf("'for' field of a 'for' loop must be a string")
	}
	list, err := args.GetList("in")
	expression := ""
	if err != nil {
		expression, err = args.GetString("in")
		if err != nil {
			return nil, fmt.Errorf("'in' field of 'for' loop must be a list or string")
		}
	}
	steps, err := ParseSteps(target, args, "do")
	if err != nil {
		return nil, err
	}
	return func() error {
		if expression != "" {
			result, err := target.Build.Context.Evaluate(expression)
			if err != nil {
				return fmt.Errorf("evaluating in field of for loop: %v", err)
			}
			list, err = util.ToList(result)
			if err != nil {
				return fmt.Errorf("'in' field of 'for' loop must be an expression that returns a list")
			}
		}
		for _, value := range list {
			target.Build.Context.SetProperty(variable, value)
			if err != nil {
				return err
			}
			err := RunSteps(steps)
			if err != nil {
				return err
			}
		}
		return nil
	}, nil
}
<commit_msg>Fixed doc<commit_after>package task

import (
	"fmt"
	"neon/build"
	"neon/util"
)

func init() {
	build.TaskMap["for"] = build.TaskDescriptor{
		Constructor: For,
		Help: `For loop.

Arguments:
- for: the name of the variable to set at each loop iteration.
- in: the list of values or expression that generates this list.
- do: the block of steps to execute at each loop iteration.

Examples:
# create empty files
- for: file
  in:  ["foo", "bar"]
  do:
  - touch: "#{file}"
# print first 10 integers
- for: i
  in: range(10)
  do:
  - print: "#{i}"`,
	}
}

func For(target *build.Target, args util.Object) (build.Task, error) {
	fields := []string{"for", "in", "do"}
	if err := CheckFields(args, fields, fields); err != nil {
		return nil, err
	}
	variable, err := args.GetString("for")
	if err != nil {
		return nil, fmt.Errorf("'for' field of a 'for' loop must be a string")
	}
	list, err := args.GetList("in")
	expression := ""
	if err != nil {
		expression, err = args.GetString("in")
		if err != nil {
			return nil, fmt.Errorf("'in' field of 'for' loop must be a list or string")
		}
	}
	steps, err := ParseSteps(target, args, "do")
	if err != nil {
		return nil, err
	}
	return func() error {
		if expression != "" {
			result, err := target.Build.Context.Evaluate(expression)
			if err != nil {
				return fmt.Errorf("evaluating in field of for loop: %v", err)
			}
			list, err = util.ToList(result)
			if err != nil {
				return fmt.Errorf("'in' field of 'for' loop must be an expression that returns a list")
			}
		}
		for _, value := range list {
			target.Build.Context.SetProperty(variable, value)
			if err != nil {
				return err
			}
			err := RunSteps(steps)
			if err != nil {
				return err
			}
		}
		return nil
	}, nil
}
<|endoftext|>


<commit_before>package commands

import (
	"encoding/gob"
	"flag"
	"fmt"
	"log"
	"net/url"
	"os"
	"path/filepath"
	"strings"
	"time"

	"encoding/json"

	"github.com/astaxie/beego"
	beegoCache "github.com/astaxie/beego/cache"
	_ "github.com/astaxie/beego/cache/memcache"
	_ "github.com/astaxie/beego/cache/redis"
	"github.com/astaxie/beego/logs"
	"github.com/astaxie/beego/orm"
	"github.com/lifei6671/gocaptcha"
	"github.com/lifei6671/mindoc/cache"
	"github.com/lifei6671/mindoc/commands/migrate"
	"github.com/lifei6671/mindoc/conf"
	"github.com/lifei6671/mindoc/models"
	"github.com/lifei6671/mindoc/utils/filetil"
)

// RegisterDataBase 注册数据库
func RegisterDataBase() {
	beego.Info("正在初始化数据库配置.")
	adapter := beego.AppConfig.String("db_adapter")

	if adapter == "mysql" {
		host := beego.AppConfig.String("db_host")
		database := beego.AppConfig.String("db_database")
		username := beego.AppConfig.String("db_username")
		password := beego.AppConfig.String("db_password")

		timezone := beego.AppConfig.String("timezone")
		location, err := time.LoadLocation(timezone)
		if err == nil {
			orm.DefaultTimeLoc = location
		} else {
			beego.Error("加载时区配置信息失败,请检查是否存在ZONEINFO环境变量:", err)
		}

		port := beego.AppConfig.String("db_port")

		dataSource := fmt.Sprintf("%s:%s@tcp(%s:%s)/%s?charset=utf8mb4&parseTime=true&loc=%s", username, password, host, port, database, url.QueryEscape(timezone))

		if err := orm.RegisterDataBase("default", "mysql", dataSource); err != nil {
			beego.Error("注册默认数据库失败:", err)
			os.Exit(1)
		}
	} else if adapter == "sqlite3" {
		orm.DefaultTimeLoc = time.UTC
		database := beego.AppConfig.String("db_database")
		if strings.HasPrefix(database, "./") {
			database = filepath.Join(conf.WorkingDirectory, string(database[1:]))
		}

		dbPath := filepath.Dir(database)
		os.MkdirAll(dbPath, 0777)

		err := orm.RegisterDataBase("default", "sqlite3", database)

		if err != nil {
			beego.Error("注册默认数据库失败:", err)
		}
	} else {
		beego.Error("不支持的数据库类型.")
		os.Exit(1)
	}
	beego.Info("数据库初始化完成.")
}

// RegisterModel 注册Model
func RegisterModel() {
	orm.RegisterModelWithPrefix(conf.GetDatabasePrefix(),
		new(models.Member),
		new(models.Book),
		new(models.Relationship),
		new(models.Option),
		new(models.Document),
		new(models.Attachment),
		new(models.Logger),
		new(models.MemberToken),
		new(models.DocumentHistory),
		new(models.Migration),
		new(models.Label),
	)
	//migrate.RegisterMigration()
}

// RegisterLogger 注册日志
func RegisterLogger(log string) {

	logs.SetLogFuncCall(true)
	logs.SetLogger("console")
	logs.EnableFuncCallDepth(true)
	logs.Async()

	logPath := filepath.Join(log, "log.log")

	if _, err := os.Stat(logPath); os.IsNotExist(err) {

		os.MkdirAll(log, 0777)

		if f, err := os.Create(logPath); err == nil {
			f.Close()
			config := make(map[string]interface{}, 1)

			config["filename"] = logPath

			b, _ := json.Marshal(config)

			beego.SetLogger("file", string(b))
		}
	}

	beego.SetLogFuncCall(true)
	beego.BeeLogger.Async()
}

// RunCommand 注册orm命令行工具
func RegisterCommand() {

	if len(os.Args) >= 2 && os.Args[1] == "install" {
		ResolveCommand(os.Args[2:])
		Install()
	} else if len(os.Args) >= 2 && os.Args[1] == "version" {
		CheckUpdate()
		os.Exit(0)
	} else if len(os.Args) >= 2 && os.Args[1] == "migrate" {
		ResolveCommand(os.Args[2:])
		migrate.RunMigration()
	}
}

//注册模板函数
func RegisterFunction() {
	beego.AddFuncMap("config", models.GetOptionValue)

	beego.AddFuncMap("cdn", func(p string) string {
		cdn := beego.AppConfig.DefaultString("cdn", "")
		if strings.HasPrefix(p, "http://") || strings.HasPrefix(p, "https://") {
			return p
		}
		//如果没有设置cdn，则使用baseURL拼接
		if cdn == "" {
			baseUrl := beego.AppConfig.DefaultString("baseurl", "")

			if strings.HasPrefix(p, "/") && strings.HasSuffix(baseUrl, "/") {
				return baseUrl + p[1:]
			}
			if !strings.HasPrefix(p, "/") && !strings.HasSuffix(baseUrl, "/") {
				return baseUrl + "/" + p
			}
			return baseUrl + p
		}
		if strings.HasPrefix(p, "/") && strings.HasSuffix(cdn, "/") {
			return cdn + string(p[1:])
		}
		if !strings.HasPrefix(p, "/") && !strings.HasSuffix(cdn, "/") {
			return cdn + "/" + p
		}
		return cdn + p
	})

	beego.AddFuncMap("cdnjs", conf.URLForWithCdnJs)
	beego.AddFuncMap("cdncss", conf.URLForWithCdnCss)
	beego.AddFuncMap("cdnimg", conf.URLForWithCdnImage)
	//重写url生成，支持配置域名以及域名前缀
	beego.AddFuncMap("urlfor", conf.URLFor)
	beego.AddFuncMap("date_format", func(t time.Time, format string) string {
		return t.Local().Format(format)
	})
}

//解析命令
func ResolveCommand(args []string) {
	flagSet := flag.NewFlagSet("MinDoc command: ", flag.ExitOnError)
	flagSet.StringVar(&conf.ConfigurationFile, "config", "", "MinDoc configuration file.")
	flagSet.StringVar(&conf.WorkingDirectory, "dir", "", "MinDoc working directory.")
	flagSet.StringVar(&conf.LogFile, "log", "", "MinDoc log file path.")

	flagSet.Parse(args)

	if conf.WorkingDirectory == "" {
		if p, err := filepath.Abs(os.Args[0]); err == nil {
			conf.WorkingDirectory = filepath.Dir(p)
		}
	}
	if conf.LogFile == "" {
		conf.LogFile = filepath.Join(conf.WorkingDirectory, "logs")
	}
	if conf.ConfigurationFile == "" {
		conf.ConfigurationFile = filepath.Join(conf.WorkingDirectory, "conf", "app.conf")
		config := filepath.Join(conf.WorkingDirectory, "conf", "app.conf.example")
		if !filetil.FileExists(conf.ConfigurationFile) && filetil.FileExists(config) {
			filetil.CopyFile(conf.ConfigurationFile, config)
		}
	}
	gocaptcha.ReadFonts(filepath.Join(conf.WorkingDirectory, "static", "fonts"), ".ttf")

	err := beego.LoadAppConfig("ini", conf.ConfigurationFile)

	if err != nil {
		log.Println("An error occurred:", err)
		os.Exit(1)
	}
	uploads := filepath.Join(conf.WorkingDirectory, "uploads")

	os.MkdirAll(uploads, 0666)

	beego.BConfig.WebConfig.StaticDir["/static"] = filepath.Join(conf.WorkingDirectory, "static")
	beego.BConfig.WebConfig.StaticDir["/uploads"] = uploads
	beego.BConfig.WebConfig.ViewsPath = filepath.Join(conf.WorkingDirectory, "views")

	fonts := filepath.Join(conf.WorkingDirectory, "static", "fonts")

	if !filetil.FileExists(fonts) {
		log.Fatal("Font path not exist.")
	}
	gocaptcha.ReadFonts(filepath.Join(conf.WorkingDirectory, "static", "fonts"), ".ttf")

	RegisterDataBase()
	RegisterCache()
	RegisterModel()
	RegisterLogger(conf.LogFile)
}

//注册缓存管道
func RegisterCache() {
	isOpenCache := beego.AppConfig.DefaultBool("cache", false)
	if !isOpenCache {
		cache.Init(&cache.NullCache{})
	}
	beego.Info("正常初始化缓存配置.")
	cacheProvider := beego.AppConfig.String("cache_provider")
	if cacheProvider == "file" {
		cacheFilePath := beego.AppConfig.DefaultString("cache_file_path", "./runtime/cache/")
		if strings.HasPrefix(cacheFilePath, "./") {
			cacheFilePath = filepath.Join(conf.WorkingDirectory, string(cacheFilePath[1:]))
		}
		fileCache := beegoCache.NewFileCache()

		fileConfig := make(map[string]string, 0)

		fileConfig["CachePath"] = cacheFilePath
		fileConfig["DirectoryLevel"] = beego.AppConfig.DefaultString("cache_file_dir_level", "2")
		fileConfig["EmbedExpiry"] = beego.AppConfig.DefaultString("cache_file_expiry", "120")
		fileConfig["FileSuffix"] = beego.AppConfig.DefaultString("cache_file_suffix", ".bin")

		bc, err := json.Marshal(&fileConfig)
		if err != nil {
			beego.Error("初始化Redis缓存失败:", err)
			os.Exit(1)
		}

		fileCache.StartAndGC(string(bc))

		cache.Init(fileCache)

	} else if cacheProvider == "memory" {
		cacheInterval := beego.AppConfig.DefaultInt("cache_memory_interval", 60)
		memory := beegoCache.NewMemoryCache()
		beegoCache.DefaultEvery = cacheInterval
		cache.Init(memory)
	} else if cacheProvider == "redis" {
		var redisConfig struct {
			Conn     string `json:"conn"`
			Password string `json:"password"`
			DbNum    int    `json:"dbNum"`
		}
		redisConfig.DbNum = 0
		redisConfig.Conn = beego.AppConfig.DefaultString("cache_redis_host", "")
		if pwd := beego.AppConfig.DefaultString("cache_redis_password", ""); pwd != "" {
			redisConfig.Password = pwd
		}
		if dbNum := beego.AppConfig.DefaultInt("cache_redis_db", 0); dbNum > 0 {
			redisConfig.DbNum = dbNum
		}

		bc, err := json.Marshal(&redisConfig)
		if err != nil {
			beego.Error("初始化Redis缓存失败:", err)
			os.Exit(1)
		}
		redisCache, err := beegoCache.NewCache("redis", string(bc))

		if err != nil {
			beego.Error("初始化Redis缓存失败:", err)
			os.Exit(1)
		}

		cache.Init(redisCache)
	} else if cacheProvider == "memcache" {

		var memcacheConfig struct {
			Conn string `json:"conn"`
		}
		memcacheConfig.Conn = beego.AppConfig.DefaultString("cache_memcache_host", "")

		bc, err := json.Marshal(&memcacheConfig)
		if err != nil {
			beego.Error("初始化Redis缓存失败:", err)
			os.Exit(1)
		}
		memcache, err := beegoCache.NewCache("memcache", string(bc))

		if err != nil {
			beego.Error("初始化Memcache缓存失败:", err)
			os.Exit(1)
		}

		cache.Init(memcache)

	} else {
		cache.Init(&cache.NullCache{})
		beego.Warn("不支持的缓存管道,缓存将禁用.")
		return
	}
	beego.Info("缓存初始化完成.")
}

func init() {

	if configPath, err := filepath.Abs(conf.ConfigurationFile); err == nil {
		conf.ConfigurationFile = configPath
	}
	gocaptcha.ReadFonts("./static/fonts", ".ttf")
	gob.Register(models.Member{})

	if p, err := filepath.Abs(os.Args[0]); err == nil {
		conf.WorkingDirectory = filepath.Dir(p)
	}
}
<commit_msg>Update command.go<commit_after>package commands

import (
	"encoding/gob"
	"flag"
	"fmt"
	"log"
	"net/url"
	"os"
	"path/filepath"
	"strings"
	"time"

	"encoding/json"

	"github.com/astaxie/beego"
	beegoCache "github.com/astaxie/beego/cache"
	_ "github.com/astaxie/beego/cache/memcache"
	_ "github.com/astaxie/beego/cache/redis"
	"github.com/astaxie/beego/logs"
	"github.com/astaxie/beego/orm"
	"github.com/lifei6671/gocaptcha"
	"github.com/lifei6671/mindoc/cache"
	"github.com/lifei6671/mindoc/commands/migrate"
	"github.com/lifei6671/mindoc/conf"
	"github.com/lifei6671/mindoc/models"
	"github.com/lifei6671/mindoc/utils/filetil"
)

// RegisterDataBase 注册数据库
func RegisterDataBase() {
	beego.Info("正在初始化数据库配置.")
	adapter := beego.AppConfig.String("db_adapter")

	if adapter == "mysql" {
		host := beego.AppConfig.String("db_host")
		database := beego.AppConfig.String("db_database")
		username := beego.AppConfig.String("db_username")
		password := beego.AppConfig.String("db_password")

		timezone := beego.AppConfig.String("timezone")
		location, err := time.LoadLocation(timezone)
		if err == nil {
			orm.DefaultTimeLoc = location
		} else {
			beego.Error("加载时区配置信息失败,请检查是否存在ZONEINFO环境变量:", err)
		}

		port := beego.AppConfig.String("db_port")

		dataSource := fmt.Sprintf("%s:%s@tcp(%s:%s)/%s?charset=utf8mb4&parseTime=true&loc=%s", username, password, host, port, database, url.QueryEscape(timezone))

		if err := orm.RegisterDataBase("default", "mysql", dataSource); err != nil {
			beego.Error("注册默认数据库失败:", err)
			os.Exit(1)
		}
	} else if adapter == "sqlite3" {
		orm.DefaultTimeLoc = time.UTC
		database := beego.AppConfig.String("db_database")
		if strings.HasPrefix(database, "./") {
			database = filepath.Join(conf.WorkingDirectory, string(database[1:]))
		}

		dbPath := filepath.Dir(database)
		os.MkdirAll(dbPath, 0777)

		err := orm.RegisterDataBase("default", "sqlite3", database)

		if err != nil {
			beego.Error("注册默认数据库失败:", err)
		}
	} else {
		beego.Error("不支持的数据库类型.")
		os.Exit(1)
	}
	beego.Info("数据库初始化完成.")
}

// RegisterModel 注册Model
func RegisterModel() {
	orm.RegisterModelWithPrefix(conf.GetDatabasePrefix(),
		new(models.Member),
		new(models.Book),
		new(models.Relationship),
		new(models.Option),
		new(models.Document),
		new(models.Attachment),
		new(models.Logger),
		new(models.MemberToken),
		new(models.DocumentHistory),
		new(models.Migration),
		new(models.Label),
	)
	//migrate.RegisterMigration()
}

// RegisterLogger 注册日志
func RegisterLogger(log string) {

	logs.SetLogFuncCall(true)
	logs.SetLogger("console")
	logs.EnableFuncCallDepth(true)
	logs.Async()

	logPath := filepath.Join(log, "log.log")

	if _, err := os.Stat(logPath); os.IsNotExist(err) {

		os.MkdirAll(log, 0777)

		if f, err := os.Create(logPath); err == nil {
			f.Close()
			config := make(map[string]interface{}, 1)

			config["filename"] = logPath

			b, _ := json.Marshal(config)

			beego.SetLogger("file", string(b))
		}
	}

	beego.SetLogFuncCall(true)
	beego.BeeLogger.Async()
}

// RunCommand 注册orm命令行工具
func RegisterCommand() {

	if len(os.Args) >= 2 && os.Args[1] == "install" {
		ResolveCommand(os.Args[2:])
		Install()
	} else if len(os.Args) >= 2 && os.Args[1] == "version" {
		CheckUpdate()
		os.Exit(0)
	} else if len(os.Args) >= 2 && os.Args[1] == "migrate" {
		ResolveCommand(os.Args[2:])
		migrate.RunMigration()
	}
}

//注册模板函数
func RegisterFunction() {
	beego.AddFuncMap("config", models.GetOptionValue)

	beego.AddFuncMap("cdn", func(p string) string {
		cdn := beego.AppConfig.DefaultString("cdn", "")
		if strings.HasPrefix(p, "http://") || strings.HasPrefix(p, "https://") {
			return p
		}
		//如果没有设置cdn，则使用baseURL拼接
		if cdn == "" {
			baseUrl := beego.AppConfig.DefaultString("baseurl", "")

			if strings.HasPrefix(p, "/") && strings.HasSuffix(baseUrl, "/") {
				return baseUrl + p[1:]
			}
			if !strings.HasPrefix(p, "/") && !strings.HasSuffix(baseUrl, "/") {
				return baseUrl + "/" + p
			}
			return baseUrl + p
		}
		if strings.HasPrefix(p, "/") && strings.HasSuffix(cdn, "/") {
			return cdn + string(p[1:])
		}
		if !strings.HasPrefix(p, "/") && !strings.HasSuffix(cdn, "/") {
			return cdn + "/" + p
		}
		return cdn + p
	})

	beego.AddFuncMap("cdnjs", conf.URLForWithCdnJs)
	beego.AddFuncMap("cdncss", conf.URLForWithCdnCss)
	beego.AddFuncMap("cdnimg", conf.URLForWithCdnImage)
	//重写url生成，支持配置域名以及域名前缀
	beego.AddFuncMap("urlfor", conf.URLFor)
	beego.AddFuncMap("date_format", func(t time.Time, format string) string {
		return t.Local().Format(format)
	})
}

//解析命令
func ResolveCommand(args []string) {
	flagSet := flag.NewFlagSet("MinDoc command: ", flag.ExitOnError)
	flagSet.StringVar(&conf.ConfigurationFile, "config", "", "MinDoc configuration file.")
	flagSet.StringVar(&conf.WorkingDirectory, "dir", "", "MinDoc working directory.")
	flagSet.StringVar(&conf.LogFile, "log", "", "MinDoc log file path.")

	flagSet.Parse(args)

	if conf.WorkingDirectory == "" {
		if p, err := filepath.Abs(os.Args[0]); err == nil {
			conf.WorkingDirectory = filepath.Dir(p)
		}
	}
	if conf.LogFile == "" {
		conf.LogFile = filepath.Join(conf.WorkingDirectory, "logs")
	}
	if conf.ConfigurationFile == "" {
		conf.ConfigurationFile = filepath.Join(conf.WorkingDirectory, "conf", "app.conf")
		config := filepath.Join(conf.WorkingDirectory, "conf", "app.conf.example")
		if !filetil.FileExists(conf.ConfigurationFile) && filetil.FileExists(config) {
			filetil.CopyFile(conf.ConfigurationFile, config)
		}
	}
	gocaptcha.ReadFonts(filepath.Join(conf.WorkingDirectory, "static", "fonts"), ".ttf")

	err := beego.LoadAppConfig("ini", conf.ConfigurationFile)

	if err != nil {
		log.Println("An error occurred:", err)
		os.Exit(1)
	}
	uploads := filepath.Join(conf.WorkingDirectory, "uploads")

	os.MkdirAll(uploads, 0666)

	beego.BConfig.WebConfig.StaticDir["/static"] = filepath.Join(conf.WorkingDirectory, "static")
	beego.BConfig.WebConfig.StaticDir["/uploads"] = uploads
	beego.BConfig.WebConfig.ViewsPath = filepath.Join(conf.WorkingDirectory, "views")

	fonts := filepath.Join(conf.WorkingDirectory, "static", "fonts")

	if !filetil.FileExists(fonts) {
		log.Fatal("Font path not exist.")
	}
	gocaptcha.ReadFonts(filepath.Join(conf.WorkingDirectory, "static", "fonts"), ".ttf")

	RegisterDataBase()
	RegisterCache()
	RegisterModel()
	RegisterLogger(conf.LogFile)
}

//注册缓存管道
func RegisterCache() {
	isOpenCache := beego.AppConfig.DefaultBool("cache", false)
	if !isOpenCache {
		cache.Init(&cache.NullCache{})
	}
	beego.Info("正常初始化缓存配置.")
	cacheProvider := beego.AppConfig.String("cache_provider")
	if cacheProvider == "file" {
		cacheFilePath := beego.AppConfig.DefaultString("cache_file_path", "./runtime/cache/")
		if strings.HasPrefix(cacheFilePath, "./") {
			cacheFilePath = filepath.Join(conf.WorkingDirectory, string(cacheFilePath[1:]))
		}
		fileCache := beegoCache.NewFileCache()

		fileConfig := make(map[string]string, 0)

		fileConfig["CachePath"] = cacheFilePath
		fileConfig["DirectoryLevel"] = beego.AppConfig.DefaultString("cache_file_dir_level", "2")
		fileConfig["EmbedExpiry"] = beego.AppConfig.DefaultString("cache_file_expiry", "120")
		fileConfig["FileSuffix"] = beego.AppConfig.DefaultString("cache_file_suffix", ".bin")

		bc, err := json.Marshal(&fileConfig)
		if err != nil {
			beego.Error("初始化file缓存失败:", err)
			os.Exit(1)
		}

		fileCache.StartAndGC(string(bc))

		cache.Init(fileCache)

	} else if cacheProvider == "memory" {
		cacheInterval := beego.AppConfig.DefaultInt("cache_memory_interval", 60)
		memory := beegoCache.NewMemoryCache()
		beegoCache.DefaultEvery = cacheInterval
		cache.Init(memory)
	} else if cacheProvider == "redis" {
		var redisConfig struct {
			Conn     string `json:"conn"`
			Password string `json:"password"`
			DbNum    int    `json:"dbNum"`
		}
		redisConfig.DbNum = 0
		redisConfig.Conn = beego.AppConfig.DefaultString("cache_redis_host", "")
		if pwd := beego.AppConfig.DefaultString("cache_redis_password", ""); pwd != "" {
			redisConfig.Password = pwd
		}
		if dbNum := beego.AppConfig.DefaultInt("cache_redis_db", 0); dbNum > 0 {
			redisConfig.DbNum = dbNum
		}

		bc, err := json.Marshal(&redisConfig)
		if err != nil {
			beego.Error("初始化Redis缓存失败:", err)
			os.Exit(1)
		}
		redisCache, err := beegoCache.NewCache("redis", string(bc))

		if err != nil {
			beego.Error("初始化Redis缓存失败:", err)
			os.Exit(1)
		}

		cache.Init(redisCache)
	} else if cacheProvider == "memcache" {

		var memcacheConfig struct {
			Conn string `json:"conn"`
		}
		memcacheConfig.Conn = beego.AppConfig.DefaultString("cache_memcache_host", "")

		bc, err := json.Marshal(&memcacheConfig)
		if err != nil {
			beego.Error("初始化Redis缓存失败:", err)
			os.Exit(1)
		}
		memcache, err := beegoCache.NewCache("memcache", string(bc))

		if err != nil {
			beego.Error("初始化Memcache缓存失败:", err)
			os.Exit(1)
		}

		cache.Init(memcache)

	} else {
		cache.Init(&cache.NullCache{})
		beego.Warn("不支持的缓存管道,缓存将禁用.")
		return
	}
	beego.Info("缓存初始化完成.")
}

func init() {

	if configPath, err := filepath.Abs(conf.ConfigurationFile); err == nil {
		conf.ConfigurationFile = configPath
	}
	gocaptcha.ReadFonts("./static/fonts", ".ttf")
	gob.Register(models.Member{})

	if p, err := filepath.Abs(os.Args[0]); err == nil {
		conf.WorkingDirectory = filepath.Dir(p)
	}
}
<|endoftext|>


<commit_before>package socks

import (
	"net"

	"github.com/v2ray/v2ray-core/common/alloc"
	"github.com/v2ray/v2ray-core/common/log"
	v2net "github.com/v2ray/v2ray-core/common/net"
	"github.com/v2ray/v2ray-core/proxy/socks/protocol"
)

const (
	bufferSize = 2 * 1024
)

var udpAddress v2net.Address

func (server *SocksServer) ListenUDP(port uint16) error {
	addr := &net.UDPAddr{
		IP:   net.IP{0, 0, 0, 0},
		Port: int(port),
		Zone: "",
	}
	conn, err := net.ListenUDP("udp", addr)
	if err != nil {
		log.Error("Socks failed to listen UDP on port %d: %v", port, err)
		return err
	}
	// TODO: make this configurable
	udpAddress = v2net.IPAddress([]byte{127, 0, 0, 1}, port)

	go server.AcceptPackets(conn)
	return nil
}

func (server *SocksServer) getUDPAddr() v2net.Address {
	return udpAddress
}

func (server *SocksServer) AcceptPackets(conn *net.UDPConn) error {
	for {
		buffer := alloc.NewBuffer()
		nBytes, addr, err := conn.ReadFromUDP(buffer.Value)
		if err != nil {
			log.Error("Socks failed to read UDP packets: %v", err)
			buffer.Release()
			continue
		}
		log.Info("Client UDP connection from %v", addr)
		request, err := protocol.ReadUDPRequest(buffer.Value[:nBytes])
		buffer.Release()
		if err != nil {
			log.Error("Socks failed to parse UDP request: %v", err)
			request.Data.Release()
			continue
		}
		if request.Fragment != 0 {
			log.Warning("Dropping fragmented UDP packets.")
			// TODO handle fragments
			request.Data.Release()
			continue
		}

		udpPacket := v2net.NewPacket(request.Destination(), request.Data, false)
		log.Info("Send packet to %s with %d bytes", udpPacket.Destination().String(), request.Data.Len())
		go server.handlePacket(conn, udpPacket, addr, request.Address)
	}
}

func (server *SocksServer) handlePacket(conn *net.UDPConn, packet v2net.Packet, clientAddr *net.UDPAddr, targetAddr v2net.Address) {
	ray := server.vPoint.DispatchToOutbound(packet)
	close(ray.InboundInput())

	if data, ok := <-ray.InboundOutput(); ok {
		response := &protocol.Socks5UDPRequest{
			Fragment: 0,
			Address:  targetAddr,
			Data:     data,
		}
		log.Info("Writing back UDP response with %d bytes from %s to %s", data.Len(), targetAddr.String(), clientAddr.String())

		udpMessage := alloc.NewSmallBuffer().Clear()
		response.Write(udpMessage)

		nBytes, err := conn.WriteToUDP(udpMessage.Value, clientAddr)
		udpMessage.Release()
		response.Data.Release()
		if err != nil {
			log.Error("Socks failed to write UDP message (%d bytes) to %s: %v", nBytes, clientAddr.String(), err)
		}
	}
}
<commit_msg>Remove unused const<commit_after>package socks

import (
	"net"

	"github.com/v2ray/v2ray-core/common/alloc"
	"github.com/v2ray/v2ray-core/common/log"
	v2net "github.com/v2ray/v2ray-core/common/net"
	"github.com/v2ray/v2ray-core/proxy/socks/protocol"
)

var udpAddress v2net.Address

func (server *SocksServer) ListenUDP(port uint16) error {
	addr := &net.UDPAddr{
		IP:   net.IP{0, 0, 0, 0},
		Port: int(port),
		Zone: "",
	}
	conn, err := net.ListenUDP("udp", addr)
	if err != nil {
		log.Error("Socks failed to listen UDP on port %d: %v", port, err)
		return err
	}
	// TODO: make this configurable
	udpAddress = v2net.IPAddress([]byte{127, 0, 0, 1}, port)

	go server.AcceptPackets(conn)
	return nil
}

func (server *SocksServer) getUDPAddr() v2net.Address {
	return udpAddress
}

func (server *SocksServer) AcceptPackets(conn *net.UDPConn) error {
	for {
		buffer := alloc.NewBuffer()
		nBytes, addr, err := conn.ReadFromUDP(buffer.Value)
		if err != nil {
			log.Error("Socks failed to read UDP packets: %v", err)
			buffer.Release()
			continue
		}
		log.Info("Client UDP connection from %v", addr)
		request, err := protocol.ReadUDPRequest(buffer.Value[:nBytes])
		buffer.Release()
		if err != nil {
			log.Error("Socks failed to parse UDP request: %v", err)
			request.Data.Release()
			continue
		}
		if request.Fragment != 0 {
			log.Warning("Dropping fragmented UDP packets.")
			// TODO handle fragments
			request.Data.Release()
			continue
		}

		udpPacket := v2net.NewPacket(request.Destination(), request.Data, false)
		log.Info("Send packet to %s with %d bytes", udpPacket.Destination().String(), request.Data.Len())
		go server.handlePacket(conn, udpPacket, addr, request.Address)
	}
}

func (server *SocksServer) handlePacket(conn *net.UDPConn, packet v2net.Packet, clientAddr *net.UDPAddr, targetAddr v2net.Address) {
	ray := server.vPoint.DispatchToOutbound(packet)
	close(ray.InboundInput())

	if data, ok := <-ray.InboundOutput(); ok {
		response := &protocol.Socks5UDPRequest{
			Fragment: 0,
			Address:  targetAddr,
			Data:     data,
		}
		log.Info("Writing back UDP response with %d bytes from %s to %s", data.Len(), targetAddr.String(), clientAddr.String())

		udpMessage := alloc.NewSmallBuffer().Clear()
		response.Write(udpMessage)

		nBytes, err := conn.WriteToUDP(udpMessage.Value, clientAddr)
		udpMessage.Release()
		response.Data.Release()
		if err != nil {
			log.Error("Socks failed to write UDP message (%d bytes) to %s: %v", nBytes, clientAddr.String(), err)
		}
	}
}
<|endoftext|>


<commit_before>//+build !travis

package pir

import (
	"fmt"
	"strings"
	"unsafe"

	"github.com/go-gl/cl/v1.2/cl"
	"github.com/privacylab/talek/common"
)

// ShardCL represents a read-only shard of the database,
// backed by an OpenCL implementation of PIR
type ShardCL struct {
	log        *common.Logger
	name       string
	context    *ContextCL
	bucketSize int
	numBuckets int
	data       []byte
	numThreads int
	clData     cl.Mem
}

// NewShardCL creates a new OpenCL-backed shard
// The data is represented as a flat byte array = append(bucket_1, bucket_2 ... bucket_n)
// Pre-conditions:
// - len(data) must be a multiple of bucketSize
// Returns: the shard, or an error if mismatched size
func NewShardCL(name string, context *ContextCL, bucketSize int, data []byte, numThreads int) (*ShardCL, error) {
	s := &ShardCL{}
	s.log = common.NewLogger(name)
	s.name = name
	s.context = context

	// GetNumBuckets will compute the number of buckets stored in the Shard
	// If len(s.data) is not cleanly divisible by s.bucketSize,
	// returns an error
	if len(data)%bucketSize != 0 {
		return nil, fmt.Errorf("NewShardCL(%v) failed: data(len=%v) not multiple of bucketSize=%v", name, len(data), bucketSize)
	}

	s.bucketSize = bucketSize
	s.numBuckets = (len(data) / bucketSize)
	s.data = data
	s.numThreads = numThreads

	/** OpenCL **/
	//  Create buffers
	var errptr *cl.ErrorCode
	s.clData = cl.CreateBuffer(s.context.Context, cl.MEM_READ_ONLY, uint64(len(data)), nil, errptr)
	if errptr != nil && cl.ErrorCode(*errptr) != cl.SUCCESS {
		return nil, fmt.Errorf("NewShardCL(%v) failed: couldnt create OpenCL buffer", name)
	}
	//Write shard data to GPU
	err := cl.EnqueueWriteBuffer(s.context.CommandQueue, s.clData, cl.TRUE, 0, uint64(len(data)), unsafe.Pointer(&data[0]), 0, nil, nil)
	if err != cl.SUCCESS {
		return nil, fmt.Errorf("NewShardCL(%v) failed: cannot write shard to GPU (OpenCL buffer)", name)
	}

	return s, nil
}

/*********************************************
 * PUBLIC METHODS
 *********************************************/

// Free releases all OpenCL buffers
func (s *ShardCL) Free() error {
	errStr := ""
	err := cl.ReleaseMemObject(s.clData)
	if err != cl.SUCCESS {
		errStr += cl.ErrToStr(err) + "\n"
	}
	if strings.Compare(errStr, "") != 0 {
		return fmt.Errorf("ContextCL.Free errors: " + errStr)
	}
	return nil
}

// GetName returns the name of the shard
func (s *ShardCL) GetName() string {
	return s.name
}

// GetBucketSize returns the size (in bytes) of a bucket
func (s *ShardCL) GetBucketSize() int {
	return s.bucketSize
}

// GetNumBuckets returns the number of buckets in the shard
func (s *ShardCL) GetNumBuckets() int {
	return s.numBuckets
}

// GetData returns a slice of the data
func (s *ShardCL) GetData() []byte {
	return s.data[:]
}

// Read handles a batch read, where each request is concatentated into `reqs`
//   each request consists of `reqLength` bytes
//   Note: every request starts on a byte boundary
// Returns: a single byte array where responses are concatenated by the order in `reqs`
//   each response consists of `s.bucketSize` bytes
func (s *ShardCL) Read(reqs []byte, reqLength int) ([]byte, error) {
	if len(reqs)%reqLength != 0 {
		return nil, fmt.Errorf("ShardCL.Read expects len(reqs)=%d to be a multiple of reqLength=%d", len(reqs), reqLength)
	}

	inputSize := len(reqs)
	batchSize := inputSize / reqLength
	outputSize := batchSize * s.bucketSize
	responses := make([]byte, outputSize)
	context := s.context.Context
	var err cl.ErrorCode
	var errptr *cl.ErrorCode

	//Create buffers
	input := cl.CreateBuffer(context, cl.MEM_READ_ONLY, uint64(inputSize), nil, errptr)
	if errptr != nil && cl.ErrorCode(*errptr) != cl.SUCCESS {
		return nil, fmt.Errorf("couldnt create input buffer")
	}
	defer cl.ReleaseMemObject(input)

	output := cl.CreateBuffer(context, cl.MEM_WRITE_ONLY, uint64(outputSize), nil, errptr)
	if errptr != nil && cl.ErrorCode(*errptr) != cl.SUCCESS {
		return nil, fmt.Errorf("couldnt create output buffer")
	}
	defer cl.ReleaseMemObject(output)

	//Write request data
	err = cl.EnqueueWriteBuffer(s.context.CommandQueue, input, cl.TRUE, 0, uint64(inputSize), unsafe.Pointer(&reqs[0]), 0, nil, nil)
	if err != cl.SUCCESS {
		return nil, fmt.Errorf("Failed to write to input requests (OpenCL buffer)")
	}

	/** START LOCK REGION **/
	s.context.KernelMutex.Lock()

	// Note: SetKernelArgs->EnqueueNDRangeKernel is not thread-safe
	//   @todo - create multiple kernels to support parallel PIR in a single context
	//   https://www.khronos.org/registry/OpenCL/sdk/1.2/docs/man/xhtml/clSetKernelArg.html
	//Set kernel args
	data := s.clData
	batchSize32 := uint32(batchSize)
	reqLength32 := uint32(reqLength)
	numBuckets32 := uint32(s.numBuckets)
	bucketSize32 := uint32(s.bucketSize / KernelDataSize)
	//global := local
	local := uint64(s.context.GetGroupSize())
	global := uint64(s.numThreads)
	if global < local {
		local = global
	}
	global32 := uint32(global)
	scratchSize32 := uint32(GPUScratchSize / KernelDataSize)
	argSizes := []uint64{8, 8, 8, GPUScratchSize, 4, 4, 4, 4, 4, 4}
	args := []unsafe.Pointer{
		unsafe.Pointer(&data),
		unsafe.Pointer(&input),
		unsafe.Pointer(&output),
		nil,
		unsafe.Pointer(&batchSize32),
		unsafe.Pointer(&reqLength32),
		unsafe.Pointer(&numBuckets32),
		unsafe.Pointer(&bucketSize32),
		unsafe.Pointer(&global32),
		unsafe.Pointer(&scratchSize32),
	}

	for i := 0; i < len(args); i++ {
		err = cl.SetKernelArg(s.context.Kernel, uint32(i), argSizes[i], args[i])
		if err != cl.SUCCESS {
			return nil, fmt.Errorf("Failed to write kernel arg %v", i)
		}
	}

	//s.log.Info.Printf("local=%v, global=%v\n", local, global)
	err = cl.EnqueueNDRangeKernel(s.context.CommandQueue, s.context.Kernel, 1, nil, &global, &local, 0, nil, nil)
	if err != cl.SUCCESS {
		return nil, fmt.Errorf("Failed to execute kernel")
	}
	cl.Finish(s.context.CommandQueue) //@todo inside or outside lock region?
	s.context.KernelMutex.Unlock()
	/** END LOCK REGION **/

	err = cl.EnqueueReadBuffer(s.context.CommandQueue, output, cl.TRUE, 0, uint64(outputSize), unsafe.Pointer(&responses[0]), 0, nil, nil)
	if err != cl.SUCCESS {
		return nil, fmt.Errorf("Failed to read output response (OpenCL buffer), err=%v", cl.ErrToStr(err))
	}

	return responses, nil
}

/*********************************************
 * PRIVATE METHODS
 *********************************************/
<commit_msg>reduce lock region<commit_after>//+build !travis

package pir

import (
	"fmt"
	"strings"
	"unsafe"

	"github.com/go-gl/cl/v1.2/cl"
	"github.com/privacylab/talek/common"
)

// ShardCL represents a read-only shard of the database,
// backed by an OpenCL implementation of PIR
type ShardCL struct {
	log        *common.Logger
	name       string
	context    *ContextCL
	bucketSize int
	numBuckets int
	data       []byte
	numThreads int
	clData     cl.Mem
}

// NewShardCL creates a new OpenCL-backed shard
// The data is represented as a flat byte array = append(bucket_1, bucket_2 ... bucket_n)
// Pre-conditions:
// - len(data) must be a multiple of bucketSize
// Returns: the shard, or an error if mismatched size
func NewShardCL(name string, context *ContextCL, bucketSize int, data []byte, numThreads int) (*ShardCL, error) {
	s := &ShardCL{}
	s.log = common.NewLogger(name)
	s.name = name
	s.context = context

	// GetNumBuckets will compute the number of buckets stored in the Shard
	// If len(s.data) is not cleanly divisible by s.bucketSize,
	// returns an error
	if len(data)%bucketSize != 0 {
		return nil, fmt.Errorf("NewShardCL(%v) failed: data(len=%v) not multiple of bucketSize=%v", name, len(data), bucketSize)
	}

	s.bucketSize = bucketSize
	s.numBuckets = (len(data) / bucketSize)
	s.data = data
	s.numThreads = numThreads

	/** OpenCL **/
	//  Create buffers
	var errptr *cl.ErrorCode
	s.clData = cl.CreateBuffer(s.context.Context, cl.MEM_READ_ONLY, uint64(len(data)), nil, errptr)
	if errptr != nil && cl.ErrorCode(*errptr) != cl.SUCCESS {
		return nil, fmt.Errorf("NewShardCL(%v) failed: couldnt create OpenCL buffer", name)
	}
	//Write shard data to GPU
	err := cl.EnqueueWriteBuffer(s.context.CommandQueue, s.clData, cl.TRUE, 0, uint64(len(data)), unsafe.Pointer(&data[0]), 0, nil, nil)
	if err != cl.SUCCESS {
		return nil, fmt.Errorf("NewShardCL(%v) failed: cannot write shard to GPU (OpenCL buffer)", name)
	}

	return s, nil
}

/*********************************************
 * PUBLIC METHODS
 *********************************************/

// Free releases all OpenCL buffers
func (s *ShardCL) Free() error {
	errStr := ""
	err := cl.ReleaseMemObject(s.clData)
	if err != cl.SUCCESS {
		errStr += cl.ErrToStr(err) + "\n"
	}
	if strings.Compare(errStr, "") != 0 {
		return fmt.Errorf("ContextCL.Free errors: " + errStr)
	}
	return nil
}

// GetName returns the name of the shard
func (s *ShardCL) GetName() string {
	return s.name
}

// GetBucketSize returns the size (in bytes) of a bucket
func (s *ShardCL) GetBucketSize() int {
	return s.bucketSize
}

// GetNumBuckets returns the number of buckets in the shard
func (s *ShardCL) GetNumBuckets() int {
	return s.numBuckets
}

// GetData returns a slice of the data
func (s *ShardCL) GetData() []byte {
	return s.data[:]
}

// Read handles a batch read, where each request is concatentated into `reqs`
//   each request consists of `reqLength` bytes
//   Note: every request starts on a byte boundary
// Returns: a single byte array where responses are concatenated by the order in `reqs`
//   each response consists of `s.bucketSize` bytes
func (s *ShardCL) Read(reqs []byte, reqLength int) ([]byte, error) {
	if len(reqs)%reqLength != 0 {
		return nil, fmt.Errorf("ShardCL.Read expects len(reqs)=%d to be a multiple of reqLength=%d", len(reqs), reqLength)
	}

	inputSize := len(reqs)
	batchSize := inputSize / reqLength
	outputSize := batchSize * s.bucketSize
	responses := make([]byte, outputSize)
	context := s.context.Context
	var err cl.ErrorCode
	var errptr *cl.ErrorCode

	//Create buffers
	input := cl.CreateBuffer(context, cl.MEM_READ_ONLY, uint64(inputSize), nil, errptr)
	if errptr != nil && cl.ErrorCode(*errptr) != cl.SUCCESS {
		return nil, fmt.Errorf("couldnt create input buffer")
	}
	defer cl.ReleaseMemObject(input)

	output := cl.CreateBuffer(context, cl.MEM_WRITE_ONLY, uint64(outputSize), nil, errptr)
	if errptr != nil && cl.ErrorCode(*errptr) != cl.SUCCESS {
		return nil, fmt.Errorf("couldnt create output buffer")
	}
	defer cl.ReleaseMemObject(output)

	//Write request data
	err = cl.EnqueueWriteBuffer(s.context.CommandQueue, input, cl.TRUE, 0, uint64(inputSize), unsafe.Pointer(&reqs[0]), 0, nil, nil)
	if err != cl.SUCCESS {
		return nil, fmt.Errorf("Failed to write to input requests (OpenCL buffer)")
	}

	//Set kernel args
	data := s.clData
	batchSize32 := uint32(batchSize)
	reqLength32 := uint32(reqLength)
	numBuckets32 := uint32(s.numBuckets)
	bucketSize32 := uint32(s.bucketSize / KernelDataSize)
	//global := local
	local := uint64(s.context.GetGroupSize())
	global := uint64(s.numThreads)
	if global < local {
		local = global
	}
	global32 := uint32(global)
	scratchSize32 := uint32(GPUScratchSize / KernelDataSize)
	argSizes := []uint64{8, 8, 8, GPUScratchSize, 4, 4, 4, 4, 4, 4}
	args := []unsafe.Pointer{
		unsafe.Pointer(&data),
		unsafe.Pointer(&input),
		unsafe.Pointer(&output),
		nil,
		unsafe.Pointer(&batchSize32),
		unsafe.Pointer(&reqLength32),
		unsafe.Pointer(&numBuckets32),
		unsafe.Pointer(&bucketSize32),
		unsafe.Pointer(&global32),
		unsafe.Pointer(&scratchSize32),
	}

	/** START LOCK REGION **/
	s.context.KernelMutex.Lock()
	// Note: SetKernelArgs->EnqueueNDRangeKernel is not thread-safe
	//   @todo - create multiple kernels to support parallel PIR in a single context
	//   https://www.khronos.org/registry/OpenCL/sdk/1.2/docs/man/xhtml/clSetKernelArg.html

	for i := 0; i < len(args); i++ {
		err = cl.SetKernelArg(s.context.Kernel, uint32(i), argSizes[i], args[i])
		if err != cl.SUCCESS {
			return nil, fmt.Errorf("Failed to write kernel arg %v", i)
		}
	}

	//s.log.Info.Printf("local=%v, global=%v\n", local, global)
	err = cl.EnqueueNDRangeKernel(s.context.CommandQueue, s.context.Kernel, 1, nil, &global, &local, 0, nil, nil)
	if err != cl.SUCCESS {
		return nil, fmt.Errorf("Failed to execute kernel")
	}
	cl.Finish(s.context.CommandQueue) //@todo inside or outside lock region?
	s.context.KernelMutex.Unlock()
	/** END LOCK REGION **/

	err = cl.EnqueueReadBuffer(s.context.CommandQueue, output, cl.TRUE, 0, uint64(outputSize), unsafe.Pointer(&responses[0]), 0, nil, nil)
	if err != cl.SUCCESS {
		return nil, fmt.Errorf("Failed to read output response (OpenCL buffer), err=%v", cl.ErrToStr(err))
	}

	return responses, nil
}

/*********************************************
 * PRIVATE METHODS
 *********************************************/
<|endoftext|>


<commit_before>package main

import (
	"encoding/json"
	"flag"
	"fmt"
	"io/ioutil"
	"log"
	"net/http"
	"os"
	"os/signal"
	"sync"
	"syscall"
	"time"

	app "github.com/microamp/pitchfork-dl/app"
)

var (
	maxPageWorkers   = 2  // Number of concurrent page workers
	maxReviewWorkers = 48 // Number of concurrent review workers

	retryDelayPage   = 5 * time.Second
	retryDelayReview = 3 * time.Second

	proxy, output       string
	pageFirst, pageLast int
)

func writeToFile(path string, review *app.Review) error {
	bytes, err := json.MarshalIndent(review, "", "  ")
	if err != nil {
		return err
	}
	return ioutil.WriteFile(path, bytes, 0644)
}

func startPageWorker(scraper *app.Scraper, chanPages <-chan int, chanReviewIDs chan<- string, chanDone chan<- bool, wg *sync.WaitGroup) {
	defer wg.Done()

	for {
		pageNumber, ok := <-chanPages
		if !ok {
			break // Termination due to channel closed
		}

		for {
			_, resp, err := scraper.ScrapePage(pageNumber)
			if err != nil {
				log.Fatalf("Error while scraping page, %d: %s", pageNumber, err)
			}
			if resp.StatusCode == http.StatusNotFound {
				log.Printf("Page %d not found", pageNumber)
				chanDone <- true
				continue
			}

			reviewIDs, err := app.ParsePage(pageNumber, resp)
			if err != nil {
				log.Fatalf("Error while parsing page, %d: %s", pageNumber, err)
			}

			if len(reviewIDs) == 0 {
				// Retry if no data
				log.Printf("Empty data received. Retrying page %d after %d seconds...", pageNumber, retryDelayPage/time.Second)
				time.Sleep(retryDelayPage)
				continue
			} else {
				// Log page info
				go log.Printf("Page %d with %d reviews", pageNumber, len(reviewIDs))
			}

			for _, reviewID := range reviewIDs {
				chanReviewIDs <- reviewID
			}
			break
		}
	}
}

func startReviewWorker(scraper *app.Scraper, chanReviewIDs <-chan string, wg *sync.WaitGroup) {
	defer wg.Done()

	for {
		reviewID, ok := <-chanReviewIDs
		if !ok {
			break // Termination due to channel closed
		}

		for {
			_, resp, err := scraper.ScrapeReview(reviewID)
			if err != nil {
				log.Fatalf("Error while scraping review, %s: %s", reviewID, err)
			}
			if resp.StatusCode == http.StatusNotFound {
				log.Printf("Review not found: %s", reviewID)
				break
			}

			review, err := app.ParseReview(reviewID, resp)
			if err != nil {
				log.Fatalf("Error while parsing review, %s: %s", reviewID, err)
			}

			if len(review.Albums) == 0 {
				// Retry if no data
				log.Printf("Empty data received. Retrying review %s after %d seconds...", reviewID, retryDelayReview/time.Second)
				time.Sleep(retryDelayReview)
				continue
			} else {
				// Log review info
				go review.PrintInfo()
			}

			// Write to file (JSON)
			filename := fmt.Sprintf("%s/%s.json", scraper.OutputDirectory, reviewID)
			err = writeToFile(filename, review)
			if err != nil {
				log.Fatalf("Error while writing to file, %s: %s", filename, err)
			}
			break
		}
	}
}

func startProcessing(scraper *app.Scraper, first, last int) {
	// Channels
	chanSigs := make(chan os.Signal)
	signal.Notify(chanSigs, os.Interrupt)
	signal.Notify(chanSigs, syscall.SIGTERM)

	chanPages := make(chan int)
	chanReviewIDs := make(chan string)

	chanDone := make(chan bool, maxPageWorkers) // Buffered!

	// Wait groups
	var wgPages, wgReviewIDs sync.WaitGroup
	wgPages.Add(maxPageWorkers)
	wgReviewIDs.Add(maxReviewWorkers)

	// Start page workers
	go func() {
		for i := 0; i < maxPageWorkers; i++ {
			go startPageWorker(scraper, chanPages, chanReviewIDs, chanDone, &wgPages)
		}
	}()

	// Start review workers
	go func() {
		for i := 0; i < maxReviewWorkers; i++ {
			go startReviewWorker(scraper, chanReviewIDs, &wgReviewIDs)
		}
	}()

	// Detect OS interrupt signals
	go func() {
		for s := range chanSigs {
			log.Printf("Shutdown signal received (%s)", s)
			chanDone <- true
			break
		}
	}()

	pageNumber := first
	for {
		select {
		case <-chanDone:
			log.Println("Waiting for workers to finish...")
			close(chanPages)
			wgPages.Wait()
			close(chanReviewIDs)
			wgReviewIDs.Wait()
			os.Exit(1)
		default:
			// If last page is unknown, scrape until the end (i.e. 404)
			if last == 0 {
				chanPages <- pageNumber
				pageNumber++
			} else {
				if pageNumber <= last {
					chanPages <- pageNumber
					pageNumber++
				} else {
					log.Println("No more pages to download")
					chanDone <- true
				}
			}
		}
	}
}

func main() {
	// Parse params
	flags := flag.NewFlagSet("pitchfork-dl", flag.ExitOnError)
	flags.StringVar(&proxy, "proxy", "socks5://127.0.0.1:9150", "Proxy server")
	flags.StringVar(&output, "output", "reviews", "Output directory")
	flags.IntVar(&pageFirst, "first", 1, "First page")
	flags.IntVar(&pageLast, "last", 0, "Last page")

	err := flags.Parse(os.Args[1:])
	if err != nil {
		log.Fatalf("Error parsing flags: %v", err)
	}

	if pageLast == 0 {
		log.Printf("Scraping pages from %d...", pageFirst)
	} else {
		log.Printf("Scraping pages from %d to %d...", pageFirst, pageLast)
	}

	// Prepare proxy client
	proxyClient, err := app.GetProxyClient(proxy)
	if err != nil {
		log.Fatalf("Error: %v", err)
	}

	// Fire up goroutines and start processing
	scraper := &app.Scraper{
		Client:          proxyClient,
		OutputDirectory: output,
	}
	startProcessing(scraper, pageFirst, pageLast)
}
<commit_msg>Break out of loop when 404 is received<commit_after>package main

import (
	"encoding/json"
	"flag"
	"fmt"
	"io/ioutil"
	"log"
	"net/http"
	"os"
	"os/signal"
	"sync"
	"syscall"
	"time"

	app "github.com/microamp/pitchfork-dl/app"
)

var (
	maxPageWorkers   = 2  // Number of concurrent page workers
	maxReviewWorkers = 48 // Number of concurrent review workers

	retryDelayPage   = 5 * time.Second
	retryDelayReview = 3 * time.Second

	proxy, output       string
	pageFirst, pageLast int
)

func writeToFile(path string, review *app.Review) error {
	bytes, err := json.MarshalIndent(review, "", "  ")
	if err != nil {
		return err
	}
	return ioutil.WriteFile(path, bytes, 0644)
}

func startPageWorker(scraper *app.Scraper, chanPages <-chan int, chanReviewIDs chan<- string, chanDone chan<- bool, wg *sync.WaitGroup) {
	defer wg.Done()

	for {
		pageNumber, ok := <-chanPages
		if !ok {
			break // Termination due to channel closed
		}

		for {
			_, resp, err := scraper.ScrapePage(pageNumber)
			if err != nil {
				log.Fatalf("Error while scraping page, %d: %s", pageNumber, err)
			}
			if resp.StatusCode == http.StatusNotFound {
				log.Printf("Page %d not found", pageNumber)
				chanDone <- true
				break
			}

			reviewIDs, err := app.ParsePage(pageNumber, resp)
			if err != nil {
				log.Fatalf("Error while parsing page, %d: %s", pageNumber, err)
			}

			if len(reviewIDs) == 0 {
				// Retry if no data
				log.Printf("Empty data received. Retrying page %d after %d seconds...", pageNumber, retryDelayPage/time.Second)
				time.Sleep(retryDelayPage)
				continue
			} else {
				// Log page info
				go log.Printf("Page %d with %d reviews", pageNumber, len(reviewIDs))
			}

			for _, reviewID := range reviewIDs {
				chanReviewIDs <- reviewID
			}
			break
		}
	}
}

func startReviewWorker(scraper *app.Scraper, chanReviewIDs <-chan string, wg *sync.WaitGroup) {
	defer wg.Done()

	for {
		reviewID, ok := <-chanReviewIDs
		if !ok {
			break // Termination due to channel closed
		}

		for {
			_, resp, err := scraper.ScrapeReview(reviewID)
			if err != nil {
				log.Fatalf("Error while scraping review, %s: %s", reviewID, err)
			}
			if resp.StatusCode == http.StatusNotFound {
				log.Printf("Review not found: %s", reviewID)
				break
			}

			review, err := app.ParseReview(reviewID, resp)
			if err != nil {
				log.Fatalf("Error while parsing review, %s: %s", reviewID, err)
			}

			if len(review.Albums) == 0 {
				// Retry if no data
				log.Printf("Empty data received. Retrying review %s after %d seconds...", reviewID, retryDelayReview/time.Second)
				time.Sleep(retryDelayReview)
				continue
			} else {
				// Log review info
				go review.PrintInfo()
			}

			// Write to file (JSON)
			filename := fmt.Sprintf("%s/%s.json", scraper.OutputDirectory, reviewID)
			err = writeToFile(filename, review)
			if err != nil {
				log.Fatalf("Error while writing to file, %s: %s", filename, err)
			}
			break
		}
	}
}

func startProcessing(scraper *app.Scraper, first, last int) {
	// Channels
	chanSigs := make(chan os.Signal)
	signal.Notify(chanSigs, os.Interrupt)
	signal.Notify(chanSigs, syscall.SIGTERM)

	chanPages := make(chan int)
	chanReviewIDs := make(chan string)

	chanDone := make(chan bool, maxPageWorkers) // Buffered!

	// Wait groups
	var wgPages, wgReviewIDs sync.WaitGroup
	wgPages.Add(maxPageWorkers)
	wgReviewIDs.Add(maxReviewWorkers)

	// Start page workers
	go func() {
		for i := 0; i < maxPageWorkers; i++ {
			go startPageWorker(scraper, chanPages, chanReviewIDs, chanDone, &wgPages)
		}
	}()

	// Start review workers
	go func() {
		for i := 0; i < maxReviewWorkers; i++ {
			go startReviewWorker(scraper, chanReviewIDs, &wgReviewIDs)
		}
	}()

	// Detect OS interrupt signals
	go func() {
		for s := range chanSigs {
			log.Printf("Shutdown signal received (%s)", s)
			chanDone <- true
			break
		}
	}()

	pageNumber := first
	for {
		select {
		case <-chanDone:
			log.Println("Waiting for workers to finish...")
			close(chanPages)
			wgPages.Wait()
			close(chanReviewIDs)
			wgReviewIDs.Wait()
			os.Exit(1)
		default:
			// If last page is unknown, scrape until the end (i.e. 404)
			if last == 0 {
				chanPages <- pageNumber
				pageNumber++
			} else {
				if pageNumber <= last {
					chanPages <- pageNumber
					pageNumber++
				} else {
					log.Println("No more pages to download")
					chanDone <- true
				}
			}
		}
	}
}

func main() {
	// Parse params
	flags := flag.NewFlagSet("pitchfork-dl", flag.ExitOnError)
	flags.StringVar(&proxy, "proxy", "socks5://127.0.0.1:9150", "Proxy server")
	flags.StringVar(&output, "output", "reviews", "Output directory")
	flags.IntVar(&pageFirst, "first", 1, "First page")
	flags.IntVar(&pageLast, "last", 0, "Last page")

	err := flags.Parse(os.Args[1:])
	if err != nil {
		log.Fatalf("Error parsing flags: %v", err)
	}

	if pageLast == 0 {
		log.Printf("Scraping pages from %d...", pageFirst)
	} else {
		log.Printf("Scraping pages from %d to %d...", pageFirst, pageLast)
	}

	// Prepare proxy client
	proxyClient, err := app.GetProxyClient(proxy)
	if err != nil {
		log.Fatalf("Error: %v", err)
	}

	// Fire up goroutines and start processing
	scraper := &app.Scraper{
		Client:          proxyClient,
		OutputDirectory: output,
	}
	startProcessing(scraper, pageFirst, pageLast)
}
<|endoftext|>


<commit_before>//  Copyright 2016 The goscope Authors
//
//  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 testutil

import "github.com/zagrodzki/goscope/scope"

// Recorder is a buffer that implements the scope.DataRecorder interface.
// Users should call Wait() to get the data that was recorded in the buffer.
type BufferRecorder struct {
	tb     scope.Duration
	i      scope.Duration
	sweeps [][]scope.Voltage
	err    error
	done   chan struct{}
}

// NewBufferRecorder creates a new test data recorder with timebase equal to tb.
func NewBufferRecorder(tb scope.Duration) *BufferRecorder {
	return &BufferRecorder{
		tb: tb,
	}
}

// TimeBase returns the configured timebase (sweep length) of the recorder.
func (r *BufferRecorder) TimeBase() scope.Duration {
	return r.tb
}

// Reset prepares a new recording with sample interval i, reading samples from ch.
func (r *BufferRecorder) Reset(i scope.Duration, ch <-chan []scope.ChannelData) {
	r.i = i
	r.done = make(chan struct{})
	tbCount := int(r.tb / r.i)
	var buf []scope.Voltage
	var l int
	go func() {
		for d := range ch {
			l += len(d[0].Samples)
			buf = append(buf, d[0].Samples...)
			if l >= tbCount {
				l = 0
				r.sweeps = append(r.sweeps, buf)
				buf = nil
			}
		}
		if len(buf) > 0 {
			r.sweeps = append(r.sweeps, buf)
		}
		close(r.done)
	}()
}

// Error reports an acquisition error to the data recorder.
func (r *BufferRecorder) Error(err error) {
	r.err = err
}

// Wait waits until the source finishes writing data to the recorder.
// It then returns recorded data and last error reported (if any).
func (r *BufferRecorder) Wait() ([][]scope.Voltage, error) {
	<-r.done
	return r.sweeps, r.err
}
<commit_msg>Fix the BufferRecorder comment (lint warning)<commit_after>//  Copyright 2016 The goscope Authors
//
//  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 testutil

import "github.com/zagrodzki/goscope/scope"

// BufferRecorder is a buffer that implements the scope.DataRecorder interface.
// Users should call Wait() to get the data that was recorded in the buffer.
type BufferRecorder struct {
	tb     scope.Duration
	i      scope.Duration
	sweeps [][]scope.Voltage
	err    error
	done   chan struct{}
}

// NewBufferRecorder creates a new test data recorder with timebase equal to tb.
func NewBufferRecorder(tb scope.Duration) *BufferRecorder {
	return &BufferRecorder{
		tb: tb,
	}
}

// TimeBase returns the configured timebase (sweep length) of the recorder.
func (r *BufferRecorder) TimeBase() scope.Duration {
	return r.tb
}

// Reset prepares a new recording with sample interval i, reading samples from ch.
func (r *BufferRecorder) Reset(i scope.Duration, ch <-chan []scope.ChannelData) {
	r.i = i
	r.done = make(chan struct{})
	tbCount := int(r.tb / r.i)
	var buf []scope.Voltage
	var l int
	go func() {
		for d := range ch {
			l += len(d[0].Samples)
			buf = append(buf, d[0].Samples...)
			if l >= tbCount {
				l = 0
				r.sweeps = append(r.sweeps, buf)
				buf = nil
			}
		}
		if len(buf) > 0 {
			r.sweeps = append(r.sweeps, buf)
		}
		close(r.done)
	}()
}

// Error reports an acquisition error to the data recorder.
func (r *BufferRecorder) Error(err error) {
	r.err = err
}

// Wait waits until the source finishes writing data to the recorder.
// It then returns recorded data and last error reported (if any).
func (r *BufferRecorder) Wait() ([][]scope.Voltage, error) {
	<-r.done
	return r.sweeps, r.err
}
<|endoftext|>


<commit_before>//
// Copyright 2016 Gregory Trubetskoy. All Rights Reserved.
//
// 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 rrd

import (
	"math"
	"time"
)

type pdp struct {
	Value    float64
	Duration time.Duration
}

func (p *pdp) SetValue(value float64, duration time.Duration) {
	p.Value = value
	p.Duration = duration
}

func (p *pdp) AddValue(value float64, duration time.Duration) {
	if !math.IsNaN(value) && duration > 0 {
		if math.IsNaN(p.Value) {
			p.Value = 0
		}
		p.Value = p.Value*float64(p.Duration)/float64(p.Duration+duration) +
			value*float64(duration)/float64(p.Duration+duration)
		p.Duration = p.Duration + duration
	}
}

func (p *pdp) Reset() float64 {
	result := p.Value
	p.Value = math.NaN()
	p.Duration = 0
	return result
}

type ClockPdp struct {
	pdp
	End time.Time
}

func (p *ClockPdp) AddValue(value float64) {
	if p.End.IsZero() {
		p.End = time.Now()
		return
	}
	duration := time.Now().Sub(p.End)
	p.pdp.AddValue(value, duration)
}
<commit_msg>ClockPdp - remember to update End<commit_after>//
// Copyright 2016 Gregory Trubetskoy. All Rights Reserved.
//
// 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 rrd

import (
	"math"
	"time"
)

type pdp struct {
	Value    float64
	Duration time.Duration
}

func (p *pdp) SetValue(value float64, duration time.Duration) {
	p.Value = value
	p.Duration = duration
}

func (p *pdp) AddValue(value float64, duration time.Duration) {
	if !math.IsNaN(value) && duration > 0 {
		if math.IsNaN(p.Value) {
			p.Value = 0
		}
		p.Value = p.Value*float64(p.Duration)/float64(p.Duration+duration) +
			value*float64(duration)/float64(p.Duration+duration)
		p.Duration = p.Duration + duration
	}
}

func (p *pdp) Reset() float64 {
	result := p.Value
	p.Value = math.NaN()
	p.Duration = 0
	return result
}

type ClockPdp struct {
	pdp
	End time.Time
}

func (p *ClockPdp) AddValue(value float64) {
	if p.End.IsZero() {
		p.End = time.Now()
		return
	}
	now := time.Now()
	duration := now.Sub(p.End)
	p.pdp.AddValue(value, duration)
	p.End = now
}
<|endoftext|>


<commit_before>package request

import (
	"encoding/json"
	"fmt"
	"net/http"
	"net/url"
	"time"
)

// GetCookies gets []*http.Cookie by domain
func (c *Client) GetCookies(domain string) (cookies []*http.Cookie, err error) {
	debug(domain)

	u, err := url.Parse(domain)
	if err != nil {
		debug("ERR(parse)", err)
		return
	}

	cookies = c.httpClient.Jar.Cookies(u)
	return
}

// SetCookies sets cookies by domain
func (c *Client) SetCookies(domain string, cookies []*http.Cookie) (err error) {
	debug(domain)
	debug(cookies)

	u, err := url.Parse(domain)
	if err != nil {
		debug("ERR(parse)", err)
		return
	}

	c.httpClient.Jar.SetCookies(u, cookies)
	return
}

// GetCookie gets cookie value by domain and cookie name
func (c *Client) GetCookie(domain, name string) (value string, err error) {
	debug(domain, name)

	u, err := url.Parse(domain)
	if err != nil {
		debug("ERR(parse)", err)
		return
	}

	cookies := c.httpClient.Jar.Cookies(u)

	for i := 0; i < len(cookies); i++ {
		if cookies[i].Name == name {
			value = cookies[i].Value

			debug("DONE", value)
			return
		}
	}

	debug("EMPTY")
	return
}

// cookie is a duplicate of http.Cookie but with json parser
type cookie struct {
	Name  string `json:"name"`
	Value string `json:"value"`

	Path       string `json:"path"`
	Domain     string `json:"domain"`
	Expires    time.Time
	RawExpires string `json:"expires"` // for reading cookies only

	MaxAge   int  `json:"expiry"`
	Secure   bool `json:"secure"`
	HttpOnly bool `json:"httponly"`
}

// debug purpose
func (c cookie) String() string {
	return fmt.Sprintf("\nName\t\t:%s\nValue\t\t:%s\nPath\t\t:%s\nDomain\t\t:%s\nExpires\t\t:%v\nRawExpires\t:%s\nMaxAge\t\t:%v\nSecure\t\t:%v\nHttpOnly\t:%v\n-------------\n", c.Name, c.Value, c.Path, c.Domain, c.Expires, c.RawExpires, c.MaxAge, c.Secure, c.HttpOnly)
}

func tohttpCookie(cookies []cookie) (httpCookies []*http.Cookie) {
	debug()

	var expires time.Time
	var err error

	for i := 0; i < len(cookies); i++ {
		// .Expires
		expires, err = time.Parse(time.RFC1123, cookies[i].RawExpires)
		if err == nil {
			cookies[i].Expires = expires
		}

		// new httpCookie
		httpCookies = append(httpCookies, &http.Cookie{
			Name:       cookies[i].Name,
			Value:      cookies[i].Value,
			Path:       cookies[i].Path,
			Domain:     cookies[i].Domain,
			Expires:    cookies[i].Expires,
			RawExpires: cookies[i].RawExpires,
			MaxAge:     cookies[i].MaxAge,
			Secure:     cookies[i].Secure,
			HttpOnly:   cookies[i].HttpOnly,
		})
	}

	return
}

func toCookie(httpCookies []*http.Cookie) (cookies []cookie) {
	debug()

	for i := 0; i < len(httpCookies); i++ {
		// new cookie
		cookies = append(cookies, cookie{
			Name:       httpCookies[i].Name,
			Value:      httpCookies[i].Value,
			Path:       httpCookies[i].Path,
			Domain:     httpCookies[i].Domain,
			Expires:    httpCookies[i].Expires,
			RawExpires: httpCookies[i].RawExpires,
			MaxAge:     httpCookies[i].MaxAge,
			Secure:     httpCookies[i].Secure,
			HttpOnly:   httpCookies[i].HttpOnly,
		})
	}

	return
}

// ImportCookie imports cookie from json
func (c *Client) ImportCookie(domain, jsonStr string) (err error) {
	debug("domain:", domain)

	var cookies []cookie

	err = json.Unmarshal([]byte(jsonStr), &cookies)
	if err != nil {
		debug("ERR(json.Unmarshal)", err)
		return
	}
	// debug(cookies)

	httpCookies := tohttpCookie(cookies)

	err = c.SetCookies(domain, httpCookies)
	if err != nil {
		return
	}
	return
}

// ExportCookie exports client cookies as json
func (c *Client) ExportCookie(domain string) (jsonStr string, err error) {
	debug("domain:", domain)

	httpCookies, err := c.GetCookies(domain)
	if err != nil {
		return
	}

	cookies := toCookie(httpCookies)
	// debug(cookies)

	jsonByte, err := json.Marshal(cookies)
	if err != nil {
		debug("ERR(json.Marshal)", err)
		return
	}

	jsonStr = string(jsonByte)
	return
}
<commit_msg>use array of cookie pointers<commit_after>package request

import (
	"encoding/json"
	"fmt"
	"net/http"
	"net/url"
	"time"
)

// GetCookies gets []*http.Cookie by domain
func (c *Client) GetCookies(domain string) (cookies []*http.Cookie, err error) {
	debug(domain)

	u, err := url.Parse(domain)
	if err != nil {
		debug("ERR(parse)", err)
		return
	}

	cookies = c.httpClient.Jar.Cookies(u)
	return
}

// SetCookies sets cookies by domain
func (c *Client) SetCookies(domain string, cookies []*http.Cookie) (err error) {
	debug(domain)
	debug(cookies)

	u, err := url.Parse(domain)
	if err != nil {
		debug("ERR(parse)", err)
		return
	}

	c.httpClient.Jar.SetCookies(u, cookies)
	return
}

// GetCookie gets cookie value by domain and cookie name
func (c *Client) GetCookie(domain, name string) (value string, err error) {
	debug(domain, name)

	u, err := url.Parse(domain)
	if err != nil {
		debug("ERR(parse)", err)
		return
	}

	cookies := c.httpClient.Jar.Cookies(u)

	for i := 0; i < len(cookies); i++ {
		if cookies[i].Name == name {
			value = cookies[i].Value

			debug("DONE", value)
			return
		}
	}

	debug("EMPTY")
	return
}

// cookie is a duplicate of http.Cookie but with json parser
type cookie struct {
	Name  string `json:"name"`
	Value string `json:"value"`

	Path       string `json:"path"`
	Domain     string `json:"domain"`
	Expires    time.Time
	RawExpires string `json:"expires"` // for reading cookies only

	MaxAge   int  `json:"expiry"`
	Secure   bool `json:"secure"`
	HttpOnly bool `json:"httponly"`
}

// debug purpose
func (c cookie) String() string {
	return fmt.Sprintf("\nName\t\t:%s\nValue\t\t:%s\nPath\t\t:%s\nDomain\t\t:%s\nExpires\t\t:%v\nRawExpires\t:%s\nMaxAge\t\t:%v\nSecure\t\t:%v\nHttpOnly\t:%v\n-------------\n", c.Name, c.Value, c.Path, c.Domain, c.Expires, c.RawExpires, c.MaxAge, c.Secure, c.HttpOnly)
}

func tohttpCookie(cookies []*cookie) (httpCookies []*http.Cookie) {
	debug()

	var expires time.Time
	var err error

	for i := 0; i < len(cookies); i++ {
		// .Expires
		expires, err = time.Parse(time.RFC1123, cookies[i].RawExpires)
		if err == nil {
			cookies[i].Expires = expires
		}

		// new httpCookie
		httpCookies = append(httpCookies, &http.Cookie{
			Name:       cookies[i].Name,
			Value:      cookies[i].Value,
			Path:       cookies[i].Path,
			Domain:     cookies[i].Domain,
			Expires:    cookies[i].Expires,
			RawExpires: cookies[i].RawExpires,
			MaxAge:     cookies[i].MaxAge,
			Secure:     cookies[i].Secure,
			HttpOnly:   cookies[i].HttpOnly,
		})
	}

	return
}

func toCookie(httpCookies []*http.Cookie) (cookies []*cookie) {
	debug()

	for i := 0; i < len(httpCookies); i++ {
		// new cookie
		cookies = append(cookies, &cookie{
			Name:       httpCookies[i].Name,
			Value:      httpCookies[i].Value,
			Path:       httpCookies[i].Path,
			Domain:     httpCookies[i].Domain,
			Expires:    httpCookies[i].Expires,
			RawExpires: httpCookies[i].RawExpires,
			MaxAge:     httpCookies[i].MaxAge,
			Secure:     httpCookies[i].Secure,
			HttpOnly:   httpCookies[i].HttpOnly,
		})
	}

	return
}

// ImportCookie imports cookie from json
func (c *Client) ImportCookie(domain, jsonStr string) (err error) {
	debug("domain:", domain)

	var cookies []*cookie

	err = json.Unmarshal([]byte(jsonStr), &cookies)
	if err != nil {
		debug("ERR(json.Unmarshal)", err)
		return
	}
	// debug(cookies)

	httpCookies := tohttpCookie(cookies)

	err = c.SetCookies(domain, httpCookies)
	if err != nil {
		return
	}
	return
}

// ExportCookie exports client cookies as json
func (c *Client) ExportCookie(domain string) (jsonStr string, err error) {
	debug("domain:", domain)

	httpCookies, err := c.GetCookies(domain)
	if err != nil {
		return
	}

	cookies := toCookie(httpCookies)
	// debug(cookies)

	jsonByte, err := json.Marshal(cookies)
	if err != nil {
		debug("ERR(json.Marshal)", err)
		return
	}

	jsonStr = string(jsonByte)
	return
}
<|endoftext|>


<commit_before>package redis

import (
	"github.com/Clever/leakybucket"
	"github.com/garyburd/redigo/redis"
	"time"
)

type bucket struct {
	name                string
	capacity, remaining uint
	reset               time.Time
	rate                time.Duration
	pool                *redis.Pool
}

func (b *bucket) Capacity() uint {
	return b.capacity
}

// Remaining space in the bucket.
func (b *bucket) Remaining() uint {
	return b.remaining
}

// Reset returns when the bucket will be drained.
func (b *bucket) Reset() time.Time {
	return b.reset
}

func (b *bucket) State() leakybucket.BucketState {
	return leakybucket.BucketState{b.Capacity(), b.Remaining(), b.Reset()}
}

// Add to the bucket.
func (b *bucket) Add(amount uint) (leakybucket.BucketState, error) {
	conn := b.pool.Get()
	defer conn.Close()

	if amount > b.remaining {
		return b.State(), leakybucket.ErrorFull
	}

	// If SETNX doesn't return nil, we just set the key. Otherwise, it already exists.
	if set, err := conn.Do("SET", b.name, amount, "NX", "EX", int(b.rate.Seconds())); err != nil {
		return b.State(), err
	} else if set != nil {
		b.remaining = b.capacity - amount
		return b.State(), nil
	}

	if count, err := conn.Do("INCRBY", b.name, amount); err != nil {
		return b.State(), err
	} else {
		b.remaining = b.capacity - uint(count.(int64))
		if b.remaining > b.capacity {
			// We overflowed because of a race condition
			b.remaining = 0
		}
		return b.State(), nil
	}
}

type Storage struct {
	pool *redis.Pool
}

// Create a bucket.
func (s *Storage) Create(name string, capacity uint, rate time.Duration) (leakybucket.Bucket, error) {
	b := &bucket{
		name:      name,
		capacity:  capacity,
		remaining: capacity,
		reset:     time.Now().Add(rate),
		rate:      rate,
		pool:      s.pool,
	}
	return b, nil
}

// New initializes the connection to redis.
func New(network, address string, readTimeout, writeTimeout time.Duration) (*Storage, error) {
	return &Storage{
		pool: redis.NewPool(func() (redis.Conn, error) {
			return redis.Dial(network, address)
		}, 5),
	}, nil
}
<commit_msg>redis: ensure we don't overflow<commit_after>package redis

import (
	"github.com/Clever/leakybucket"
	"github.com/garyburd/redigo/redis"
	"time"
)

type bucket struct {
	name                string
	capacity, remaining uint
	reset               time.Time
	rate                time.Duration
	pool                *redis.Pool
}

func (b *bucket) Capacity() uint {
	return b.capacity
}

// Remaining space in the bucket.
func (b *bucket) Remaining() uint {
	return b.remaining
}

// Reset returns when the bucket will be drained.
func (b *bucket) Reset() time.Time {
	return b.reset
}

func (b *bucket) State() leakybucket.BucketState {
	return leakybucket.BucketState{b.Capacity(), b.Remaining(), b.Reset()}
}

// Add to the bucket.
func (b *bucket) Add(amount uint) (leakybucket.BucketState, error) {
	conn := b.pool.Get()
	defer conn.Close()

	if amount > b.remaining {
		return b.State(), leakybucket.ErrorFull
	}

	// If SETNX doesn't return nil, we just set the key. Otherwise, it already exists.
	if set, err := conn.Do("SET", b.name, amount, "NX", "EX", int(b.rate.Seconds())); err != nil {
		return b.State(), err
	} else if set != nil {
		b.remaining = b.capacity - amount
		return b.State(), nil
	}

	if count, err := conn.Do("INCRBY", b.name, amount); err != nil {
		return b.State(), err
	} else {
		// Ensure we can't overflow
		b.remaining = b.capacity - min(uint(count.(int64)), b.capacity)
		return b.State(), nil
	}
}

type Storage struct {
	pool *redis.Pool
}

// Create a bucket.
func (s *Storage) Create(name string, capacity uint, rate time.Duration) (leakybucket.Bucket, error) {
	b := &bucket{
		name:      name,
		capacity:  capacity,
		remaining: capacity,
		reset:     time.Now().Add(rate),
		rate:      rate,
		pool:      s.pool,
	}
	return b, nil
}

// New initializes the connection to redis.
func New(network, address string, readTimeout, writeTimeout time.Duration) (*Storage, error) {
	return &Storage{
		pool: redis.NewPool(func() (redis.Conn, error) {
			return redis.Dial(network, address)
		}, 5),
	}, nil
}

func min(a, b uint) uint {
	if a < b {
		return a
	} else {
		return b
	}
}
<|endoftext|>