<commit_before>package main

import (
	"context"
	"encoding/json"
	"errors"
	"fmt"
	"io/ioutil"
	"net/http"
	"os"
	"os/exec"
	"strings"
	"time"

	log "github.com/Sirupsen/logrus"
	"github.com/fatih/structs"
	"github.com/gorilla/mux"
	"github.com/maliceio/go-plugin-utils/database/elasticsearch"
	"github.com/maliceio/go-plugin-utils/utils"
	"github.com/maliceio/malice/utils/clitable"
	"github.com/parnurzeal/gorequest"
	"github.com/urfave/cli"
)

// Version stores the plugin's version
var Version string

// BuildTime stores the plugin's build time
var BuildTime string

const (
	name     = "avg"
	category = "av"
)

type pluginResults struct {
	ID   string      `json:"id" structs:"id,omitempty"`
	Data ResultsData `json:"avast" structs:"avg"`
}

// AVG json object
type AVG struct {
	Results ResultsData `json:"avg"`
}

// ResultsData json object
type ResultsData struct {
	Infected bool   `json:"infected" structs:"infected"`
	Result   string `json:"result" structs:"result"`
	Engine   string `json:"engine" structs:"engine"`
	Database string `json:"database" structs:"database"`
	Updated  string `json:"updated" structs:"updated"`
}

// AvScan performs antivirus scan
func AvScan(path string, timeout int) AVG {

	// Give avgd 10 seconds to finish
	avgdCtx, avgdCancel := context.WithTimeout(context.Background(), time.Duration(10)*time.Second)
	defer avgdCancel()
	// AVG needs to have the daemon started first
	_, err := utils.RunCommand(avgdCtx, "/etc/init.d/avgd", "start")
	utils.Assert(err)

	var results ResultsData

	ctx, cancel := context.WithTimeout(context.Background(), time.Duration(timeout)*time.Second)
	defer cancel()

	output, err := utils.RunCommand(ctx, "/usr/bin/avgscan", path)
	results, err = ParseAVGOutput(output, err, path)

	if err != nil {
		// If fails try a second time
		output, err := utils.RunCommand(ctx, "/usr/bin/avgscan", path)
		results, err = ParseAVGOutput(output, err, path)
		utils.Assert(err)
	}

	return AVG{
		Results: results,
	}
}

// ParseAVGOutput convert avg output into ResultsData struct
func ParseAVGOutput(avgout string, err error, path string) (ResultsData, error) {

	if err != nil {
		return ResultsData{}, err
	}

	log.Debug("AVG Output: ", avgout)

	avg := ResultsData{
		Infected: false,
		Engine:   getAvgVersion(),
	}
	colonSeparated := []string{}

	lines := strings.Split(avgout, "\n")
	// Extract Virus string and extract colon separated lines into an slice
	for _, line := range lines {
		if len(line) != 0 {
			if strings.Contains(line, ":") {
				colonSeparated = append(colonSeparated, line)
			}
			if strings.Contains(line, path) {
				pathVirusString := strings.Split(line, "  ")
				avg.Result = strings.TrimSpace(pathVirusString[1])
			}
		}
	}
	// fmt.Println(lines)

	// Extract AVG Details from scan output
	if len(colonSeparated) != 0 {
		for _, line := range colonSeparated {
			if len(line) != 0 {
				keyvalue := strings.Split(line, ":")
				if len(keyvalue) != 0 {
					switch {
					case strings.Contains(line, "Virus database version"):
						avg.Database = strings.TrimSpace(keyvalue[1])
					case strings.Contains(line, "Virus database release date"):
						date := strings.TrimSpace(strings.TrimPrefix(line, "Virus database release date:"))
						avg.Updated = parseUpdatedDate(date)
					case strings.Contains(line, "Infections found"):
						if strings.Contains(keyvalue[1], "1") {
							avg.Infected = true
						}
					}
				}
			}
		}
	} else {
		log.Error("[ERROR] colonSeparated was empty: ", colonSeparated)
		log.Errorf("[ERROR] AVG output was: \n%s", avgout)
		// fmt.Println("[ERROR] colonSeparated was empty: ", colonSeparated)
		// fmt.Printf("[ERROR] AVG output was: \n%s", avgout)
		return ResultsData{}, errors.New("Unable to parse AVG output")
	}

	return avg, nil
}

// Get Anti-Virus scanner version
func getAvgVersion() string {
	versionOut, err := utils.RunCommand(nil, "/usr/bin/avgscan", "-v")
	utils.Assert(err)

	log.Debug("AVG Version: ", versionOut)

	lines := strings.Split(versionOut, "\n")
	for _, line := range lines {
		if len(line) != 0 {
			keyvalue := strings.Split(line, ":")
			if len(keyvalue) != 0 {
				if strings.Contains(keyvalue[0], "Anti-Virus scanner version") {
					return strings.TrimSpace(keyvalue[1])
				}
			}
		}
	}
	return ""
}

func parseUpdatedDate(date string) string {
	layout := "Mon, 02 Jan 2006 15:04:05 +0000"
	t, _ := time.Parse(layout, date)
	return fmt.Sprintf("%d%02d%02d", t.Year(), t.Month(), t.Day())
}

func getUpdatedDate() string {
	if _, err := os.Stat("/opt/malice/UPDATED"); os.IsNotExist(err) {
		return BuildTime
	}
	updated, err := ioutil.ReadFile("/opt/malice/UPDATED")
	utils.Assert(err)
	return string(updated)
}

func updateAV(ctx context.Context) error {
	fmt.Println("Updating AVG...")
	// AVG needs to have the daemon started first
	exec.Command("/etc/init.d/avgd", "start").Output()

	fmt.Println(utils.RunCommand(nil, "avgupdate"))
	// Update UPDATED file
	t := time.Now().Format("20060102")
	err := ioutil.WriteFile("/opt/malice/UPDATED", []byte(t), 0644)
	return err
}

func printMarkDownTable(avg AVG) {

	fmt.Println("#### AVG")
	table := clitable.New([]string{"Infected", "Result", "Engine", "Updated"})
	table.AddRow(map[string]interface{}{
		"Infected": avg.Results.Infected,
		"Result":   avg.Results.Result,
		"Engine":   avg.Results.Engine,
		"Updated":  avg.Results.Updated,
	})
	table.Markdown = true
	table.Print()
}

func printStatus(resp gorequest.Response, body string, errs []error) {
	fmt.Println(body)
}

func webService() {
	router := mux.NewRouter().StrictSlash(true)
	router.HandleFunc("/scan", webAvScan).Methods("POST")
	log.Info("web service listening on port :3993")
	log.Fatal(http.ListenAndServe(":3993", router))
}

func webAvScan(w http.ResponseWriter, r *http.Request) {

	r.ParseMultipartForm(32 << 20)
	file, header, err := r.FormFile("malware")
	if err != nil {
		w.WriteHeader(http.StatusBadRequest)
		fmt.Fprintln(w, "Please supply a valid file to scan.")
		log.Error(err)
	}
	defer file.Close()

	log.Debug("Uploaded fileName: ", header.Filename)

	tmpfile, err := ioutil.TempFile("/malware", "web_")
	if err != nil {
		log.Fatal(err)
	}
	defer os.Remove(tmpfile.Name()) // clean up

	data, err := ioutil.ReadAll(file)

	if _, err = tmpfile.Write(data); err != nil {
		log.Fatal(err)
	}
	if err = tmpfile.Close(); err != nil {
		log.Fatal(err)
	}

	// Do AV scan
	avg := AvScan(tmpfile.Name(), 60)

	w.Header().Set("Content-Type", "application/json; charset=UTF-8")
	w.WriteHeader(http.StatusOK)

	if err := json.NewEncoder(w).Encode(avg); err != nil {
		log.Fatal(err)
	}
}

func main() {

	var elastic string

	cli.AppHelpTemplate = utils.AppHelpTemplate
	app := cli.NewApp()

	app.Name = "avg"
	app.Author = "blacktop"
	app.Email = "https://github.com/blacktop"
	app.Version = Version + ", BuildTime: " + BuildTime
	app.Compiled, _ = time.Parse("20060102", BuildTime)
	app.Usage = "Malice AVG AntiVirus Plugin"
	app.Flags = []cli.Flag{
		cli.BoolFlag{
			Name:  "verbose, V",
			Usage: "verbose output",
		},
		cli.StringFlag{
			Name:        "elasitcsearch",
			Value:       "",
			Usage:       "elasitcsearch address for Malice to store results",
			EnvVar:      "MALICE_ELASTICSEARCH",
			Destination: &elastic,
		},
		cli.BoolFlag{
			Name:  "table, t",
			Usage: "output as Markdown table",
		},
		cli.BoolFlag{
			Name:   "callback, c",
			Usage:  "POST results back to Malice webhook",
			EnvVar: "MALICE_ENDPOINT",
		},
		cli.BoolFlag{
			Name:   "proxy, x",
			Usage:  "proxy settings for Malice webhook endpoint",
			EnvVar: "MALICE_PROXY",
		},
		cli.IntFlag{
			Name:   "timeout",
			Value:  60,
			Usage:  "malice plugin timeout (in seconds)",
			EnvVar: "MALICE_TIMEOUT",
		},
	}
	app.Commands = []cli.Command{
		{
			Name:  "update",
			Usage: "Update virus definitions",
			Action: func(c *cli.Context) error {
				ctx, cancel := context.WithTimeout(context.Background(), time.Duration(c.Int("timeout"))*time.Second)
				defer cancel()

				return updateAV(ctx)
			},
		},
		{
			Name:  "web",
			Usage: "Create a AVG scan web service",
			Action: func(c *cli.Context) error {
				// ctx, cancel := context.WithTimeout(context.Background(), time.Duration(c.Int("timeout"))*time.Second)
				// defer cancel()

				webService()

				return nil
			},
		},
	}
	app.Action = func(c *cli.Context) error {

		if c.Bool("verbose") {
			log.SetLevel(log.DebugLevel)
		}

		if c.Args().Present() {

			path := c.Args().First()

			if _, err := os.Stat(path); os.IsNotExist(err) {
				utils.Assert(err)
			}

			avg := AvScan(path, c.Int("timeout"))

			// upsert into Database
			elasticsearch.InitElasticSearch(elastic)
			elasticsearch.WritePluginResultsToDatabase(elasticsearch.PluginResults{
				ID:       utils.Getopt("MALICE_SCANID", utils.GetSHA256(path)),
				Name:     name,
				Category: category,
				Data:     structs.Map(avg.Results),
			})

			if c.Bool("table") {
				printMarkDownTable(avg)
			} else {
				avgJSON, err := json.Marshal(avg)
				utils.Assert(err)
				if c.Bool("callback") {
					request := gorequest.New()
					if c.Bool("proxy") {
						request = gorequest.New().Proxy(os.Getenv("MALICE_PROXY"))
					}
					request.Post(os.Getenv("MALICE_ENDPOINT")).
						Set("X-Malice-ID", utils.Getopt("MALICE_SCANID", utils.GetSHA256(path))).
						Send(string(avgJSON)).
						End(printStatus)

					return nil
				}
				fmt.Println(string(avgJSON))
			}
		} else {
			log.Fatal(fmt.Errorf("Please supply a file to scan with malice/avg"))
		}
		return nil
	}

	err := app.Run(os.Args)
	utils.Assert(err)
}
<commit_msg>add abs path<commit_after>package main

import (
	"context"
	"encoding/json"
	"errors"
	"fmt"
	"io/ioutil"
	"net/http"
	"os"
	"os/exec"
	"path/filepath"
	"strings"
	"time"

	log "github.com/Sirupsen/logrus"
	"github.com/fatih/structs"
	"github.com/gorilla/mux"
	"github.com/maliceio/go-plugin-utils/database/elasticsearch"
	"github.com/maliceio/go-plugin-utils/utils"
	"github.com/maliceio/malice/utils/clitable"
	"github.com/parnurzeal/gorequest"
	"github.com/urfave/cli"
)

// Version stores the plugin's version
var Version string

// BuildTime stores the plugin's build time
var BuildTime string

const (
	name     = "avg"
	category = "av"
)

type pluginResults struct {
	ID   string      `json:"id" structs:"id,omitempty"`
	Data ResultsData `json:"avast" structs:"avg"`
}

// AVG json object
type AVG struct {
	Results ResultsData `json:"avg"`
}

// ResultsData json object
type ResultsData struct {
	Infected bool   `json:"infected" structs:"infected"`
	Result   string `json:"result" structs:"result"`
	Engine   string `json:"engine" structs:"engine"`
	Database string `json:"database" structs:"database"`
	Updated  string `json:"updated" structs:"updated"`
}

// AvScan performs antivirus scan
func AvScan(path string, timeout int) AVG {

	// Give avgd 10 seconds to finish
	avgdCtx, avgdCancel := context.WithTimeout(context.Background(), time.Duration(10)*time.Second)
	defer avgdCancel()
	// AVG needs to have the daemon started first
	_, err := utils.RunCommand(avgdCtx, "/etc/init.d/avgd", "start")
	utils.Assert(err)

	var results ResultsData

	ctx, cancel := context.WithTimeout(context.Background(), time.Duration(timeout)*time.Second)
	defer cancel()

	output, err := utils.RunCommand(ctx, "/usr/bin/avgscan", path)
	results, err = ParseAVGOutput(output, err, path)

	if err != nil {
		// If fails try a second time
		output, err := utils.RunCommand(ctx, "/usr/bin/avgscan", path)
		results, err = ParseAVGOutput(output, err, path)
		utils.Assert(err)
	}

	return AVG{
		Results: results,
	}
}

// ParseAVGOutput convert avg output into ResultsData struct
func ParseAVGOutput(avgout string, err error, path string) (ResultsData, error) {

	if err != nil {
		return ResultsData{}, err
	}

	log.Debug("AVG Output: ", avgout)

	avg := ResultsData{
		Infected: false,
		Engine:   getAvgVersion(),
	}
	colonSeparated := []string{}

	lines := strings.Split(avgout, "\n")
	// Extract Virus string and extract colon separated lines into an slice
	for _, line := range lines {
		if len(line) != 0 {
			if strings.Contains(line, ":") {
				colonSeparated = append(colonSeparated, line)
			}
			if strings.Contains(line, path) {
				pathVirusString := strings.Split(line, "  ")
				avg.Result = strings.TrimSpace(pathVirusString[1])
			}
		}
	}
	// fmt.Println(lines)

	// Extract AVG Details from scan output
	if len(colonSeparated) != 0 {
		for _, line := range colonSeparated {
			if len(line) != 0 {
				keyvalue := strings.Split(line, ":")
				if len(keyvalue) != 0 {
					switch {
					case strings.Contains(line, "Virus database version"):
						avg.Database = strings.TrimSpace(keyvalue[1])
					case strings.Contains(line, "Virus database release date"):
						date := strings.TrimSpace(strings.TrimPrefix(line, "Virus database release date:"))
						avg.Updated = parseUpdatedDate(date)
					case strings.Contains(line, "Infections found"):
						if strings.Contains(keyvalue[1], "1") {
							avg.Infected = true
						}
					}
				}
			}
		}
	} else {
		log.Error("[ERROR] colonSeparated was empty: ", colonSeparated)
		log.Errorf("[ERROR] AVG output was: \n%s", avgout)
		// fmt.Println("[ERROR] colonSeparated was empty: ", colonSeparated)
		// fmt.Printf("[ERROR] AVG output was: \n%s", avgout)
		return ResultsData{}, errors.New("Unable to parse AVG output")
	}

	return avg, nil
}

// Get Anti-Virus scanner version
func getAvgVersion() string {
	versionOut, err := utils.RunCommand(nil, "/usr/bin/avgscan", "-v")
	utils.Assert(err)

	log.Debug("AVG Version: ", versionOut)

	lines := strings.Split(versionOut, "\n")
	for _, line := range lines {
		if len(line) != 0 {
			keyvalue := strings.Split(line, ":")
			if len(keyvalue) != 0 {
				if strings.Contains(keyvalue[0], "Anti-Virus scanner version") {
					return strings.TrimSpace(keyvalue[1])
				}
			}
		}
	}
	return ""
}

func parseUpdatedDate(date string) string {
	layout := "Mon, 02 Jan 2006 15:04:05 +0000"
	t, _ := time.Parse(layout, date)
	return fmt.Sprintf("%d%02d%02d", t.Year(), t.Month(), t.Day())
}

func getUpdatedDate() string {
	if _, err := os.Stat("/opt/malice/UPDATED"); os.IsNotExist(err) {
		return BuildTime
	}
	updated, err := ioutil.ReadFile("/opt/malice/UPDATED")
	utils.Assert(err)
	return string(updated)
}

func updateAV(ctx context.Context) error {
	fmt.Println("Updating AVG...")
	// AVG needs to have the daemon started first
	exec.Command("/etc/init.d/avgd", "start").Output()

	fmt.Println(utils.RunCommand(nil, "avgupdate"))
	// Update UPDATED file
	t := time.Now().Format("20060102")
	err := ioutil.WriteFile("/opt/malice/UPDATED", []byte(t), 0644)
	return err
}

func printMarkDownTable(avg AVG) {

	fmt.Println("#### AVG")
	table := clitable.New([]string{"Infected", "Result", "Engine", "Updated"})
	table.AddRow(map[string]interface{}{
		"Infected": avg.Results.Infected,
		"Result":   avg.Results.Result,
		"Engine":   avg.Results.Engine,
		"Updated":  avg.Results.Updated,
	})
	table.Markdown = true
	table.Print()
}

func printStatus(resp gorequest.Response, body string, errs []error) {
	fmt.Println(body)
}

func webService() {
	router := mux.NewRouter().StrictSlash(true)
	router.HandleFunc("/scan", webAvScan).Methods("POST")
	log.Info("web service listening on port :3993")
	log.Fatal(http.ListenAndServe(":3993", router))
}

func webAvScan(w http.ResponseWriter, r *http.Request) {

	r.ParseMultipartForm(32 << 20)
	file, header, err := r.FormFile("malware")
	if err != nil {
		w.WriteHeader(http.StatusBadRequest)
		fmt.Fprintln(w, "Please supply a valid file to scan.")
		log.Error(err)
	}
	defer file.Close()

	log.Debug("Uploaded fileName: ", header.Filename)

	tmpfile, err := ioutil.TempFile("/malware", "web_")
	if err != nil {
		log.Fatal(err)
	}
	defer os.Remove(tmpfile.Name()) // clean up

	data, err := ioutil.ReadAll(file)

	if _, err = tmpfile.Write(data); err != nil {
		log.Fatal(err)
	}
	if err = tmpfile.Close(); err != nil {
		log.Fatal(err)
	}

	// Do AV scan
	avg := AvScan(tmpfile.Name(), 60)

	w.Header().Set("Content-Type", "application/json; charset=UTF-8")
	w.WriteHeader(http.StatusOK)

	if err := json.NewEncoder(w).Encode(avg); err != nil {
		log.Fatal(err)
	}
}

func main() {

	var elastic string

	cli.AppHelpTemplate = utils.AppHelpTemplate
	app := cli.NewApp()

	app.Name = "avg"
	app.Author = "blacktop"
	app.Email = "https://github.com/blacktop"
	app.Version = Version + ", BuildTime: " + BuildTime
	app.Compiled, _ = time.Parse("20060102", BuildTime)
	app.Usage = "Malice AVG AntiVirus Plugin"
	app.Flags = []cli.Flag{
		cli.BoolFlag{
			Name:  "verbose, V",
			Usage: "verbose output",
		},
		cli.StringFlag{
			Name:        "elasitcsearch",
			Value:       "",
			Usage:       "elasitcsearch address for Malice to store results",
			EnvVar:      "MALICE_ELASTICSEARCH",
			Destination: &elastic,
		},
		cli.BoolFlag{
			Name:  "table, t",
			Usage: "output as Markdown table",
		},
		cli.BoolFlag{
			Name:   "callback, c",
			Usage:  "POST results back to Malice webhook",
			EnvVar: "MALICE_ENDPOINT",
		},
		cli.BoolFlag{
			Name:   "proxy, x",
			Usage:  "proxy settings for Malice webhook endpoint",
			EnvVar: "MALICE_PROXY",
		},
		cli.IntFlag{
			Name:   "timeout",
			Value:  60,
			Usage:  "malice plugin timeout (in seconds)",
			EnvVar: "MALICE_TIMEOUT",
		},
	}
	app.Commands = []cli.Command{
		{
			Name:  "update",
			Usage: "Update virus definitions",
			Action: func(c *cli.Context) error {
				ctx, cancel := context.WithTimeout(context.Background(), time.Duration(c.Int("timeout"))*time.Second)
				defer cancel()

				return updateAV(ctx)
			},
		},
		{
			Name:  "web",
			Usage: "Create a AVG scan web service",
			Action: func(c *cli.Context) error {
				// ctx, cancel := context.WithTimeout(context.Background(), time.Duration(c.Int("timeout"))*time.Second)
				// defer cancel()

				webService()

				return nil
			},
		},
	}
	app.Action = func(c *cli.Context) error {

		if c.Bool("verbose") {
			log.SetLevel(log.DebugLevel)
		}

		if c.Args().Present() {
			path, err := filepath.Abs(c.Args().First())
			utils.Assert(err)

			if _, err := os.Stat(path); os.IsNotExist(err) {
				utils.Assert(err)
			}

			avg := AvScan(path, c.Int("timeout"))

			// upsert into Database
			elasticsearch.InitElasticSearch(elastic)
			elasticsearch.WritePluginResultsToDatabase(elasticsearch.PluginResults{
				ID:       utils.Getopt("MALICE_SCANID", utils.GetSHA256(path)),
				Name:     name,
				Category: category,
				Data:     structs.Map(avg.Results),
			})

			if c.Bool("table") {
				printMarkDownTable(avg)
			} else {
				avgJSON, err := json.Marshal(avg)
				utils.Assert(err)
				if c.Bool("callback") {
					request := gorequest.New()
					if c.Bool("proxy") {
						request = gorequest.New().Proxy(os.Getenv("MALICE_PROXY"))
					}
					request.Post(os.Getenv("MALICE_ENDPOINT")).
						Set("X-Malice-ID", utils.Getopt("MALICE_SCANID", utils.GetSHA256(path))).
						Send(string(avgJSON)).
						End(printStatus)

					return nil
				}
				fmt.Println(string(avgJSON))
			}
		} else {
			log.Fatal(fmt.Errorf("Please supply a file to scan with malice/avg"))
		}
		return nil
	}

	err := app.Run(os.Args)
	utils.Assert(err)
}
<|endoftext|>


<commit_before>package gosctp

/*
#cgo CFLAGS: -DINET -DINET6 -Wno-deprecated
#cgo LDFLAGS: -lusrsctp -lpthread

#include <stdlib.h>
#include <string.h>
#include <usrsctp.h>

static int g_sctp_ref = 0;

typedef struct {
  struct socket *sock;
  void *udata;
} sctp_transport;

extern void go_sctp_data_ready_cb(sctp_transport *sctp, void *data, size_t len);
static int sctp_data_ready_cb(void *addr, void *data, size_t len, uint8_t tos, uint8_t set_df) {
  go_sctp_data_ready_cb((sctp_transport *)addr, data, len);
  return 0;
}

extern void go_sctp_data_received_cb(sctp_transport *sctp, void *data, size_t len, int sid, int ppid);
extern void go_sctp_notification_received_cb(sctp_transport *sctp, void *data, size_t len);
static int sctp_data_received_cb(struct socket *sock, union sctp_sockstore addr, void *data,
                                 size_t len, struct sctp_rcvinfo recv_info, int flags, void *udata) {
  if (flags & MSG_NOTIFICATION)
    go_sctp_notification_received_cb((sctp_transport *)udata, data, len);
  else
    go_sctp_data_received_cb((sctp_transport *)udata, data, len, recv_info.rcv_sid, ntohl(recv_info.rcv_ppid));

  free(data);
  return 0;
}

static sctp_transport *new_sctp_transport(int port, void *udata) {
  sctp_transport *sctp = (sctp_transport *)calloc(1, sizeof *sctp);
  if (sctp == NULL)
    return NULL;
  sctp->udata = udata;

  if (g_sctp_ref == 0) {
    usrsctp_init(0, sctp_data_ready_cb, NULL);
    usrsctp_sysctl_set_sctp_ecn_enable(0);
  }
  g_sctp_ref++;

  usrsctp_register_address(sctp);
  struct socket *s = usrsctp_socket(AF_CONN, SOCK_STREAM, IPPROTO_SCTP,
                                    sctp_data_received_cb, NULL, 0, sctp);
  if (s == NULL)
    goto trans_err;
  sctp->sock = s;

  struct linger lopt;
  lopt.l_onoff = 1;
  lopt.l_linger = 0;
  usrsctp_setsockopt(s, SOL_SOCKET, SO_LINGER, &lopt, sizeof lopt);

  struct sctp_paddrparams addr_param;
  memset(&addr_param, 0, sizeof addr_param);
  addr_param.spp_flags = SPP_PMTUD_DISABLE;
  addr_param.spp_pathmtu = 1200;
  usrsctp_setsockopt(s, IPPROTO_SCTP, SCTP_PEER_ADDR_PARAMS, &addr_param, sizeof addr_param);

  struct sctp_assoc_value av;
  av.assoc_id = SCTP_ALL_ASSOC;
  av.assoc_value = 1;
  usrsctp_setsockopt(s, IPPROTO_SCTP, SCTP_ENABLE_STREAM_RESET, &av, sizeof av);

  uint32_t nodelay = 1;
  usrsctp_setsockopt(s, IPPROTO_SCTP, SCTP_NODELAY, &nodelay, sizeof nodelay);

  struct sctp_initmsg init_msg;
  memset(&init_msg, 0, sizeof init_msg);
  init_msg.sinit_num_ostreams = 1024;
  init_msg.sinit_max_instreams = 1023;
  usrsctp_setsockopt(s, IPPROTO_SCTP, SCTP_INITMSG, &init_msg, sizeof init_msg);

  struct sockaddr_conn sconn;
  memset(&sconn, 0, sizeof sconn);
  sconn.sconn_family = AF_CONN;
  sconn.sconn_port = htons(port);
  sconn.sconn_addr = (void *)sctp;
#if defined(__APPLE__) || defined(__FreeBSD__) || defined(__OpenBSD__)
  sconn.sconn_len = sizeof *sctp;
#endif
  if (usrsctp_bind(s, (struct sockaddr *)&sconn, sizeof sconn) < 0)
    goto trans_err;

  if (0) {
trans_err:
    usrsctp_finish();
    free(sctp);
    sctp = NULL;
  }

  return sctp;
}

static void release_usrsctp() {
  if (--g_sctp_ref <= 0) {
    g_sctp_ref = 0;
    usrsctp_finish();
  }
}

static ssize_t send_data(sctp_transport *sctp,
                         void *data, size_t len, uint16_t sid, uint32_t ppid)
{
  struct sctp_sndinfo info;
  memset(&info, 0, sizeof info);
  info.snd_sid = sid;
  info.snd_flags = SCTP_EOR;
  info.snd_ppid = htonl(ppid);
  return usrsctp_sendv(sctp->sock, data, len, NULL, 0,
                       &info, sizeof info, SCTP_SENDV_SNDINFO, 0);
}

static int connect_sctp(sctp_transport *sctp, int port) {
  struct sockaddr_conn sconn;
  memset(&sconn, 0, sizeof sconn);
  sconn.sconn_family = AF_CONN;
  sconn.sconn_port = htons(port);
  sconn.sconn_addr = (void *)sctp;
#if defined(__APPLE__) || defined(__FreeBSD__) || defined(__OpenBSD__)
  sconn.sconn_len = sizeof *sctp;
#endif
  if (usrsctp_connect(sctp->sock, (struct sockaddr *)&sconn, sizeof sconn) < 0)
    return -1;

  return 0;
}

static int accept_sctp(sctp_transport *sctp, int port) {
  struct sockaddr_conn sconn;
  memset(&sconn, 0, sizeof sconn);
  sconn.sconn_family = AF_CONN;
  sconn.sconn_port = htons(port);
  sconn.sconn_addr = (void *)sctp;
#if defined(__APPLE__) || defined(__FreeBSD__) || defined(__OpenBSD__)
  sconn.sconn_len = sizeof *sctp;
#endif
  usrsctp_listen(sctp->sock, 1);
  socklen_t len = sizeof sconn;
  struct socket *s = usrsctp_accept(sctp->sock, (struct sockaddr *)&sconn, &len);
  if (s) {
    struct socket *t = sctp->sock;
    sctp->sock = s;
    usrsctp_close(t);
    return 0;
  }

  return -1;
}
*/
import "C"

import (
  "errors"
  "unsafe"
  "sync"
)

type SctpData struct {
  Sid, Ppid int
  Data []byte
}

type SctpTransport struct {
  sctp *C.sctp_transport
  Port int
  mtx sync.Mutex
  BytesToWrite chan []byte
  DataToRead chan *SctpData
}

func NewTransport(port int) (*SctpTransport, error) {
  sctp := C.new_sctp_transport(C.int(port), nil)
  if sctp == nil {
    return nil, errors.New("failed to create SCTP transport")
  }
  s := &SctpTransport{sctp: sctp, Port: port}
  s.BytesToWrite = make(chan []byte, 16)
  s.DataToRead = make(chan *SctpData, 16)
  sctp.udata = unsafe.Pointer(s)
  return s, nil
}

func (s *SctpTransport) Destroy() {
  C.usrsctp_close(s.sctp.sock)
  C.usrsctp_deregister_address(unsafe.Pointer(s.sctp))
  C.free(unsafe.Pointer(s.sctp))
  C.release_usrsctp()
}

//export go_sctp_data_ready_cb
func go_sctp_data_ready_cb(sctp *C.sctp_transport, data unsafe.Pointer, length C.size_t) {
  s := (*SctpTransport)(sctp.udata)
  b := C.GoBytes(data, C.int(length))
  s.BytesToWrite <- b
}

//export go_sctp_data_received_cb
func go_sctp_data_received_cb(sctp *C.sctp_transport, data unsafe.Pointer, length C.size_t, sid, ppid C.int) {
  s := (*SctpTransport)(sctp.udata)
  b := C.GoBytes(data, C.int(length))
  d := &SctpData{int(sid), int(ppid), b}
  s.DataToRead <- d
}

//export go_sctp_notification_received_cb
func go_sctp_notification_received_cb(sctp *C.sctp_transport, data unsafe.Pointer, length C.size_t) {
  // TODO: add interested events
}

func (s *SctpTransport) Feed(data []byte) {
  s.mtx.Lock()
  defer s.mtx.Unlock()
  C.usrsctp_conninput(unsafe.Pointer(s.sctp), unsafe.Pointer(&data[0]), C.size_t(len(data)), 0)
}

func (s *SctpTransport) Send(data []byte, sid, ppid int) (int, error) {
  s.mtx.Lock()
  defer s.mtx.Unlock()
  rv := C.send_data(s.sctp, unsafe.Pointer(&data[0]), C.size_t(len(data)), C.uint16_t(sid), C.uint32_t(ppid))
  if rv < 0 {
    return 0, errors.New("failed to send data")
  }
  return int(rv), nil
}

func (s *SctpTransport) Connect(port int) error {
  rv := C.connect_sctp(s.sctp, C.int(port))
  if rv < 0 {
    return errors.New("failed to connect SCTP transport")
  }
  return nil
}

func (s *SctpTransport) Accept() error {
  rv := C.accept_sctp(s.sctp, C.int(s.Port))
  if rv < 0 {
    return errors.New("failed to accept SCTP transport")
  }
  return nil
}
<commit_msg>change names<commit_after>package gosctp

/*
#cgo CFLAGS: -DINET -DINET6 -Wno-deprecated
#cgo LDFLAGS: -lusrsctp -lpthread

#include <stdlib.h>
#include <string.h>
#include <usrsctp.h>

static int g_sctp_ref = 0;

typedef struct {
  struct socket *sock;
  void *udata;
} sctp_transport;

extern void go_sctp_data_ready_cb(sctp_transport *sctp, void *data, size_t len);
static int sctp_data_ready_cb(void *addr, void *data, size_t len, uint8_t tos, uint8_t set_df) {
  go_sctp_data_ready_cb((sctp_transport *)addr, data, len);
  return 0;
}

extern void go_sctp_data_received_cb(sctp_transport *sctp, void *data, size_t len, int sid, int ppid);
extern void go_sctp_notification_received_cb(sctp_transport *sctp, void *data, size_t len);
static int sctp_data_received_cb(struct socket *sock, union sctp_sockstore addr, void *data,
                                 size_t len, struct sctp_rcvinfo recv_info, int flags, void *udata) {
  if (flags & MSG_NOTIFICATION)
    go_sctp_notification_received_cb((sctp_transport *)udata, data, len);
  else
    go_sctp_data_received_cb((sctp_transport *)udata, data, len, recv_info.rcv_sid, ntohl(recv_info.rcv_ppid));

  free(data);
  return 0;
}

static sctp_transport *new_sctp_transport(int port, void *udata) {
  sctp_transport *sctp = (sctp_transport *)calloc(1, sizeof *sctp);
  if (sctp == NULL)
    return NULL;
  sctp->udata = udata;

  if (g_sctp_ref == 0) {
    usrsctp_init(0, sctp_data_ready_cb, NULL);
    usrsctp_sysctl_set_sctp_ecn_enable(0);
  }
  g_sctp_ref++;

  usrsctp_register_address(sctp);
  struct socket *s = usrsctp_socket(AF_CONN, SOCK_STREAM, IPPROTO_SCTP,
                                    sctp_data_received_cb, NULL, 0, sctp);
  if (s == NULL)
    goto trans_err;
  sctp->sock = s;

  struct linger lopt;
  lopt.l_onoff = 1;
  lopt.l_linger = 0;
  usrsctp_setsockopt(s, SOL_SOCKET, SO_LINGER, &lopt, sizeof lopt);

  struct sctp_paddrparams addr_param;
  memset(&addr_param, 0, sizeof addr_param);
  addr_param.spp_flags = SPP_PMTUD_DISABLE;
  addr_param.spp_pathmtu = 1200;
  usrsctp_setsockopt(s, IPPROTO_SCTP, SCTP_PEER_ADDR_PARAMS, &addr_param, sizeof addr_param);

  struct sctp_assoc_value av;
  av.assoc_id = SCTP_ALL_ASSOC;
  av.assoc_value = 1;
  usrsctp_setsockopt(s, IPPROTO_SCTP, SCTP_ENABLE_STREAM_RESET, &av, sizeof av);

  uint32_t nodelay = 1;
  usrsctp_setsockopt(s, IPPROTO_SCTP, SCTP_NODELAY, &nodelay, sizeof nodelay);

  struct sctp_initmsg init_msg;
  memset(&init_msg, 0, sizeof init_msg);
  init_msg.sinit_num_ostreams = 1024;
  init_msg.sinit_max_instreams = 1023;
  usrsctp_setsockopt(s, IPPROTO_SCTP, SCTP_INITMSG, &init_msg, sizeof init_msg);

  struct sockaddr_conn sconn;
  memset(&sconn, 0, sizeof sconn);
  sconn.sconn_family = AF_CONN;
  sconn.sconn_port = htons(port);
  sconn.sconn_addr = (void *)sctp;
#if defined(__APPLE__) || defined(__FreeBSD__) || defined(__OpenBSD__)
  sconn.sconn_len = sizeof *sctp;
#endif
  if (usrsctp_bind(s, (struct sockaddr *)&sconn, sizeof sconn) < 0)
    goto trans_err;

  if (0) {
trans_err:
    usrsctp_finish();
    free(sctp);
    sctp = NULL;
  }

  return sctp;
}

static void release_usrsctp() {
  if (--g_sctp_ref <= 0) {
    g_sctp_ref = 0;
    usrsctp_finish();
  }
}

static ssize_t send_data(sctp_transport *sctp,
                         void *data, size_t len, uint16_t sid, uint32_t ppid)
{
  struct sctp_sndinfo info;
  memset(&info, 0, sizeof info);
  info.snd_sid = sid;
  info.snd_flags = SCTP_EOR;
  info.snd_ppid = htonl(ppid);
  return usrsctp_sendv(sctp->sock, data, len, NULL, 0,
                       &info, sizeof info, SCTP_SENDV_SNDINFO, 0);
}

static int connect_sctp(sctp_transport *sctp, int port) {
  struct sockaddr_conn sconn;
  memset(&sconn, 0, sizeof sconn);
  sconn.sconn_family = AF_CONN;
  sconn.sconn_port = htons(port);
  sconn.sconn_addr = (void *)sctp;
#if defined(__APPLE__) || defined(__FreeBSD__) || defined(__OpenBSD__)
  sconn.sconn_len = sizeof *sctp;
#endif
  if (usrsctp_connect(sctp->sock, (struct sockaddr *)&sconn, sizeof sconn) < 0)
    return -1;

  return 0;
}

static int accept_sctp(sctp_transport *sctp, int port) {
  struct sockaddr_conn sconn;
  memset(&sconn, 0, sizeof sconn);
  sconn.sconn_family = AF_CONN;
  sconn.sconn_port = htons(port);
  sconn.sconn_addr = (void *)sctp;
#if defined(__APPLE__) || defined(__FreeBSD__) || defined(__OpenBSD__)
  sconn.sconn_len = sizeof *sctp;
#endif
  usrsctp_listen(sctp->sock, 1);
  socklen_t len = sizeof sconn;
  struct socket *s = usrsctp_accept(sctp->sock, (struct sockaddr *)&sconn, &len);
  if (s) {
    struct socket *t = sctp->sock;
    sctp->sock = s;
    usrsctp_close(t);
    return 0;
  }

  return -1;
}
*/
import "C"

import (
  "errors"
  "unsafe"
  "sync"
)

type SctpData struct {
  Sid, Ppid int
  Data []byte
}

type SctpTransport struct {
  sctp *C.sctp_transport
  Port int
  mtx sync.Mutex
  BufferChannel chan []byte
  DataChannel chan *SctpData
}

func NewTransport(port int) (*SctpTransport, error) {
  sctp := C.new_sctp_transport(C.int(port), nil)
  if sctp == nil {
    return nil, errors.New("failed to create SCTP transport")
  }
  s := &SctpTransport{sctp: sctp, Port: port}
  s.BufferChannel = make(chan []byte, 16)
  s.DataChannel = make(chan *SctpData, 16)
  sctp.udata = unsafe.Pointer(s)
  return s, nil
}

func (s *SctpTransport) Destroy() {
  C.usrsctp_close(s.sctp.sock)
  C.usrsctp_deregister_address(unsafe.Pointer(s.sctp))
  C.free(unsafe.Pointer(s.sctp))
  C.release_usrsctp()
}

//export go_sctp_data_ready_cb
func go_sctp_data_ready_cb(sctp *C.sctp_transport, data unsafe.Pointer, length C.size_t) {
  s := (*SctpTransport)(sctp.udata)
  b := C.GoBytes(data, C.int(length))
  s.BufferChannel <- b
}

//export go_sctp_data_received_cb
func go_sctp_data_received_cb(sctp *C.sctp_transport, data unsafe.Pointer, length C.size_t, sid, ppid C.int) {
  s := (*SctpTransport)(sctp.udata)
  b := C.GoBytes(data, C.int(length))
  d := &SctpData{int(sid), int(ppid), b}
  s.DataChannel <- d
}

//export go_sctp_notification_received_cb
func go_sctp_notification_received_cb(sctp *C.sctp_transport, data unsafe.Pointer, length C.size_t) {
  // TODO: add interested events
}

func (s *SctpTransport) Feed(data []byte) {
  s.mtx.Lock()
  defer s.mtx.Unlock()
  C.usrsctp_conninput(unsafe.Pointer(s.sctp), unsafe.Pointer(&data[0]), C.size_t(len(data)), 0)
}

func (s *SctpTransport) Send(data []byte, sid, ppid int) (int, error) {
  s.mtx.Lock()
  defer s.mtx.Unlock()
  rv := C.send_data(s.sctp, unsafe.Pointer(&data[0]), C.size_t(len(data)), C.uint16_t(sid), C.uint32_t(ppid))
  if rv < 0 {
    return 0, errors.New("failed to send data")
  }
  return int(rv), nil
}

func (s *SctpTransport) Connect(port int) error {
  rv := C.connect_sctp(s.sctp, C.int(port))
  if rv < 0 {
    return errors.New("failed to connect SCTP transport")
  }
  return nil
}

func (s *SctpTransport) Accept() error {
  rv := C.accept_sctp(s.sctp, C.int(s.Port))
  if rv < 0 {
    return errors.New("failed to accept SCTP transport")
  }
  return nil
}
<|endoftext|>


<commit_before>package main

import (
	"encoding/base64"
	"flag"
	"fmt"
	"golang.org/x/crypto/sha3"
	"strings"
	"time"
)

func main() {
	var before string
	var after [64]byte
	var base64Encoding string
	var result chan string = make(chan string)
	var resultString string
	var addOnStart int64
	var concurrencyFlag = flag.Int("concurrency", 1, "Number of goroutines to run simultaneously")
	var start time.Time

	flag.Parse()
	concurrency := *concurrencyFlag
	if concurrency == 1 {
		addOn := 0
		start = time.Now()

		for {
			before = fmt.Sprintf("Message%d", addOn)
			after = sha3.Sum512([]byte(before))
			base64Encoding = base64.StdEncoding.EncodeToString(after[:])
			if strings.HasPrefix(base64Encoding, "TEST") {
				fmt.Printf("%d seconds\n", int64(time.Since(start)/time.Second))
				break
			}
			addOn++
		}
		fmt.Printf("%d: %s\n", addOn, base64Encoding)

	} else {
		addOnStart = 0
		start = time.Now()

		sem := make(chan bool, concurrency)
	SEARCHY:
		for {
			sem <- true
			go scan1000000(addOnStart, result, sem)

			select {
			case resultString, _ = <-result:
				break SEARCHY
			default:
			}
			addOnStart++
		}

		fmt.Printf("%d seconds\n", int64(time.Since(start)/time.Second))
		fmt.Print(resultString)
	}
}

func scan1000000(addOnStart int64, result chan string, sem chan bool) {
	var before string
	var after [64]byte
	var base64Encoding string

	defer func() { <-sem }()

	for i := addOnStart * 1000000; i < (addOnStart+1)*1000000; i++ {
		before = fmt.Sprintf("Message%d", i)
		after = sha3.Sum512([]byte(before))
		base64Encoding = base64.StdEncoding.EncodeToString(after[:])

		if strings.HasPrefix(base64Encoding, "TEST") {
			result <- fmt.Sprintf("%d: %s\n", i, base64Encoding)
			break
		}
	}
}
<commit_msg>Add parameter for test string<commit_after>package main

import (
	"encoding/base64"
	"flag"
	"fmt"
	"golang.org/x/crypto/sha3"
	"strings"
	"time"
)

var testString string

func main() {
	var before string
	var after [64]byte
	var base64Encoding string
	var result chan string = make(chan string)
	var resultString string
	var addOnStart int64
	var concurrencyFlag = flag.Int("concurrency", 1, "Number of goroutines to run simultaneously")
	var testStringFlag = flag.String("search", "TEST", "String to search for")
	var start time.Time

	flag.Parse()
	concurrency := *concurrencyFlag
	testString = *testStringFlag
	if concurrency == 1 {
		addOn := 0
		start = time.Now()

		for {
			before = fmt.Sprintf("Message%d", addOn)
			after = sha3.Sum512([]byte(before))
			base64Encoding = base64.StdEncoding.EncodeToString(after[:])
			if strings.HasPrefix(base64Encoding, testString) {
				fmt.Printf("%d seconds\n", int64(time.Since(start)/time.Second))
				break
			}
			addOn++
		}
		fmt.Printf("%d: %s\n", addOn, base64Encoding)

	} else {
		addOnStart = 0
		start = time.Now()

		sem := make(chan bool, concurrency)
	SEARCHY:
		for {
			sem <- true
			go scan1000000(addOnStart, result, sem)

			select {
			case resultString, _ = <-result:
				break SEARCHY
			default:
			}
			addOnStart++
		}

		fmt.Printf("%d seconds\n", int64(time.Since(start)/time.Second))
		fmt.Print(resultString)
	}
}

func scan1000000(addOnStart int64, result chan string, sem chan bool) {
	var before string
	var after [64]byte
	var base64Encoding string

	defer func() { <-sem }()

	for i := addOnStart * 1000000; i < (addOnStart+1)*1000000; i++ {
		before = fmt.Sprintf("Message%d", i)
		after = sha3.Sum512([]byte(before))
		base64Encoding = base64.StdEncoding.EncodeToString(after[:])

		if strings.HasPrefix(base64Encoding, testString) {
			result <- fmt.Sprintf("%d: %s\n", i, base64Encoding)
			break
		}
	}
}
<|endoftext|>


<commit_before>package main

// copyright(c) 2014, Jason E. Aten
//
// goq : a simple queueing system in go; qsub replacement.
//

import (
	"bytes"
	"fmt"
	"os"
	"os/exec"
	"strings"
	"time"
)

// pass by value to avoid races
func (w *Worker) Shepard(jobPtr *Job) {

	// to avoid data races, make a copy of the job value
	j := *jobPtr

	// reset our input channel
	w.DrainTellShepPidKilled()

	go func() {
		jid := j.Id
		dir := j.Dir
		cmd := j.Cmd
		args := j.Args
		env := j.Env
		if j.Out == nil {
			j.Out = make([]string, 0)
		}
		var origdir string
		var err error
		origdir, err = os.Getwd()
		if err != nil {
			panic(err)
		}

		if dir != "" {
			err = os.Chdir(dir)
			if err != nil {
				j.Out = append(j.Out, fmt.Sprintf("Shepard got error trying to move to submit directory with os.Chdir('%s'): %s", dir, err))
				return
			}
			// go back to our starting dir at the end of shepding this job.
			defer os.Chdir(origdir)
		}

		//		if w.ShepStarted != nil {
		//			w.ShepStarted <- true
		//		}

		c := exec.Command(cmd, args...)
		c.Dir = dir
		c.Env = env

		var oe bytes.Buffer
		c.Stdout = &oe
		c.Stderr = &oe

		//	oe, err = c.CombinedOutput()
		err = c.Start()
		if err != nil {
			j.Out = append(j.Out, fmt.Sprintf("Shepard finds non-nil err on trying to Start() cmd '%s' in dir '%s': %s", cmd, dir, err))
			w.ShepSaysJobStarted <- 0
			w.ShepSaysJobDone <- &j
			return
		}

		myPid := c.Process.Pid
		j.Pid = int64(myPid)
		VPrintf("\n SHEP Shepard goroutine about to block on w.ShepSaysJobStarted <- j\n")
		w.ShepSaysJobStarted <- myPid
		VPrintf("\n SHEP Shepard goroutine about to block on c.Wait()\n")

		// this c.Wait() can be 15-20 seconds *slooooow*, so also wait on TellShepPidKilled
		//  to speed things up.
		err = nil
		// waitDone is buffered so this next short goro can exit immediately after c.Wait() finishes.
		// i.e. if TellShepPidKilled arrives first, then there will never be a receiver, so
		// without the buffering the channel and goroutine would be blocked, uncollectable, waiting-forever garbage/leak.
		waitDone := make(chan error, 1)
		go func() {
			err = c.Wait()
			waitDone <- err
		}()

		// back in shep goroutine:
		select {
		case err = <-waitDone:
		case killedPid := <-w.TellShepPidKilled:
			if killedPid != myPid {
				panic(fmt.Sprintf("SHEP error: mismatch in myPid(%d) vs killedPid(%d) received on w.TellShepPidKilled", myPid, killedPid))
			}
			j.Cancelled = true
			WPrintf("\n SHEP got notice from w.TellShepPidKilled, setting j.Cancelled = true\n")
		}

		// Now set j.Out based on which of the two cases we just saw:
		//  Either we saw w.TellShepPidKilled, in which case j.Cancelled == true and we want to exit quickly.
		//  Otherwise, we had a normal or fast c.Wait() exit, and we want to gather and send output on j.Out.
		//
		if j.Cancelled {
			j.Out = append(j.Out, fmt.Sprintf("cancelled/killed: job %d / pid %d ; cmd '%s' in dir '%s' on worker '%s' at '%s'", jid, myPid, cmd, dir, j.Workeraddr, time.Now()))

			// Don't wait around for output/etc.
			// Just skip down to ShepSaysJobDone and get out of here fast.

		} else {
			// Normal/Fast c.Wait() exit:
			WPrintf("\n SHEP DONE with WAIT, err: '%s'\n", err)
			if err != nil && err.Error() == "signal: killed" {
				WPrintf("\n SHEP found 'signal:killed', setting j.Cancelled = true\n")
				j.Cancelled = true
			}
			s := string(oe.Bytes())
			strings.Trim(s, "\n")
			slen := len(s)
			out := strings.Split(s, "\n")
			// if file ended in '\n' then we now have an extra empty line to eliminate.
			N := len(out)
			if slen > 0 && s[slen-1] == '\n' && N > 0 && out[N-1] == "" {
				out = out[:N-1]
			}
			j.Out = append(j.Out, out...)

			if err != nil {
				j.Out = append(j.Out, fmt.Sprintf("Shepard finds non-nil err on trying to Wait() on cmd '%s' in dir '%s': %s", cmd, dir, err))
			}
		}
		WPrintf("end of SHEP: just before w.ShepSaysJobDone <- j\n")
		w.ShepSaysJobDone <- &j
		WPrintf("end of SHEP: just after w.ShepSaysJobDone <- j\n")
	}()
}
<commit_msg>shep uses a defer to always signal worker even on error exit<commit_after>package main

// copyright(c) 2014, Jason E. Aten
//
// goq : a simple queueing system in go; qsub replacement.
//

import (
	"bytes"
	"fmt"
	"os"
	"os/exec"
	"strings"
	"time"
)

// pass by value to avoid races
func (w *Worker) Shepard(jobPtr *Job) {

	// to avoid data races, make a copy of the job value
	j := *jobPtr

	// reset our input channel
	w.DrainTellShepPidKilled()

	go func() {
		myPid := 0

		defer func() {
			// if myPid set then ShepSaysJobStarted already signalled.
			// otherwise indicate error:
			if myPid == 0 {
				w.ShepSaysJobStarted <- 0
			}
			WPrintf("end of SHEP: just before w.ShepSaysJobDone <- j\n")
			w.ShepSaysJobDone <- &j
			WPrintf("end of SHEP: just after w.ShepSaysJobDone <- j\n")
		}()

		jid := j.Id
		dir := j.Dir
		cmd := j.Cmd
		args := j.Args
		env := j.Env
		if j.Out == nil {
			j.Out = make([]string, 0)
		}
		var origdir string
		var err error
		origdir, err = os.Getwd()
		if err != nil {
			panic(err)
		}

		if dir != "" {
			err = os.Chdir(dir)
			if err != nil {
				j.Out = append(j.Out, fmt.Sprintf("Shepard got error trying to move to submit directory with os.Chdir('%s'): %s", dir, err))
				return
			}
			// go back to our starting dir at the end of shepding this job.
			defer os.Chdir(origdir)
		}

		c := exec.Command(cmd, args...)
		c.Dir = dir
		c.Env = env

		var oe bytes.Buffer
		c.Stdout = &oe
		c.Stderr = &oe

		//	oe, err = c.CombinedOutput()
		err = c.Start()
		if err != nil {
			j.Out = append(j.Out, fmt.Sprintf("Shepard finds non-nil err on trying to Start() cmd '%s' in dir '%s': %s", cmd, dir, err))
			return
		}

		// no error/return should be possible between
		// setting myPid and w.ShepSaysJobStarted <- myPid
		myPid = c.Process.Pid
		j.Pid = int64(myPid)
		VPrintf("\n SHEP Shepard goroutine about to block on w.ShepSaysJobStarted <- j\n")
		w.ShepSaysJobStarted <- myPid
		VPrintf("\n SHEP Shepard goroutine about to block on c.Wait()\n")

		// this c.Wait() can be 15-20 seconds *slooooow*, so also wait on TellShepPidKilled
		//  to speed things up.
		err = nil
		// waitDone is buffered so this next short goro can exit immediately after c.Wait() finishes.
		// i.e. if TellShepPidKilled arrives first, then there will never be a receiver, so
		// without the buffering the channel and goroutine would be blocked, uncollectable, waiting-forever garbage/leak.
		waitDone := make(chan error, 1)
		go func() {
			err = c.Wait()
			waitDone <- err
		}()

		// back in shep goroutine:
		select {
		case err = <-waitDone:
		case killedPid := <-w.TellShepPidKilled:
			if killedPid != myPid {
				panic(fmt.Sprintf("SHEP error: mismatch in myPid(%d) vs killedPid(%d) received on w.TellShepPidKilled", myPid, killedPid))
			}
			j.Cancelled = true
			WPrintf("\n SHEP got notice from w.TellShepPidKilled, setting j.Cancelled = true\n")
		}

		// Now set j.Out based on which of the two cases we just saw:
		//  Either we saw w.TellShepPidKilled, in which case j.Cancelled == true and we want to exit quickly.
		//  Otherwise, we had a normal or fast c.Wait() exit, and we want to gather and send output on j.Out.
		//
		if j.Cancelled {
			j.Out = append(j.Out, fmt.Sprintf("cancelled/killed: job %d / pid %d ; cmd '%s' in dir '%s' on worker '%s' at '%s'", jid, myPid, cmd, dir, j.Workeraddr, time.Now()))

			// Don't wait around for output/etc.
			// Just skip down to ShepSaysJobDone and get out of here fast.

		} else {
			// Normal/Fast c.Wait() exit:
			WPrintf("\n SHEP DONE with WAIT, err: '%s'\n", err)
			if err != nil && err.Error() == "signal: killed" {
				WPrintf("\n SHEP found 'signal:killed', setting j.Cancelled = true\n")
				j.Cancelled = true
			}
			s := string(oe.Bytes())
			strings.Trim(s, "\n")
			slen := len(s)
			out := strings.Split(s, "\n")
			// if file ended in '\n' then we now have an extra empty line to eliminate.
			N := len(out)
			if slen > 0 && s[slen-1] == '\n' && N > 0 && out[N-1] == "" {
				out = out[:N-1]
			}
			j.Out = append(j.Out, out...)

			if err != nil {
				j.Out = append(j.Out, fmt.Sprintf("Shepard finds non-nil err on trying to Wait() on cmd '%s' in dir '%s': %s", cmd, dir, err))
			}
		}
	}()
}
<|endoftext|>


<commit_before>//  Copyright (c) 2013 Couchbase, Inc.
//  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 slab

import (
	"encoding/binary"
	"fmt"
	"math"
	"math/rand"
)

type Arena struct {
	growthFactor float64
	slabClasses  []slabClass // The chunkSizes of slabClasses grows by growthFactor.
	slabMagic    int32       // Magic number at the end of each slab memory []byte.
	slabSize     int
	malloc       func(size int) []byte // App-specific allocator; may be nil.
}

type slabClass struct {
	slabs     []*slab  // A growing array of slabs.
	chunkSize int      // Each slab is sliced into fixed-sized chunks.
	chunkFree chunkLoc // Chunks are tracked in a free-list per slabClass.
}

type slab struct {
	memory []byte  // len(memory) == slabSize + SLAB_MEMORY_FOOTER_LEN.
	chunks []chunk // Parallel array of chunk metadata.
}

const SLAB_MEMORY_FOOTER_LEN int = 4 + 4 + 4 // slabClassIndex + slabIndex + slabMagic.

type chunkLoc struct {
	slabClassIndex int
	slabIndex      int
	chunkIndex     int
}

var empty_chunkLoc = chunkLoc{-1, -1, -1} // A sentinel.

func (cl *chunkLoc) isEmpty() bool {
	return cl.slabClassIndex == -1 && cl.slabIndex == -1 && cl.chunkIndex == -1
}

type chunk struct {
	refs int32    // Ref-count.
	self chunkLoc // The self is the chunkLoc for this chunk.
	next chunkLoc // Used when the chunk is in the free-list.
}

// Returns an Arena based on a slab allocator implementation.
// The startChunkSize and slabSize should be > 0.
// The growthFactor should be > 1.0.
// The malloc() func will be invoked when the Arena needs more memory for a new slab.
// The malloc() may be nil, in which case the Arena defaults to make([]byte, size).
func NewArena(startChunkSize int, slabSize int, growthFactor float64,
	malloc func(size int) []byte) *Arena {
	s := &Arena{
		growthFactor: growthFactor,
		slabMagic:    rand.Int31(),
		slabSize:     slabSize,
		malloc:       malloc,
	}
	s.addSlabClass(startChunkSize)
	return s
}

// The input buf must be a buf returned by Alloc().  Once
// the buf's ref-count drops to 0, the Arena may re-use the buf.
// Alloc() may return nil on errors, such as if no more free chunks
// are available and new slab memory was not allocatable (such as if
// malloc() returns nil).
func (s *Arena) Alloc(bufSize int) (buf []byte) {
	if bufSize > s.slabSize {
		return nil
	}
	chunkMem := s.assignChunkMem(s.findSlabClassIndex(bufSize))
	if chunkMem == nil {
		return nil
	}
	return chunkMem[0:bufSize]
}

// The input buf must be a buf returned by Alloc().
func (s *Arena) AddRef(buf []byte) {
	sc, c := s.bufContainer(buf)
	if sc == nil || c == nil {
		panic("buf not from this arena")
	}
	c.addRef()
}

// The buf must be from an Alloc() from the same Arena.
func (s *Arena) DecRef(buf []byte) {
	sc, c := s.bufContainer(buf)
	if sc == nil || c == nil {
		panic("buf not from this arena")
	}
	s.decRef(sc, c)
}

// Returns true if this Arena owns the buf.
func (s *Arena) Owns(buf []byte) bool {
	sc, c := s.bufContainer(buf)
	return sc != nil && c != nil
}

// The buf's from an Arena can be chained.  The returned bufNext may
// be nil.  When the returned bufNext is non-nil, the caller owns a
// ref-count on bufNext and must invoke DecRef(bufNext) when the
// caller is finished using bufNext.
func (s *Arena) GetNext(buf []byte) (bufNext []byte) {
	sc, c := s.bufContainer(buf)
	if sc == nil || c == nil {
		panic("buf not from this arena")
	}
	if c.refs <= 0 {
		panic(fmt.Sprintf("unexpected ref-count during GetNext: %#v", c))
	}
	scNext, cNext := s.chunk(c.next)
	if scNext == nil || cNext == nil {
		return nil
	}
	cNext.addRef()
	return s.chunkMem(cNext)
}

// The buf's from an Arena can be chained, where buf will own an
// AddRef() on bufNext.  When buf's ref-count goes to zero, it will
// call DecRef() on bufNext.  The bufNext may be nil.
func (s *Arena) SetNext(buf, bufNext []byte) {
	sc, c := s.bufContainer(buf)
	if sc == nil || c == nil {
		panic("buf not from this arena")
	}
	if c.refs <= 0 {
		panic(fmt.Sprintf("unexpected ref-count during SetNext: %#v", c))
	}
	scOldNext, cOldNext := s.chunk(c.next)
	if scOldNext != nil && cOldNext != nil {
		s.decRef(scOldNext, cOldNext)
	}
	c.next = empty_chunkLoc
	if bufNext != nil {
		scNewNext, cNewNext := s.bufContainer(bufNext)
		if scNewNext == nil || cNewNext == nil {
			panic("bufNext not from this arena")
		}
		cNewNext.addRef()
		c.next = cNewNext.self
	}
}

func (s *Arena) addSlabClass(chunkSize int) {
	s.slabClasses = append(s.slabClasses, slabClass{
		chunkSize: chunkSize,
		chunkFree: empty_chunkLoc,
	})
}

func (s *Arena) findSlabClassIndex(bufSize int) int {
	curr := 0
	for {
		// TODO: Use binary search instead of linear walk.
		slabClass := &(s.slabClasses[curr])
		if bufSize <= slabClass.chunkSize {
			return curr
		}
		if curr+1 >= len(s.slabClasses) {
			nextChunkSize := float64(slabClass.chunkSize) * s.growthFactor
			s.addSlabClass(int(math.Ceil(nextChunkSize)))
		}
		curr++
	}
}

func (s *Arena) assignChunkMem(slabClassIndex int) (chunkMem []byte) {
	sc := &(s.slabClasses[slabClassIndex])
	if sc.chunkFree.isEmpty() {
		if !s.addSlab(slabClassIndex, s.slabSize, s.slabMagic) {
			return nil
		}
	}
	return sc.chunkMem(sc.popFreeChunk())
}

func (s *Arena) addSlab(slabClassIndex, slabSize int, slabMagic int32) bool {
	sc := &(s.slabClasses[slabClassIndex])
	chunksPerSlab := slabSize / sc.chunkSize
	slabIndex := len(sc.slabs)
	memorySize := (sc.chunkSize * chunksPerSlab) + SLAB_MEMORY_FOOTER_LEN
	var memory []byte
	if s.malloc != nil {
		memory = s.malloc(memorySize)
		if memory == nil {
			return false
		}
	} else {
		memory = make([]byte, memorySize)
	}
	slab := &slab{
		// Re-multiplying to avoid any extra fractional chunk memory.
		memory: memory,
		chunks: make([]chunk, chunksPerSlab),
	}
	footer := slab.memory[len(slab.memory)-SLAB_MEMORY_FOOTER_LEN:]
	binary.BigEndian.PutUint32(footer[0:4], uint32(slabClassIndex))
	binary.BigEndian.PutUint32(footer[4:8], uint32(slabIndex))
	binary.BigEndian.PutUint32(footer[8:12], uint32(slabMagic))
	sc.slabs = append(sc.slabs, slab)
	for i := 0; i < len(slab.chunks); i++ {
		c := &(slab.chunks[i])
		c.self.slabClassIndex = slabClassIndex
		c.self.slabIndex = slabIndex
		c.self.chunkIndex = i
		sc.pushFreeChunk(c)
	}
	return true
}

func (sc *slabClass) pushFreeChunk(c *chunk) {
	if c.refs != 0 {
		panic(fmt.Sprintf("pushFreeChunk() when non-zero refs: %v", c.refs))
	}
	c.next = sc.chunkFree
	sc.chunkFree = c.self
}

func (sc *slabClass) popFreeChunk() *chunk {
	if sc.chunkFree.isEmpty() {
		panic("popFreeChunk() when chunkFree is empty")
	}
	c := sc.chunk(sc.chunkFree)
	if c.refs != 0 {
		panic(fmt.Sprintf("popFreeChunk() when non-zero refs: %v", c.refs))
	}
	c.refs = 1
	sc.chunkFree = c.next
	c.next = empty_chunkLoc
	return c
}

func (sc *slabClass) chunkMem(c *chunk) []byte {
	if c == nil || c.self.isEmpty() {
		return nil
	}
	beg := sc.chunkSize * c.self.chunkIndex
	return sc.slabs[c.self.slabIndex].memory[beg : beg+sc.chunkSize]
}

func (sc *slabClass) chunk(cl chunkLoc) *chunk {
	if cl.isEmpty() {
		return nil
	}
	return &(sc.slabs[cl.slabIndex].chunks[cl.chunkIndex])
}

func (s *Arena) chunkMem(c *chunk) []byte {
	if c == nil || c.self.isEmpty() {
		return nil
	}
	return s.slabClasses[c.self.slabClassIndex].chunkMem(c)
}

func (s *Arena) chunk(cl chunkLoc) (*slabClass, *chunk) {
	if cl.isEmpty() {
		return nil, nil
	}
	sc := &(s.slabClasses[cl.slabClassIndex])
	return sc, sc.chunk(cl)
}

// Determine the slabClass & chunk for a buf []byte.
func (s *Arena) bufContainer(buf []byte) (*slabClass, *chunk) {
	if buf == nil || cap(buf) <= SLAB_MEMORY_FOOTER_LEN {
		return nil, nil
	}
	rest := buf[:cap(buf)]
	footerDistance := len(rest) - SLAB_MEMORY_FOOTER_LEN
	footer := rest[footerDistance:]
	slabClassIndex := binary.BigEndian.Uint32(footer[0:4])
	slabIndex := binary.BigEndian.Uint32(footer[4:8])
	slabMagic := binary.BigEndian.Uint32(footer[8:12])
	if slabMagic != uint32(s.slabMagic) {
		return nil, nil
	}
	sc := &(s.slabClasses[slabClassIndex])
	slab := sc.slabs[slabIndex]
	chunkIndex := len(slab.chunks) - (footerDistance / sc.chunkSize)
	return sc, &(slab.chunks[chunkIndex])
}

func (c *chunk) addRef() *chunk {
	c.refs++
	if c.refs <= 1 {
		panic(fmt.Sprintf("unexpected ref-count during addRef: %#v", c))
	}
	return c
}

func (s *Arena) decRef(sc *slabClass, c *chunk) *chunk {
	c.refs--
	if c.refs < 0 {
		panic(fmt.Sprintf("unexpected ref-count during decRef: %#v", c))
	}
	if c.refs == 0 {
		scNext, cNext := s.chunk(c.next)
		if scNext != nil && cNext != nil {
			s.decRef(scNext, cNext)
		}
		c.next = empty_chunkLoc
		sc.pushFreeChunk(c)
	}
	return c
}
<commit_msg>No special case malloc in the allocation.<commit_after>//  Copyright (c) 2013 Couchbase, Inc.
//  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 slab

import (
	"encoding/binary"
	"fmt"
	"math"
	"math/rand"
)

type Arena struct {
	growthFactor float64
	slabClasses  []slabClass // The chunkSizes of slabClasses grows by growthFactor.
	slabMagic    int32       // Magic number at the end of each slab memory []byte.
	slabSize     int
	malloc       func(size int) []byte // App-specific allocator; may be nil.
}

type slabClass struct {
	slabs     []*slab  // A growing array of slabs.
	chunkSize int      // Each slab is sliced into fixed-sized chunks.
	chunkFree chunkLoc // Chunks are tracked in a free-list per slabClass.
}

type slab struct {
	memory []byte  // len(memory) == slabSize + SLAB_MEMORY_FOOTER_LEN.
	chunks []chunk // Parallel array of chunk metadata.
}

const SLAB_MEMORY_FOOTER_LEN int = 4 + 4 + 4 // slabClassIndex + slabIndex + slabMagic.

type chunkLoc struct {
	slabClassIndex int
	slabIndex      int
	chunkIndex     int
}

var empty_chunkLoc = chunkLoc{-1, -1, -1} // A sentinel.

func (cl *chunkLoc) isEmpty() bool {
	return cl.slabClassIndex == -1 && cl.slabIndex == -1 && cl.chunkIndex == -1
}

type chunk struct {
	refs int32    // Ref-count.
	self chunkLoc // The self is the chunkLoc for this chunk.
	next chunkLoc // Used when the chunk is in the free-list.
}

// Returns an Arena based on a slab allocator implementation.
// The startChunkSize and slabSize should be > 0.
// The growthFactor should be > 1.0.
// The malloc() func will be invoked when the Arena needs more memory for a new slab.
// The malloc() may be nil, in which case the Arena defaults to make([]byte, size).
func NewArena(startChunkSize int, slabSize int, growthFactor float64,
	malloc func(size int) []byte) *Arena {
	if malloc == nil {
		malloc = defaultMalloc
	}
	s := &Arena{
		growthFactor: growthFactor,
		slabMagic:    rand.Int31(),
		slabSize:     slabSize,
		malloc:       malloc,
	}
	s.addSlabClass(startChunkSize)
	return s
}

// The input buf must be a buf returned by Alloc().  Once
// the buf's ref-count drops to 0, the Arena may re-use the buf.
// Alloc() may return nil on errors, such as if no more free chunks
// are available and new slab memory was not allocatable (such as if
// malloc() returns nil).
func (s *Arena) Alloc(bufSize int) (buf []byte) {
	if bufSize > s.slabSize {
		return nil
	}
	chunkMem := s.assignChunkMem(s.findSlabClassIndex(bufSize))
	if chunkMem == nil {
		return nil
	}
	return chunkMem[0:bufSize]
}

// The input buf must be a buf returned by Alloc().
func (s *Arena) AddRef(buf []byte) {
	sc, c := s.bufContainer(buf)
	if sc == nil || c == nil {
		panic("buf not from this arena")
	}
	c.addRef()
}

// The buf must be from an Alloc() from the same Arena.
func (s *Arena) DecRef(buf []byte) {
	sc, c := s.bufContainer(buf)
	if sc == nil || c == nil {
		panic("buf not from this arena")
	}
	s.decRef(sc, c)
}

// Returns true if this Arena owns the buf.
func (s *Arena) Owns(buf []byte) bool {
	sc, c := s.bufContainer(buf)
	return sc != nil && c != nil
}

// The buf's from an Arena can be chained.  The returned bufNext may
// be nil.  When the returned bufNext is non-nil, the caller owns a
// ref-count on bufNext and must invoke DecRef(bufNext) when the
// caller is finished using bufNext.
func (s *Arena) GetNext(buf []byte) (bufNext []byte) {
	sc, c := s.bufContainer(buf)
	if sc == nil || c == nil {
		panic("buf not from this arena")
	}
	if c.refs <= 0 {
		panic(fmt.Sprintf("unexpected ref-count during GetNext: %#v", c))
	}
	scNext, cNext := s.chunk(c.next)
	if scNext == nil || cNext == nil {
		return nil
	}
	cNext.addRef()
	return s.chunkMem(cNext)
}

// The buf's from an Arena can be chained, where buf will own an
// AddRef() on bufNext.  When buf's ref-count goes to zero, it will
// call DecRef() on bufNext.  The bufNext may be nil.
func (s *Arena) SetNext(buf, bufNext []byte) {
	sc, c := s.bufContainer(buf)
	if sc == nil || c == nil {
		panic("buf not from this arena")
	}
	if c.refs <= 0 {
		panic(fmt.Sprintf("unexpected ref-count during SetNext: %#v", c))
	}
	scOldNext, cOldNext := s.chunk(c.next)
	if scOldNext != nil && cOldNext != nil {
		s.decRef(scOldNext, cOldNext)
	}
	c.next = empty_chunkLoc
	if bufNext != nil {
		scNewNext, cNewNext := s.bufContainer(bufNext)
		if scNewNext == nil || cNewNext == nil {
			panic("bufNext not from this arena")
		}
		cNewNext.addRef()
		c.next = cNewNext.self
	}
}

func (s *Arena) addSlabClass(chunkSize int) {
	s.slabClasses = append(s.slabClasses, slabClass{
		chunkSize: chunkSize,
		chunkFree: empty_chunkLoc,
	})
}

func (s *Arena) findSlabClassIndex(bufSize int) int {
	curr := 0
	for {
		// TODO: Use binary search instead of linear walk.
		slabClass := &(s.slabClasses[curr])
		if bufSize <= slabClass.chunkSize {
			return curr
		}
		if curr+1 >= len(s.slabClasses) {
			nextChunkSize := float64(slabClass.chunkSize) * s.growthFactor
			s.addSlabClass(int(math.Ceil(nextChunkSize)))
		}
		curr++
	}
}

func (s *Arena) assignChunkMem(slabClassIndex int) (chunkMem []byte) {
	sc := &(s.slabClasses[slabClassIndex])
	if sc.chunkFree.isEmpty() {
		if !s.addSlab(slabClassIndex, s.slabSize, s.slabMagic) {
			return nil
		}
	}
	return sc.chunkMem(sc.popFreeChunk())
}

func defaultMalloc(size int) []byte {
	return make([]byte, size)
}

func (s *Arena) addSlab(slabClassIndex, slabSize int, slabMagic int32) bool {
	sc := &(s.slabClasses[slabClassIndex])
	chunksPerSlab := slabSize / sc.chunkSize
	slabIndex := len(sc.slabs)
	memorySize := (sc.chunkSize * chunksPerSlab) + SLAB_MEMORY_FOOTER_LEN
	memory := s.malloc(memorySize)
	if memory == nil {
		return false
	}
	slab := &slab{
		// Re-multiplying to avoid any extra fractional chunk memory.
		memory: memory,
		chunks: make([]chunk, chunksPerSlab),
	}
	footer := slab.memory[len(slab.memory)-SLAB_MEMORY_FOOTER_LEN:]
	binary.BigEndian.PutUint32(footer[0:4], uint32(slabClassIndex))
	binary.BigEndian.PutUint32(footer[4:8], uint32(slabIndex))
	binary.BigEndian.PutUint32(footer[8:12], uint32(slabMagic))
	sc.slabs = append(sc.slabs, slab)
	for i := 0; i < len(slab.chunks); i++ {
		c := &(slab.chunks[i])
		c.self.slabClassIndex = slabClassIndex
		c.self.slabIndex = slabIndex
		c.self.chunkIndex = i
		sc.pushFreeChunk(c)
	}
	return true
}

func (sc *slabClass) pushFreeChunk(c *chunk) {
	if c.refs != 0 {
		panic(fmt.Sprintf("pushFreeChunk() when non-zero refs: %v", c.refs))
	}
	c.next = sc.chunkFree
	sc.chunkFree = c.self
}

func (sc *slabClass) popFreeChunk() *chunk {
	if sc.chunkFree.isEmpty() {
		panic("popFreeChunk() when chunkFree is empty")
	}
	c := sc.chunk(sc.chunkFree)
	if c.refs != 0 {
		panic(fmt.Sprintf("popFreeChunk() when non-zero refs: %v", c.refs))
	}
	c.refs = 1
	sc.chunkFree = c.next
	c.next = empty_chunkLoc
	return c
}

func (sc *slabClass) chunkMem(c *chunk) []byte {
	if c == nil || c.self.isEmpty() {
		return nil
	}
	beg := sc.chunkSize * c.self.chunkIndex
	return sc.slabs[c.self.slabIndex].memory[beg : beg+sc.chunkSize]
}

func (sc *slabClass) chunk(cl chunkLoc) *chunk {
	if cl.isEmpty() {
		return nil
	}
	return &(sc.slabs[cl.slabIndex].chunks[cl.chunkIndex])
}

func (s *Arena) chunkMem(c *chunk) []byte {
	if c == nil || c.self.isEmpty() {
		return nil
	}
	return s.slabClasses[c.self.slabClassIndex].chunkMem(c)
}

func (s *Arena) chunk(cl chunkLoc) (*slabClass, *chunk) {
	if cl.isEmpty() {
		return nil, nil
	}
	sc := &(s.slabClasses[cl.slabClassIndex])
	return sc, sc.chunk(cl)
}

// Determine the slabClass & chunk for a buf []byte.
func (s *Arena) bufContainer(buf []byte) (*slabClass, *chunk) {
	if buf == nil || cap(buf) <= SLAB_MEMORY_FOOTER_LEN {
		return nil, nil
	}
	rest := buf[:cap(buf)]
	footerDistance := len(rest) - SLAB_MEMORY_FOOTER_LEN
	footer := rest[footerDistance:]
	slabClassIndex := binary.BigEndian.Uint32(footer[0:4])
	slabIndex := binary.BigEndian.Uint32(footer[4:8])
	slabMagic := binary.BigEndian.Uint32(footer[8:12])
	if slabMagic != uint32(s.slabMagic) {
		return nil, nil
	}
	sc := &(s.slabClasses[slabClassIndex])
	slab := sc.slabs[slabIndex]
	chunkIndex := len(slab.chunks) - (footerDistance / sc.chunkSize)
	return sc, &(slab.chunks[chunkIndex])
}

func (c *chunk) addRef() *chunk {
	c.refs++
	if c.refs <= 1 {
		panic(fmt.Sprintf("unexpected ref-count during addRef: %#v", c))
	}
	return c
}

func (s *Arena) decRef(sc *slabClass, c *chunk) *chunk {
	c.refs--
	if c.refs < 0 {
		panic(fmt.Sprintf("unexpected ref-count during decRef: %#v", c))
	}
	if c.refs == 0 {
		scNext, cNext := s.chunk(c.next)
		if scNext != nil && cNext != nil {
			s.decRef(scNext, cNext)
		}
		c.next = empty_chunkLoc
		sc.pushFreeChunk(c)
	}
	return c
}
<|endoftext|>


<commit_before>// Copyright 2013 by Dobrosław Żybort. All rights reserved.
// This Source Code Form is subject to the terms of the Mozilla Public
// License, v. 2.0. If a copy of the MPL was not distributed with this
// file, You can obtain one at http://mozilla.org/MPL/2.0/.

package slug

import (
	"regexp"
	"strings"

	"github.com/rainycape/unidecode"
)

var (
	// Custom substitution map
	CustomSub map[string]string
	// Custom rune substitution map
	CustomRuneSub map[rune]string

	// Maximum slug length. It's smart so it will cat slug after full word.
	// By default slugs aren't shortened.
	// If MaxLength is smaller than length of the first word, then returned
	// slug will contain only substring from the first word truncated
	// after MaxLength.
	MaxLength int
)

//=============================================================================

// Make returns slug generated from provided string. Will use "en" as language
// substitution.
func Make(s string) (slug string) {
	return MakeLang(s, "en")
}

// MakeLang returns slug generated from provided string and will use provided
// language for chars substitution.
func MakeLang(s string, lang string) (slug string) {
	slug = strings.TrimSpace(s)

	// Custom substitutions
	// Always substitute runes first
	slug = SubstituteRune(slug, CustomRuneSub)
	slug = Substitute(slug, CustomSub)

	// Process string with selected substitution language
	switch lang {
	case "de":
		slug = SubstituteRune(slug, deSub)
	case "en":
		slug = SubstituteRune(slug, enSub)
	case "pl":
		slug = SubstituteRune(slug, plSub)
	case "es":
		slug = SubstituteRune(slug, esSub)
	default: // fallback to "en" if lang not found
		slug = SubstituteRune(slug, enSub)
	}

	slug = SubstituteRune(slug, defaultSub)

	// Process all non ASCII symbols
	slug = unidecode.Unidecode(slug)

	slug = strings.ToLower(slug)

	// Process all remaining symbols
	slug = regexp.MustCompile("[^a-z0-9-_]").ReplaceAllString(slug, "-")
	slug = regexp.MustCompile("-+").ReplaceAllString(slug, "-")
	slug = strings.Trim(slug, "-")

	if MaxLength > 0 {
		slug = smartTruncate(slug)
	}

	return slug
}

// Substitute returns string with superseded all substrings from
// provided substitution map.
func Substitute(s string, sub map[string]string) (buf string) {
	buf = s
	for key, val := range sub {
		buf = strings.Replace(s, key, val, -1)
	}
	return
}

// SubstituteRune substitutes string chars with provided rune
// substitution map.
func SubstituteRune(s string, sub map[rune]string) (buf string) {
	for _, c := range s {
		if d, ok := sub[c]; ok {
			buf += d
		} else {
			buf += string(c)
		}
	}
	return
}

func smartTruncate(text string) string {
	if len(text) < MaxLength {
		return text
	}

	var truncated string
	words := strings.SplitAfter(text, "-")
	// If MaxLength is smaller than length of the first word return word
	// truncated after MaxLength.
	if len(words[0]) > MaxLength {
		return words[0][:MaxLength]
	}
	for _, word := range words {
		if len(truncated)+len(word)-1 <= MaxLength {
			truncated = truncated + word
		} else {
			break
		}
	}
	return strings.Trim(truncated, "-")
}
<commit_msg>Use a buffer instead of naive concatenation in SubstituteRune<commit_after>// Copyright 2013 by Dobrosław Żybort. All rights reserved.
// This Source Code Form is subject to the terms of the Mozilla Public
// License, v. 2.0. If a copy of the MPL was not distributed with this
// file, You can obtain one at http://mozilla.org/MPL/2.0/.

package slug

import (
	"regexp"
	"strings"

	"bytes"

	"github.com/rainycape/unidecode"
)

var (
	// Custom substitution map
	CustomSub map[string]string
	// Custom rune substitution map
	CustomRuneSub map[rune]string

	// Maximum slug length. It's smart so it will cat slug after full word.
	// By default slugs aren't shortened.
	// If MaxLength is smaller than length of the first word, then returned
	// slug will contain only substring from the first word truncated
	// after MaxLength.
	MaxLength int
)

//=============================================================================

// Make returns slug generated from provided string. Will use "en" as language
// substitution.
func Make(s string) (slug string) {
	return MakeLang(s, "en")
}

// MakeLang returns slug generated from provided string and will use provided
// language for chars substitution.
func MakeLang(s string, lang string) (slug string) {
	slug = strings.TrimSpace(s)

	// Custom substitutions
	// Always substitute runes first
	slug = SubstituteRune(slug, CustomRuneSub)
	slug = Substitute(slug, CustomSub)

	// Process string with selected substitution language
	switch lang {
	case "de":
		slug = SubstituteRune(slug, deSub)
	case "en":
		slug = SubstituteRune(slug, enSub)
	case "pl":
		slug = SubstituteRune(slug, plSub)
	case "es":
		slug = SubstituteRune(slug, esSub)
	default: // fallback to "en" if lang not found
		slug = SubstituteRune(slug, enSub)
	}

	slug = SubstituteRune(slug, defaultSub)

	// Process all non ASCII symbols
	slug = unidecode.Unidecode(slug)

	slug = strings.ToLower(slug)

	// Process all remaining symbols
	slug = regexp.MustCompile("[^a-z0-9-_]").ReplaceAllString(slug, "-")
	slug = regexp.MustCompile("-+").ReplaceAllString(slug, "-")
	slug = strings.Trim(slug, "-")

	if MaxLength > 0 {
		slug = smartTruncate(slug)
	}

	return slug
}

// Substitute returns string with superseded all substrings from
// provided substitution map.
func Substitute(s string, sub map[string]string) (buf string) {
	buf = s
	for key, val := range sub {
		buf = strings.Replace(s, key, val, -1)
	}
	return
}

// SubstituteRune substitutes string chars with provided rune
// substitution map.
func SubstituteRune(s string, sub map[rune]string) (result string) {
	var buf bytes.Buffer
	for _, c := range s {
		if d, ok := sub[c]; ok {
			buf.WriteString(d)
		} else {
			buf.WriteRune(c)
		}
	}
	return buf.String()
}

func smartTruncate(text string) string {
	if len(text) < MaxLength {
		return text
	}

	var truncated string
	words := strings.SplitAfter(text, "-")
	// If MaxLength is smaller than length of the first word return word
	// truncated after MaxLength.
	if len(words[0]) > MaxLength {
		return words[0][:MaxLength]
	}
	for _, word := range words {
		if len(truncated)+len(word)-1 <= MaxLength {
			truncated = truncated + word
		} else {
			break
		}
	}
	return strings.Trim(truncated, "-")
}
<|endoftext|>


<commit_before>package main

import (
	"fmt"
	"log"
	"os"
	"path/filepath"
	"strings"
	"sync"
	"time"

	"github.com/soniah/gosnmp"
)

type Control struct {
	nameOid string
	found   map[string]string
	labels  map[string]string
	usable  map[string]struct{}
}

var (
	errorSNMP int
	nameOid   = "1.3.6.1.2.1.31.1.1.1.1" // ifName
)

const (
	maxOids = 60 // const in gosnmp
)

type pduValue struct {
	name, column string
	value        interface{}
}

func getPoint(cfg *SnmpConfig, pdu gosnmp.SnmpPDU) *pduValue {
	i := strings.LastIndex(pdu.Name, ".")
	root := pdu.Name[1:i]
	suffix := pdu.Name[i+1:]
	col := cfg.labels[cfg.asOID[suffix]]
	name, ok := oidToName[root]
	if verbose {
		log.Println("ROOT:", root, "SUFFIX:", suffix, "COL:", col, "NAME:", "VALUE:", pdu.Value)
	}
	if !ok {
		log.Printf("Invalid oid: %s\n", pdu.Name)
		return nil
	}
	if len(col) == 0 {
		log.Println("empty col for:", cfg.asOID[suffix])
		return nil // not an OID of interest
	}
	return &pduValue{name, col, pdu.Value}
}

func bulkPoint(cfg *SnmpConfig, pdu gosnmp.SnmpPDU) *pduValue {
	i := strings.LastIndex(pdu.Name, ".")
	root := pdu.Name[1:i]
	suffix := pdu.Name[i+1:]
	col := cfg.asOID[suffix]
	name, ok := oidToName[root]
	if verbose {
		log.Println("ROOT:", root, "SUFFIX:", suffix, "COL:", col, "NAME:", "VALUE:", pdu.Value)
	}
	if !ok {
		log.Printf("Invalid oid: %s\n", pdu.Name)
		return nil
	}
	if len(col) == 0 {
		log.Println("empty col for:", suffix)
		return nil // not an OID of interest
	}
	return &pduValue{name, col, pdu.Value}
}

func snmpStats(snmp *gosnmp.GoSNMP, cfg *SnmpConfig) error {
	now := time.Now()
	if cfg == nil {
		log.Fatal("cfg is nil")
	}
	if cfg.Influx == nil {
		log.Fatal("influx cfg is nil")
	}
	bps := cfg.Influx.BP()
	// we can only get 'maxOids' worth of snmp requests at a time
	for i := 0; i < len(cfg.oids); i += maxOids {
		end := i + maxOids
		if end > len(cfg.oids) {
			end = len(cfg.oids)
		}
		cfg.incRequests()
		pkt, err := snmp.Get(cfg.oids[i:end])
		if err != nil {
			errLog("SNMP (%s) get error: %s\n", cfg.Host, err)
			cfg.incErrors()
			cfg.LastError = now
			return err
		}
		cfg.incGets()
		if verbose {
			log.Println("SNMP GET CNT:", len(pkt.Variables))
		}
		for _, pdu := range pkt.Variables {
			val := getPoint(cfg, pdu)
			if val == nil {
				continue
			}
			if val.value == nil {
				continue
			}
			pt := makePoint(cfg.Host, val, now)
			bps.Points = append(bps.Points, pt)

		}
	}
	cfg.Influx.Send(bps)
	return nil
}

func bulkStats(snmp *gosnmp.GoSNMP, cfg *SnmpConfig) error {
	now := time.Now()
	if cfg == nil {
		log.Fatal("cfg is nil")
	}
	if cfg.Influx == nil {
		log.Fatal("influx cfg is nil")
	}
	bps := cfg.Influx.BP()
	addPacket := func(pdu gosnmp.SnmpPDU) error {
		val := bulkPoint(cfg, pdu)
		if val != nil && val.value != nil {
			pt := makePoint(cfg.Host, val, now)
			bps.Points = append(bps.Points, pt)
		}
		return nil
	}
	for i := 0; i < len(cfg.oids); i += 1 {
		cfg.incRequests()
		if err := snmp.BulkWalk(cfg.oids[i], addPacket); err != nil {
			errLog("SNMP (%s) get error: %s\n", cfg.Host, err)
			cfg.incErrors()
			cfg.LastError = now
			return err
		}
	}
	cfg.Influx.Send(bps)
	return nil
}

func printSnmpNames(c *SnmpConfig) {
	client, err := snmpClient(c)
	if err != nil {
		fatal(err)
	}
	defer client.Conn.Close()
	pdus, err := client.BulkWalkAll(nameOid)
	if err != nil {
		fatal("SNMP bulkwalk error", err)
	}
	for _, pdu := range pdus {
		switch pdu.Type {
		case gosnmp.OctetString:
			fmt.Println(string(pdu.Value.([]byte)), pdu.Name)
		}
	}
}

func snmpClient(s *SnmpConfig) (*gosnmp.GoSNMP, error) {
	client := &gosnmp.GoSNMP{
		Target:    s.Host,
		Port:      uint16(s.Port),
		Community: s.Public,
		Version:   gosnmp.Version2c,
		Timeout:   time.Duration(s.Timeout) * time.Second,
		Retries:   s.Retries,
	}
	err := client.Connect()
	if err != nil {
		fmt.Fprintln(os.Stderr, err)
	}
	return client, err
}

func (s *SnmpConfig) DebugLog() *log.Logger {
	name := filepath.Join(logDir, "debug_"+strings.Replace(s.Host, ".", "-", -1)+".log")
	if l, err := os.OpenFile(name, os.O_CREATE|os.O_RDWR|os.O_APPEND, 0664); err == nil {
		return log.New(l, "", 0)
	} else {
		fmt.Fprintln(os.Stderr, err)
		return nil
	}
}

func (s *SnmpConfig) Gather(count int, wg *sync.WaitGroup) {
	debug := false
	client, err := snmpClient(s)
	if err != nil {
		fatal(err)
	}
	defer client.Conn.Close()
	spew(strings.Join(s.oids, "\n"))
	fn := snmpStats
	if len(s.PortFile) == 0 {
		fn = bulkStats
	}
	c := time.Tick(time.Duration(s.Freq) * time.Second)
	for {
		err := fn(client, s)
		if count > 0 {
			count--
			if count == 0 {
				break
			}
		}
		// was seeing clients getting "wedged" -- so just restart
		if err != nil {
			errLog("snmp error - reloading snmp client: %s", err)
			client.Conn.Close()
			for {
				if client, err = snmpClient(s); err == nil {
					break
				}
				errLog("snmp client connect error: %s", err)
				time.Sleep(time.Duration(s.Timeout) * time.Second)
			}
		}

		// pause for interval period and have optional debug toggling
	LOOP:
		for {
			select {
			case <-c:
				break LOOP
			case debug := <-s.debugging:
				log.Println("debugging:", debug)
				if debug && client.Logger == nil {
					client.Logger = s.DebugLog()
				} else {
					client.Logger = nil
				}
			case status := <-s.enabled:
				status <- debug
			}
		}
	}
	wg.Done()
}
<commit_msg>use external package for snmp handling, add snmp v3 support, regexp filtering, bulkwalk of any table<commit_after><|endoftext|>


<commit_before>package zipkintracer

import (
	"fmt"
	"sync"
	"time"

	opentracing "github.com/opentracing/opentracing-go"
	"github.com/opentracing/opentracing-go/ext"

	"github.com/openzipkin/zipkin-go-opentracing/_thrift/gen-go/zipkincore"
)

// Span provides access to the essential details of the span, for use
// by zipkintracer consumers.  These methods may only be called prior
// to (*opentracing.Span).Finish().
type Span interface {
	opentracing.Span

	// Context contains trace identifiers
	Context() Context

	// Operation names the work done by this span instance
	Operation() string

	// Start indicates when the span began
	Start() time.Time
}

// Implements the `Span` interface. Created via tracerImpl (see
// `zipkintracer.NewTracer()`).
type spanImpl struct {
	tracer     *tracerImpl
	event      func(SpanEvent)
	sync.Mutex // protects the fields below
	raw        RawSpan
	Endpoint   *zipkincore.Endpoint
	sampled    bool
}

var spanPool = &sync.Pool{New: func() interface{} {
	return &spanImpl{}
}}

func (s *spanImpl) reset() {
	s.tracer, s.event = nil, nil
	// Note: Would like to do the following, but then the consumer of RawSpan
	// (the recorder) needs to make sure that they're not holding on to the
	// baggage or logs when they return (i.e. they need to copy if they care):
	//
	// logs, baggage := s.raw.Logs[:0], s.raw.Baggage
	// for k := range baggage {
	// 	delete(baggage, k)
	// }
	// s.raw.Logs, s.raw.Baggage = logs, baggage
	//
	// That's likely too much to ask for. But there is some magic we should
	// be able to do with `runtime.SetFinalizer` to reclaim that memory into
	// a buffer pool when GC considers them unreachable, which should ease
	// some of the load. Hard to say how quickly that would be in practice
	// though.
	s.raw = RawSpan{}
}

func (s *spanImpl) SetOperationName(operationName string) opentracing.Span {
	s.Lock()
	defer s.Unlock()
	s.raw.Operation = operationName
	return s
}

func (s *spanImpl) trim() bool {
	return !s.raw.Sampled && s.tracer.options.trimUnsampledSpans
}

func (s *spanImpl) SetTag(key string, value interface{}) opentracing.Span {
	defer s.onTag(key, value)
	s.Lock()
	defer s.Unlock()
	if key == string(ext.SamplingPriority) {
		s.raw.Sampled = true
		return s
	}
	if s.trim() {
		return s
	}

	if s.raw.Tags == nil {
		s.raw.Tags = opentracing.Tags{}
	}
	s.raw.Tags[key] = value
	return s
}

func (s *spanImpl) LogEvent(event string) {
	s.Log(opentracing.LogData{
		Event: event,
	})
}

func (s *spanImpl) LogEventWithPayload(event string, payload interface{}) {
	s.Log(opentracing.LogData{
		Event:   event,
		Payload: payload,
	})
}

func (s *spanImpl) Log(ld opentracing.LogData) {
	defer s.onLog(ld)
	s.Lock()
	defer s.Unlock()
	if s.trim() {
		return
	}

	if ld.Timestamp.IsZero() {
		ld.Timestamp = time.Now()
	}

	s.raw.Logs = append(s.raw.Logs, ld)
}

func (s *spanImpl) Finish() {
	s.FinishWithOptions(opentracing.FinishOptions{})
}

func (s *spanImpl) FinishWithOptions(opts opentracing.FinishOptions) {
	finishTime := opts.FinishTime
	if finishTime.IsZero() {
		finishTime = time.Now()
	}
	duration := finishTime.Sub(s.raw.Start)

	s.Lock()
	defer s.Unlock()
	if opts.BulkLogData != nil {
		s.raw.Logs = append(s.raw.Logs, opts.BulkLogData...)
	}
	s.raw.Duration = duration

	s.onFinish(s.raw)
	s.tracer.options.recorder.RecordSpan(s.raw)
	if s.tracer.options.debugAssertUseAfterFinish {
		// This makes it much more likely to catch a panic on any subsequent
		// operation since s.tracer is accessed on every call to `Lock`.
		s.reset()
	}
	spanPool.Put(s)
}

func (s *spanImpl) SetBaggageItem(restrictedKey, val string) opentracing.Span {
	canonicalKey, valid := opentracing.CanonicalizeBaggageKey(restrictedKey)
	if !valid {
		panic(fmt.Errorf("Invalid key: %q", restrictedKey))
	}

	s.Lock()
	defer s.Unlock()
	s.onBaggage(canonicalKey, val)
	if s.trim() {
		return s
	}

	if s.raw.Baggage == nil {
		s.raw.Baggage = make(map[string]string)
	}
	s.raw.Baggage[canonicalKey] = val
	return s
}

func (s *spanImpl) BaggageItem(restrictedKey string) string {
	canonicalKey, valid := opentracing.CanonicalizeBaggageKey(restrictedKey)
	if !valid {
		panic(fmt.Errorf("Invalid key: %q", restrictedKey))
	}

	s.Lock()
	defer s.Unlock()

	return s.raw.Baggage[canonicalKey]
}

func (s *spanImpl) Tracer() opentracing.Tracer {
	return s.tracer
}

func (s *spanImpl) Context() Context {
	return s.raw.Context
}

func (s *spanImpl) Operation() string {
	return s.raw.Operation
}

func (s *spanImpl) Start() time.Time {
	return s.raw.Start
}
<commit_msg>support for new Open Tracing method ForeachBaggageItem<commit_after>package zipkintracer

import (
	"fmt"
	"sync"
	"time"

	opentracing "github.com/opentracing/opentracing-go"
	"github.com/opentracing/opentracing-go/ext"

	"github.com/openzipkin/zipkin-go-opentracing/_thrift/gen-go/zipkincore"
)

// Span provides access to the essential details of the span, for use
// by zipkintracer consumers.  These methods may only be called prior
// to (*opentracing.Span).Finish().
type Span interface {
	opentracing.Span

	// Context contains trace identifiers
	Context() Context

	// Operation names the work done by this span instance
	Operation() string

	// Start indicates when the span began
	Start() time.Time
}

// Implements the `Span` interface. Created via tracerImpl (see
// `zipkintracer.NewTracer()`).
type spanImpl struct {
	tracer     *tracerImpl
	event      func(SpanEvent)
	sync.Mutex // protects the fields below
	raw        RawSpan
	Endpoint   *zipkincore.Endpoint
	sampled    bool
}

var spanPool = &sync.Pool{New: func() interface{} {
	return &spanImpl{}
}}

func (s *spanImpl) reset() {
	s.tracer, s.event = nil, nil
	// Note: Would like to do the following, but then the consumer of RawSpan
	// (the recorder) needs to make sure that they're not holding on to the
	// baggage or logs when they return (i.e. they need to copy if they care):
	//
	// logs, baggage := s.raw.Logs[:0], s.raw.Baggage
	// for k := range baggage {
	// 	delete(baggage, k)
	// }
	// s.raw.Logs, s.raw.Baggage = logs, baggage
	//
	// That's likely too much to ask for. But there is some magic we should
	// be able to do with `runtime.SetFinalizer` to reclaim that memory into
	// a buffer pool when GC considers them unreachable, which should ease
	// some of the load. Hard to say how quickly that would be in practice
	// though.
	s.raw = RawSpan{}
}

func (s *spanImpl) SetOperationName(operationName string) opentracing.Span {
	s.Lock()
	defer s.Unlock()
	s.raw.Operation = operationName
	return s
}

func (s *spanImpl) trim() bool {
	return !s.raw.Sampled && s.tracer.options.trimUnsampledSpans
}

func (s *spanImpl) SetTag(key string, value interface{}) opentracing.Span {
	defer s.onTag(key, value)
	s.Lock()
	defer s.Unlock()
	if key == string(ext.SamplingPriority) {
		s.raw.Sampled = true
		return s
	}
	if s.trim() {
		return s
	}

	if s.raw.Tags == nil {
		s.raw.Tags = opentracing.Tags{}
	}
	s.raw.Tags[key] = value
	return s
}

func (s *spanImpl) LogEvent(event string) {
	s.Log(opentracing.LogData{
		Event: event,
	})
}

func (s *spanImpl) LogEventWithPayload(event string, payload interface{}) {
	s.Log(opentracing.LogData{
		Event:   event,
		Payload: payload,
	})
}

func (s *spanImpl) Log(ld opentracing.LogData) {
	defer s.onLog(ld)
	s.Lock()
	defer s.Unlock()
	if s.trim() {
		return
	}

	if ld.Timestamp.IsZero() {
		ld.Timestamp = time.Now()
	}

	s.raw.Logs = append(s.raw.Logs, ld)
}

func (s *spanImpl) Finish() {
	s.FinishWithOptions(opentracing.FinishOptions{})
}

func (s *spanImpl) FinishWithOptions(opts opentracing.FinishOptions) {
	finishTime := opts.FinishTime
	if finishTime.IsZero() {
		finishTime = time.Now()
	}
	duration := finishTime.Sub(s.raw.Start)

	s.Lock()
	defer s.Unlock()
	if opts.BulkLogData != nil {
		s.raw.Logs = append(s.raw.Logs, opts.BulkLogData...)
	}
	s.raw.Duration = duration

	s.onFinish(s.raw)
	s.tracer.options.recorder.RecordSpan(s.raw)
	if s.tracer.options.debugAssertUseAfterFinish {
		// This makes it much more likely to catch a panic on any subsequent
		// operation since s.tracer is accessed on every call to `Lock`.
		s.reset()
	}
	spanPool.Put(s)
}

func (s *spanImpl) SetBaggageItem(restrictedKey, val string) opentracing.Span {
	canonicalKey, valid := opentracing.CanonicalizeBaggageKey(restrictedKey)
	if !valid {
		panic(fmt.Errorf("Invalid key: %q", restrictedKey))
	}

	s.Lock()
	defer s.Unlock()
	s.onBaggage(canonicalKey, val)
	if s.trim() {
		return s
	}

	if s.raw.Baggage == nil {
		s.raw.Baggage = make(map[string]string)
	}
	s.raw.Baggage[canonicalKey] = val
	return s
}

func (s *spanImpl) BaggageItem(restrictedKey string) string {
	canonicalKey, valid := opentracing.CanonicalizeBaggageKey(restrictedKey)
	if !valid {
		panic(fmt.Errorf("Invalid key: %q", restrictedKey))
	}

	s.Lock()
	defer s.Unlock()

	return s.raw.Baggage[canonicalKey]
}

func (s *spanImpl) ForeachBaggageItem(handler func(k, v string) bool) {
	s.Lock()
	defer s.Unlock()
	for k, v := range s.raw.Baggage {
		if !handler(k, v) {
			break
		}
	}
}

func (s *spanImpl) Tracer() opentracing.Tracer {
	return s.tracer
}

func (s *spanImpl) Context() Context {
	return s.raw.Context
}

func (s *spanImpl) Operation() string {
	return s.raw.Operation
}

func (s *spanImpl) Start() time.Time {
	return s.raw.Start
}
<|endoftext|>


<commit_before>package main

import (
	jwt "github.com/dgrijalva/jwt-go"
)

const (
	LoginEvent                = "gusher.login"
	SubscribeEvent            = "gusher.subscribe"
	UnSubscribeEvent          = "gusher.unsubscribe"
	SubscribeReplySucceeded   = "subscribe_succeeded"
	SubscribeReplyError       = "subscribe_error"
	UnSubscribeReplySucceeded = "unsubscribe_succeeded"
	UnSubscribeReplyError     = "unsubscribe_error"
)

type InternalCommand struct {
	Event string `json:"event"`
}

type ChannelCommand struct {
	InternalCommand
	Data ChannelData `json:"data"`
}
type ChannelData struct {
	Channel string `json:"channel"`
}

type CommonMessage struct {
	Channel string      `json:"channel"`
	Event   string      `json:"event"`
	Data    interface{} `json:"data"`
}

type JwtPack struct {
	Gusher Auth `json:"gusher"`
	jwt.StandardClaims
}
type Auth struct {
	Channels []string `json:"channels"`
	UserId   string   `json:"user_id"`
	AppKey   string   `json:"app_key"`
}

/*rpc use*/
type ServerInfo struct {
	Ip             string `json:"ip"`
	LocalListen    string `json:"local_listen"`
	Version        string `json:"version"`
	RunTimeVersion string `json:"runtime_version"`
	NumCpu         int    `json:"cpu"`
	MemAllcoated   uint64 `json:"usage-memory"`
	Goroutines     int    `json:"goroutines"`
	Connections    int    `json:"connections"`
	SendInterval   string `json:"send_interval"`
	UpdateTime     int64  `json:"update_time"`
}
<commit_msg>remove<commit_after>package main

import (
	jwt "github.com/dgrijalva/jwt-go"
)

const (
	LoginEvent                = "gusher.login"
	SubscribeEvent            = "gusher.subscribe"
	UnSubscribeEvent          = "gusher.unsubscribe"
	SubscribeReplySucceeded   = "subscribe_succeeded"
	SubscribeReplyError       = "subscribe_error"
	UnSubscribeReplySucceeded = "unsubscribe_succeeded"
	UnSubscribeReplyError     = "unsubscribe_error"
)

type InternalCommand struct {
	Event string `json:"event"`
}

type ChannelCommand struct {
	InternalCommand
	Data ChannelData `json:"data"`
}
type ChannelData struct {
	Channel string `json:"channel"`
}

type JwtPack struct {
	Gusher Auth `json:"gusher"`
	jwt.StandardClaims
}
type Auth struct {
	Channels []string `json:"channels"`
	UserId   string   `json:"user_id"`
	AppKey   string   `json:"app_key"`
}

/*rpc use*/
type ServerInfo struct {
	Ip             string `json:"ip"`
	LocalListen    string `json:"local_listen"`
	Version        string `json:"version"`
	RunTimeVersion string `json:"runtime_version"`
	NumCpu         int    `json:"cpu"`
	MemAllcoated   uint64 `json:"usage-memory"`
	Goroutines     int    `json:"goroutines"`
	Connections    int    `json:"connections"`
	SendInterval   string `json:"send_interval"`
	UpdateTime     int64  `json:"update_time"`
}
<|endoftext|>


<commit_before>package hotshell

import (
	"fmt"
	"github.com/julienmoumne/hotshell/item"
	pkgterm "github.com/pkg/term"
)

const DEFAULT_TTY = "/dev/tty"

var Tty = DEFAULT_TTY

type term struct {
	term *pkgterm.Term
}

func NewTerm() (*term, error) {

	t, err := pkgterm.Open(Tty)

	if err != nil {
		return nil, err
	}

	return &term{t}, err
}

func (t *term) close() error {

	return t.term.Close()
}

func (t *term) restore() {

	// term.Restore() blocks when running tests with pseudo term (termios.Pty())
	if Tty != DEFAULT_TTY {
		return
	}

	err := t.term.Restore()
	if err == nil {
		return
	}

	fmt.Printf("An error occurred while restoring your terminal default values : %s\n", err)
	fmt.Println("Your terminal may behave differently than usual.")
	fmt.Println("If it is the case, you can close and start it again.")
	fmt.Println("Please file a bug report at https://github.com/julienmoumne/hotshell/issues")
}

func (t *term) readUserChoice() (item.Key, error) {

	err := pkgterm.CBreakMode(t.term)
	defer t.restore()

	if err != nil {
		return item.NUL_KEY, err
	}

	bytes := make([]byte, 1)
	_, err = t.term.Read(bytes)

	return item.MakeKey(string(bytes)), err
}
<commit_msg>use issues/new<commit_after>package hotshell

import (
	"fmt"
	"github.com/julienmoumne/hotshell/item"
	pkgterm "github.com/pkg/term"
)

const DEFAULT_TTY = "/dev/tty"

var Tty = DEFAULT_TTY

type term struct {
	term *pkgterm.Term
}

func NewTerm() (*term, error) {

	t, err := pkgterm.Open(Tty)

	if err != nil {
		return nil, err
	}

	return &term{t}, err
}

func (t *term) close() error {

	return t.term.Close()
}

func (t *term) restore() {

	// term.Restore() blocks when running tests with pseudo term (termios.Pty())
	if Tty != DEFAULT_TTY {
		return
	}

	err := t.term.Restore()
	if err == nil {
		return
	}

	fmt.Printf("An error occurred while restoring your terminal default values : %s\n", err)
	fmt.Println("Your terminal may behave differently than usual.")
	fmt.Println("If it is the case, you can close and start it again.")
	fmt.Println("Please file a bug report at https://github.com/julienmoumne/hotshell/issues/new")
}

func (t *term) readUserChoice() (item.Key, error) {

	err := pkgterm.CBreakMode(t.term)
	defer t.restore()

	if err != nil {
		return item.NUL_KEY, err
	}

	bytes := make([]byte, 1)
	_, err = t.term.Read(bytes)

	return item.MakeKey(string(bytes)), err
}
<|endoftext|>


<commit_before>package main

import (
	"io"
	"io/ioutil"
	"log"
	"os"
	"os/exec"
)

// Errors & warnings are deliberately suppressed as tidy throws warnings very easily
func Tidy(r io.Reader, xmlIn bool) ([]byte, error) {
	f, err := ioutil.TempFile("/tmp", "sajari-convert-")
	if err != nil {
		log.Println("TempFile:", err)
		return nil, err
	}
	defer os.Remove(f.Name())
	io.Copy(f, r)

	var output []byte
	if xmlIn {
		output, err = exec.Command("tidy", "-xml", "-numeric", "-asxml", "-quiet", "-utf8", f.Name()).Output()
	} else {
		output, err = exec.Command("tidy", "-numeric", "-asxml", "-quiet", "-utf8", f.Name()).Output()
	}

	if err != nil && err.Error() != "exit status 1" {
		return nil, err
	}
	return output, nil
}
<commit_msg>Remove unnecessary log from Tidy.<commit_after>package main

import (
	"io"
	"io/ioutil"
	"os"
	"os/exec"
)

// Errors & warnings are deliberately suppressed as tidy throws warnings very easily
func Tidy(r io.Reader, xmlIn bool) ([]byte, error) {
	f, err := ioutil.TempFile("/tmp", "sajari-convert-")
	if err != nil {
		return nil, err
	}
	defer os.Remove(f.Name())
	io.Copy(f, r)

	var output []byte
	if xmlIn {
		output, err = exec.Command("tidy", "-xml", "-numeric", "-asxml", "-quiet", "-utf8", f.Name()).Output()
	} else {
		output, err = exec.Command("tidy", "-numeric", "-asxml", "-quiet", "-utf8", f.Name()).Output()
	}

	if err != nil && err.Error() != "exit status 1" {
		return nil, err
	}
	return output, nil
}
<|endoftext|>


<commit_before>package main
import("math")

type ToneGenerator struct {
	sampleRate float64
	step float64 //delta t
}

func NewToneGenerator(sampleRate float64) *ToneGenerator {
	if sampleRate < 1 {
		return nil
	}
	t := new(ToneGenerator)
	t.sampleRate = sampleRate
	t.step = 1/sampleRate
	return t
}

func (t *ToneGenerator) Tone(freq, seconds float64, vol int32) []int32{
	return t.SinTone(freq, seconds, vol)
}

//Generates a sin wave of a certain freq for a specific volume
func (t *ToneGenerator) SinTone(freq, seconds float64, vol int32) []int32{
	var synthArray = make([]int32, int(seconds*t.sampleRate)) //duration/step = dur*sR
	for i:=0; i < len(synthArray); i++{
		synthArray[i] = int32(float64(vol)*math.Sin(freq *2* math.Pi * float64(i) * t.step))
	}
	return synthArray
}

//Generates a square wave
func (t *ToneGenerator) SquareTone(freq, seconds float64, vol int32) []int32{
	var synthArray = make([]int32, int(seconds*t.sampleRate)) //duration/step = dur*sR
	var period = period(freq)

	//How many samples are there in a period
	var samplesPerPeriod = int(period/t.step)

	for i:=0; i < len(synthArray); i++{
		if i%samplesPerPeriod < samplesPerPeriod/2 { //Width: first half = 1, second = 0
			synthArray[i] = vol
		} else {
			synthArray[i] = -vol
		}
	}
	return synthArray
}

//Generates a saw wave
func (t *ToneGenerator) SawTone(freq, seconds float64, vol int32) []int32{
	var synthArray = make([]int32, int(seconds*t.sampleRate)) //duration/step = dur*sR
	//Find one period
	//Width: first half = 1, second = 0
	//Repeat n.m times
	for i:=0; i < len(synthArray); i++{
		synthArray[i] = int32(float64(vol)*math.Sin(freq *2* math.Pi * float64(i) * t.step))
	}
	return synthArray
}

//Generates a triangle wave
func (t *ToneGenerator) TriTone(freq, seconds float64, vol int32) []int32{
	var synthArray = make([]int32, int(seconds*t.sampleRate)) //duration/step = dur*sR
	//Find one period
	//Width: first half = 1, second = 0
	//Repeat n.m times
	for i:=0; i < len(synthArray); i++{
		synthArray[i] = int32(float64(vol)*math.Sin(freq *2* math.Pi * float64(i) * t.step))
	}
	return synthArray
}

func period(freq float64) float64{
	return 1.0/freq
}<commit_msg>Saw should work<commit_after>package main
import("math")

type ToneGenerator struct {
	sampleRate float64
	step float64 //delta t
}

func NewToneGenerator(sampleRate float64) *ToneGenerator {
	if sampleRate < 1 {
		return nil
	}
	t := new(ToneGenerator)
	t.sampleRate = sampleRate
	t.step = 1/sampleRate
	return t
}

func (t *ToneGenerator) Tone(freq, seconds float64, vol int32) []int32{
	return t.SinTone(freq, seconds, vol)
}

//Generates a sin wave of a certain freq for a specific volume
func (t *ToneGenerator) SinTone(freq, seconds float64, vol int32) []int32{
	var synthArray = make([]int32, int(seconds*t.sampleRate)) //duration/step = dur*sR
	for i:=0; i < len(synthArray); i++{
		synthArray[i] = int32(float64(vol)*math.Sin(freq *2* math.Pi * float64(i) * t.step))
	}
	return synthArray
}

//Generates a square wave
func (t *ToneGenerator) SquareTone(freq, seconds float64, vol int32) []int32{
	var synthArray = make([]int32, int(seconds*t.sampleRate)) //duration/step = dur*sR
	var period, samplesPerPeriod = period(freq)

	for i:=0; i < len(synthArray); i++{
		if i%samplesPerPeriod < samplesPerPeriod/2 { //Width: first half = 1, second = 0
			synthArray[i] = vol
		} else {
			synthArray[i] = -vol
		}
	}
	return synthArray
}

//Generates a saw wave
func (t *ToneGenerator) SawTone(freq, seconds float64, vol int32) []int32{
	var synthArray = make([]int32, int(seconds*t.sampleRate)) //duration/step = dur*sR
	var period, samplesPerPeriod = period(freq)

	//from -vol to vol linearly for one period
	var ramp = 2*vol/samplesPerPeriod
	var val = -vol
	for i:=0; i < len(synthArray); i++{
		if val < samplesPerPeriod{
			val = val + ramp
		} else{ //Reset
			val = -vol
		}
		synthArray[i] = val
	}
	return synthArray
}

//Generates a triangle wave
func (t *ToneGenerator) TriTone(freq, seconds float64, vol int32) []int32{
	var synthArray = make([]int32, int(seconds*t.sampleRate)) //duration/step = dur*sR
	var period, samplesPerPeriod = period(freq)
	//Find one period
	//Width: first half = 1, second = 0
	//Repeat n.m times
	for i:=0; i < len(synthArray); i++{
		synthArray[i] = int32(float64(vol)*math.Sin(freq *2* math.Pi * float64(i) * t.step))
	}
	return synthArray
}

func period(freq float64) (float64, int){
	return 1.0/freq, int(period/t.step)
}<|endoftext|>


<commit_before>// Package rwcount implements a simple counter that wraps an io.Reader or io.Writer.
// Useful for functions (like binary.Read/Write) which do not return read/write counts.
package rwcount

import "io"

// CountReader is used to wrap a io.Reader for counting
type CountReader struct {
	io.Reader
	bytesRead int64
	StickyErr bool
	Err       error
}

// Read implements the io.Reader interface
func (c *CountReader) Read(d []byte) (int, error) {
	if c.StickyErr && c.Err != nil {
		return 0, c.Err
	}
	total, err := c.Reader.Read(d)
	c.bytesRead += int64(total)
	c.Err = err
	return total, err
}

// BytesRead returns the number of bytes read
func (c CountReader) BytesRead() int64 {
	return c.bytesRead
}

// CountWriter is used to wrap a io.Writer for counting
type CountWriter struct {
	io.Writer
	bytesWritten int64
	StickyErr    bool
	Err          error
}

// Write implements the io.Writer interface
func (c *CountWriter) Write(d []byte) (int, error) {
	if c.StickyErr && c.Err != nil {
		return 0, c.Err
	}
	total, err := c.Writer.Write(d)
	c.bytesWritten += int64(total)
	c.Err = err
	return
}

// BytesWritten returns the number of bytes written
func (c CountWriter) BytesWritten() int64 {
	return c.bytesWritten
}
<commit_msg>StickyErr removed - err is sticky by default<commit_after>// Package rwcount implements a simple counter that wraps an io.Reader or io.Writer.
// Useful for functions (like binary.Read/Write) which do not return read/write counts.
package rwcount

import "io"

// CountReader is used to wrap a io.Reader for counting
type CountReader struct {
	io.Reader
	bytesRead int64
	Err       error
}

// Read implements the io.Reader interface
func (c *CountReader) Read(d []byte) (int, error) {
	if c.Err != nil {
		return 0, c.Err
	}
	total, err := c.Reader.Read(d)
	c.bytesRead += int64(total)
	c.Err = err
	return total, err
}

// BytesRead returns the number of bytes read
func (c CountReader) BytesRead() int64 {
	return c.bytesRead
}

// CountWriter is used to wrap a io.Writer for counting
type CountWriter struct {
	io.Writer
	bytesWritten int64
	Err          error
}

// Write implements the io.Writer interface
func (c *CountWriter) Write(d []byte) (int, error) {
	if c.Err != nil {
		return 0, c.Err
	}
	total, err := c.Writer.Write(d)
	c.bytesWritten += int64(total)
	c.Err = err
	return total, err
}

// BytesWritten returns the number of bytes written
func (c CountWriter) BytesWritten() int64 {
	return c.bytesWritten
}
<|endoftext|>