<commit_before>package postgres

import (
	"database/sql"
	"fmt"
	. "github.com/aktau/gomig/db/common"
	"log"
	"strings"
)

var PG_W_VERBOSE = true

var (
	postgresInit = []string{
		"SET client_encoding = 'UTF8';",
		"SET standard_conforming_strings = off;",
		"SET check_function_bodies = false;",
		"SET client_min_messages = warning;",
	}
)

const (
	explainQuery = `
SELECT col.column_name AS field,
       CASE
        WHEN col.character_maximum_length IS NOT NULL THEN col.data_type || '(' || col.character_maximum_length || ')'
        ELSE col.data_type
       END AS type,
       col.is_nullable AS null,
       CASE
        WHEN tc.constraint_type = 'PRIMARY KEY' THEN 'PRI'
        ELSE ''
       END AS key,
       '' AS default,
       '' AS extra
       --kcu.constraint_name AS constraint_name
       --kcu.*,
       --tc.*
FROM   information_schema.columns col
LEFT JOIN   information_schema.key_column_usage kcu ON (kcu.table_name = col.table_name AND kcu.column_name = col.column_name)
LEFT JOIN   information_schema.table_constraints AS tc ON (kcu.constraint_name = tc.constraint_name)
WHERE  col.table_name = '%v'
ORDER BY col.ordinal_position;`
)

type genericPostgresWriter struct {
	e               Executor
	insertBulkLimit int
}

func (w *genericPostgresWriter) bulkTransfer(src *Table, dstName string, rows *sql.Rows) (err error) {
	ex := w.e

	colnames := make([]string, 0, len(src.Columns))
	for _, col := range src.Columns {
		colnames = append(colnames, col.Name)
	}

	if err = ex.BulkInit(dstName, colnames...); err != nil {
		return
	}
	defer func() {
		berr := ex.BulkFinish()
		if err == nil {
			/* if there was no earlier error, set the one from BulkFinish */
			err = berr
		}
	}()

	/* create a slice with the right types to extract into, and let the SQL
	 * driver take care of the conversion */
	vals := NewTypedSlice(src)

	for rows.Next() {
		if err = rows.Scan(vals...); err != nil {
			return fmt.Errorf("postgres: error while reading from source:", err)
		}

		if err = ex.BulkAddRecord(vals...); err != nil {
			return fmt.Errorf("postgres: error during bulk insert:", err)
		}
	}

	return
}

func (w *genericPostgresWriter) normalTransfer(src *Table, dstName string, rows *sql.Rows) error {
	/* an alternate way to do this, with type assertions
	 * but possibly less accurately: http://go-database-sql.org/varcols.html */
	pointers := make([]interface{}, len(src.Columns))
	containers := make([]sql.RawBytes, len(src.Columns))
	for i, _ := range pointers {
		pointers[i] = &containers[i]
	}
	stringrep := make([]string, 0, len(src.Columns))
	insertLines := make([]string, 0, 32)
	for rows.Next() {
		err := rows.Scan(pointers...)
		if err != nil {
			log.Println("postgres: error while reading from source:", err)
			return err
		}

		for idx, val := range containers {
			str, err := RawToPostgres(val, src.Columns[idx].Type)
			if err != nil {
				return err
			}
			stringrep = append(stringrep, str)
		}

		insertLines = append(insertLines, "("+strings.Join(stringrep, ",")+")")
		stringrep = stringrep[:0]

		if len(insertLines) >= w.insertBulkLimit {
			err = w.e.Submit(fmt.Sprintf("INSERT INTO %v VALUES\n\t%v;\n",
				dstName, strings.Join(insertLines, ",\n\t")))
			if err != nil {
				return err
			}

			insertLines = insertLines[:0]
		}
	}

	if len(insertLines) > 0 {
		err := w.e.Submit(fmt.Sprintf("INSERT INTO %v VALUES\n\t%v;\n",
			dstName, strings.Join(insertLines, ",\n\t")))
		if err != nil {
			return err
		}
	}

	return nil
}

func (w *genericPostgresWriter) transferTable(src *Table, dstName string, r Reader) error {
	/* bulk insert values */
	rows, err := r.Read(src)
	if err != nil {
		return err
	}
	defer rows.Close()

	if PG_W_VERBOSE {
		log.Print("postgres: query done, scanning rows...")
	}

	if w.e.HasCapability(CapBulkTransfer) {
		if PG_W_VERBOSE {
			log.Print("postgres: bulk capability detected, performing bulk transfer...")
		}

		err = w.bulkTransfer(src, dstName, rows)
	} else {
		if PG_W_VERBOSE {
			log.Print("postgres: no bulk capability detected, performing normal transfer...")
		}

		err = w.normalTransfer(src, dstName, rows)
	}
	if err != nil {
		return err
	}

	return rows.Err()
}

/* how to do an UPSERT/MERGE in PostgreSQL
 * http://stackoverflow.com/questions/17267417/how-do-i-do-an-upsert-merge-insert-on-duplicate-update-in-postgresq */
func (w *genericPostgresWriter) MergeTable(src *Table, dstName, extraDstCond string, r Reader) error {
	tmpName := "gomig_tmp"

	err := w.e.Begin(fmt.Sprintf("merge table %v into table %v",
		src.Name, dstName))
	if err != nil {
		return err
	}

	/* create temporary table */
	err = w.e.Submit(fmt.Sprintf("CREATE TEMPORARY TABLE %v (\n\t%v\n)\nON COMMIT DROP;\n", tmpName, ColumnsSql(src)))
	if err != nil {
		return err
	}

	if PG_W_VERBOSE {
		log.Println("postgres: preparing to read values from source db")
	}

	err = w.transferTable(src, tmpName, r)
	if err != nil {
		return err
	}

	if PG_W_VERBOSE {
		log.Print("postgres: rowscan done, creating merge statements")
	}

	/* analyze the temp table, for performance */
	err = w.e.Submit(fmt.Sprintf("ANALYZE %v;\n", tmpName))
	if err != nil {
		return err
	}

	/* lock the target table */
	err = w.e.Submit(fmt.Sprintf("LOCK TABLE %v IN EXCLUSIVE MODE;", dstName))
	if err != nil {
		return err
	}

	colnames := make([]string, 0, len(src.Columns))
	srccol := make([]string, 0, len(src.Columns))
	pkWhere := make([]string, 0, len(src.Columns))
	pkIsNull := make([]string, 0, len(src.Columns))
	colassign := make([]string, 0, len(src.Columns))
	for _, col := range src.Columns {
		colnames = append(colnames, col.Name)
		srccol = append(srccol, "src."+col.Name)
		if col.PrimaryKey {
			pkWhere = append(pkWhere, fmt.Sprintf("dst.%[1]v = src.%[1]v", col.Name))
			pkIsNull = append(pkIsNull, fmt.Sprintf("dst.%[1]v IS NULL", col.Name))
		} else {
			colassign = append(colassign, fmt.Sprintf("%[1]v = src.%[1]v", col.Name))
		}
	}
	pkWherePart := strings.Join(pkWhere, "\nAND    ")
	pkIsNullPart := strings.Join(pkIsNull, "\nAND    ")
	srccolPart := strings.Join(srccol, ",\n       ")

	/* UPDATE from temp table to target table based on PK */
	err = w.e.Submit(fmt.Sprintf(`
UPDATE %v AS dst
SET    %v
FROM   %v AS src
WHERE  %v;`, dstName, strings.Join(colassign, ",\n       "), tmpName, pkWherePart))
	if err != nil {
		return err
	}

	/* if there is an extra condition, make sure it attaches cleanly to the
	 * rest of the query. */
	if extraDstCond != "" {
		extraDstCond = strings.TrimSpace(extraDstCond)
		if !strings.HasPrefix(extraDstCond, "and") &&
			!strings.HasPrefix(extraDstCond, "AND") &&
			!strings.HasPrefix(extraDstCond, "or") &&
			!strings.HasPrefix(extraDstCond, "OR") {
			extraDstCond = "AND " + extraDstCond
		}
		extraDstCond = "\n" + extraDstCond
	}

	/* INSERT from temp table to target table based on PK */
	err = w.e.Submit(fmt.Sprintf(`
INSERT INTO %[1]v (%[3]v)
SELECT %[4]v
FROM   %[2]v AS src
LEFT OUTER JOIN %[1]v AS dst ON (
       %[5]v
)
WHERE  %[6]v%[7]v;
`, dstName, tmpName, strings.Join(colnames, ", "), srccolPart,
		pkWherePart, pkIsNullPart, extraDstCond))
	if err != nil {
		return err
	}

	if PG_W_VERBOSE {
		log.Print("postgres: statements completed, executing transaction")
	}

	err = w.e.Commit()
	return err
}

func (w *genericPostgresWriter) Close() error {
	return w.e.Close()
}

type PostgresWriter struct {
	genericPostgresWriter
}

func NewPostgresWriter(conf *Config) (*PostgresWriter, error) {
	db, err := openDB(conf)
	if err != nil {
		return nil, err
	}

	executor, err := NewPgDbExecutor(db)
	if err != nil {
		db.Close()
		return nil, err
	}

	errors := executor.Multiple("initializing DB connection (WARNING: connection pooling might mess with this)", postgresInit)
	if len(errors) > 0 {
		executor.Close()
		for _, err := range errors {
			log.Println("postgres error:", err)
		}
		return nil, errors[0]
	}

	return &PostgresWriter{genericPostgresWriter{executor, 64}}, nil
}

type PostgresFileWriter struct {
	genericPostgresWriter
}

func NewPostgresFileWriter(filename string) (*PostgresFileWriter, error) {
	executor, err := NewFileExecutor(filename)
	if err != nil {
		return nil, err
	}

	errors := executor.Multiple("initializing DB connection", postgresInit)
	if len(errors) > 0 {
		executor.Close()
		for _, err := range errors {
			log.Println("postgres error:", err)
		}
		return nil, errors[0]
	}

	return &PostgresFileWriter{genericPostgresWriter{executor, 256}}, err
}

func ColumnsSql(table *Table) string {
	colSql := make([]string, 0, len(table.Columns))

	for _, col := range table.Columns {
		colSql = append(colSql, fmt.Sprintf("%v %v", col.Name, GenericToPostgresType(col.Type)))
	}

	pkCols := make([]string, 0, len(table.Columns))
	for _, col := range table.Columns {
		if col.PrimaryKey {
			pkCols = append(pkCols, col.Name)
		}
	}

	/* add the primary key */
	colSql = append(colSql, fmt.Sprintf("PRIMARY KEY (%v)",
		strings.Join(pkCols, ", ")))

	return strings.Join(colSql, ",\n\t")
}
<commit_msg>pg/writer: don't try to update all-PK tables<commit_after>package postgres

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

	. "github.com/aktau/gomig/db/common"
)

var PG_W_VERBOSE = true

var (
	postgresInit = []string{
		"SET client_encoding = 'UTF8';",
		"SET standard_conforming_strings = off;",
		"SET check_function_bodies = false;",
		"SET client_min_messages = warning;",
	}
)

const (
	explainQuery = `
SELECT col.column_name AS field,
       CASE
        WHEN col.character_maximum_length IS NOT NULL THEN col.data_type || '(' || col.character_maximum_length || ')'
        ELSE col.data_type
       END AS type,
       col.is_nullable AS null,
       CASE
        WHEN tc.constraint_type = 'PRIMARY KEY' THEN 'PRI'
        ELSE ''
       END AS key,
       '' AS default,
       '' AS extra
       --kcu.constraint_name AS constraint_name
       --kcu.*,
       --tc.*
FROM   information_schema.columns col
LEFT JOIN   information_schema.key_column_usage kcu ON (kcu.table_name = col.table_name AND kcu.column_name = col.column_name)
LEFT JOIN   information_schema.table_constraints AS tc ON (kcu.constraint_name = tc.constraint_name)
WHERE  col.table_name = '%v'
ORDER BY col.ordinal_position;`
)

type genericPostgresWriter struct {
	e               Executor
	insertBulkLimit int
}

func (w *genericPostgresWriter) bulkTransfer(src *Table, dstName string, rows *sql.Rows) (err error) {
	ex := w.e

	colnames := make([]string, 0, len(src.Columns))
	for _, col := range src.Columns {
		colnames = append(colnames, col.Name)
	}

	if err = ex.BulkInit(dstName, colnames...); err != nil {
		return
	}
	defer func() {
		berr := ex.BulkFinish()
		if err == nil {
			/* if there was no earlier error, set the one from BulkFinish */
			err = berr
		}
	}()

	/* create a slice with the right types to extract into, and let the SQL
	 * driver take care of the conversion */
	vals := NewTypedSlice(src)

	for rows.Next() {
		if err = rows.Scan(vals...); err != nil {
			return fmt.Errorf("postgres: error while reading from source:", err)
		}

		if err = ex.BulkAddRecord(vals...); err != nil {
			return fmt.Errorf("postgres: error during bulk insert:", err)
		}
	}

	return
}

func (w *genericPostgresWriter) normalTransfer(src *Table, dstName string, rows *sql.Rows) error {
	/* an alternate way to do this, with type assertions
	 * but possibly less accurately: http://go-database-sql.org/varcols.html */
	pointers := make([]interface{}, len(src.Columns))
	containers := make([]sql.RawBytes, len(src.Columns))
	for i, _ := range pointers {
		pointers[i] = &containers[i]
	}
	stringrep := make([]string, 0, len(src.Columns))
	insertLines := make([]string, 0, 32)
	for rows.Next() {
		err := rows.Scan(pointers...)
		if err != nil {
			log.Println("postgres: error while reading from source:", err)
			return err
		}

		for idx, val := range containers {
			str, err := RawToPostgres(val, src.Columns[idx].Type)
			if err != nil {
				return err
			}
			stringrep = append(stringrep, str)
		}

		insertLines = append(insertLines, "("+strings.Join(stringrep, ",")+")")
		stringrep = stringrep[:0]

		if len(insertLines) >= w.insertBulkLimit {
			err = w.e.Submit(fmt.Sprintf("INSERT INTO %v VALUES\n\t%v;\n",
				dstName, strings.Join(insertLines, ",\n\t")))
			if err != nil {
				return err
			}

			insertLines = insertLines[:0]
		}
	}

	if len(insertLines) > 0 {
		err := w.e.Submit(fmt.Sprintf("INSERT INTO %v VALUES\n\t%v;\n",
			dstName, strings.Join(insertLines, ",\n\t")))
		if err != nil {
			return err
		}
	}

	return nil
}

func (w *genericPostgresWriter) transferTable(src *Table, dstName string, r Reader) error {
	/* bulk insert values */
	rows, err := r.Read(src)
	if err != nil {
		return err
	}
	defer rows.Close()

	if PG_W_VERBOSE {
		log.Print("postgres: query done, scanning rows...")
	}

	if w.e.HasCapability(CapBulkTransfer) {
		if PG_W_VERBOSE {
			log.Print("postgres: bulk capability detected, performing bulk transfer...")
		}

		err = w.bulkTransfer(src, dstName, rows)
	} else {
		if PG_W_VERBOSE {
			log.Print("postgres: no bulk capability detected, performing normal transfer...")
		}

		err = w.normalTransfer(src, dstName, rows)
	}
	if err != nil {
		return err
	}

	return rows.Err()
}

/* how to do an UPSERT/MERGE in PostgreSQL
 * http://stackoverflow.com/questions/17267417/how-do-i-do-an-upsert-merge-insert-on-duplicate-update-in-postgresq */
func (w *genericPostgresWriter) MergeTable(src *Table, dstName, extraDstCond string, r Reader) error {
	tmpName := "gomig_tmp"

	err := w.e.Begin(fmt.Sprintf("merge table %v into table %v",
		src.Name, dstName))
	if err != nil {
		return err
	}

	/* create temporary table */
	err = w.e.Submit(fmt.Sprintf("CREATE TEMPORARY TABLE %v (\n\t%v\n)\nON COMMIT DROP;\n", tmpName, ColumnsSql(src)))
	if err != nil {
		return err
	}

	if PG_W_VERBOSE {
		log.Println("postgres: preparing to read values from source db")
	}

	err = w.transferTable(src, tmpName, r)
	if err != nil {
		return err
	}

	if PG_W_VERBOSE {
		log.Print("postgres: rowscan done, creating merge statements")
	}

	/* analyze the temp table, for performance */
	err = w.e.Submit(fmt.Sprintf("ANALYZE %v;\n", tmpName))
	if err != nil {
		return err
	}

	/* lock the target table */
	err = w.e.Submit(fmt.Sprintf("LOCK TABLE %v IN EXCLUSIVE MODE;", dstName))
	if err != nil {
		return err
	}

	colnames := make([]string, 0, len(src.Columns))
	srccol := make([]string, 0, len(src.Columns))
	pkWhere := make([]string, 0, len(src.Columns))
	pkIsNull := make([]string, 0, len(src.Columns))
	colassign := make([]string, 0, len(src.Columns))
	for _, col := range src.Columns {
		colnames = append(colnames, col.Name)
		srccol = append(srccol, "src."+col.Name)
		if col.PrimaryKey {
			pkWhere = append(pkWhere, fmt.Sprintf("dst.%[1]v = src.%[1]v", col.Name))
			pkIsNull = append(pkIsNull, fmt.Sprintf("dst.%[1]v IS NULL", col.Name))
		} else {
			colassign = append(colassign, fmt.Sprintf("%[1]v = src.%[1]v", col.Name))
		}
	}
	pkWherePart := strings.Join(pkWhere, "\nAND    ")
	pkIsNullPart := strings.Join(pkIsNull, "\nAND    ")
	srccolPart := strings.Join(srccol, ",\n       ")

	/* it's possible this table is all primary key (as far as we know),
	 * in which case we don't need the update part */
	if len(colassign) != 0 {
		/* UPDATE from temp table to target table based on PK */
		err = w.e.Submit(fmt.Sprintf(`
UPDATE %v AS dst
SET    %v
FROM   %v AS src
WHERE  %v;`, dstName, strings.Join(colassign, ",\n       "), tmpName, pkWherePart))
		if err != nil {
			return err
		}
	}

	/* if there is an extra condition, make sure it attaches cleanly to the
	 * rest of the query. */
	if extraDstCond != "" {
		extraDstCond = strings.TrimSpace(extraDstCond)
		if !strings.HasPrefix(extraDstCond, "and") &&
			!strings.HasPrefix(extraDstCond, "AND") &&
			!strings.HasPrefix(extraDstCond, "or") &&
			!strings.HasPrefix(extraDstCond, "OR") {
			extraDstCond = "AND " + extraDstCond
		}
		extraDstCond = "\n" + extraDstCond
	}

	/* INSERT from temp table to target table based on PK */
	err = w.e.Submit(fmt.Sprintf(`
INSERT INTO %[1]v (%[3]v)
SELECT %[4]v
FROM   %[2]v AS src
LEFT OUTER JOIN %[1]v AS dst ON (
       %[5]v
)
WHERE  %[6]v%[7]v;
`, dstName, tmpName, strings.Join(colnames, ", "), srccolPart,
		pkWherePart, pkIsNullPart, extraDstCond))
	if err != nil {
		return err
	}

	if PG_W_VERBOSE {
		log.Print("postgres: statements completed, executing transaction")
	}

	err = w.e.Commit()
	return err
}

func (w *genericPostgresWriter) Close() error {
	return w.e.Close()
}

type PostgresWriter struct {
	genericPostgresWriter
}

func NewPostgresWriter(conf *Config) (*PostgresWriter, error) {
	db, err := openDB(conf)
	if err != nil {
		return nil, err
	}

	executor, err := NewPgDbExecutor(db)
	if err != nil {
		db.Close()
		return nil, err
	}

	errors := executor.Multiple("initializing DB connection (WARNING: connection pooling might mess with this)", postgresInit)
	if len(errors) > 0 {
		executor.Close()
		for _, err := range errors {
			log.Println("postgres error:", err)
		}
		return nil, errors[0]
	}

	return &PostgresWriter{genericPostgresWriter{executor, 64}}, nil
}

type PostgresFileWriter struct {
	genericPostgresWriter
}

func NewPostgresFileWriter(filename string) (*PostgresFileWriter, error) {
	executor, err := NewFileExecutor(filename)
	if err != nil {
		return nil, err
	}

	errors := executor.Multiple("initializing DB connection", postgresInit)
	if len(errors) > 0 {
		executor.Close()
		for _, err := range errors {
			log.Println("postgres error:", err)
		}
		return nil, errors[0]
	}

	return &PostgresFileWriter{genericPostgresWriter{executor, 256}}, err
}

func ColumnsSql(table *Table) string {
	colSql := make([]string, 0, len(table.Columns))

	for _, col := range table.Columns {
		colSql = append(colSql, fmt.Sprintf("%v %v", col.Name, GenericToPostgresType(col.Type)))
	}

	pkCols := make([]string, 0, len(table.Columns))
	for _, col := range table.Columns {
		if col.PrimaryKey {
			pkCols = append(pkCols, col.Name)
		}
	}

	/* add the primary key */
	colSql = append(colSql, fmt.Sprintf("PRIMARY KEY (%v)",
		strings.Join(pkCols, ", ")))

	return strings.Join(colSql, ",\n\t")
}
<|endoftext|>


<commit_before>package kubernetes

import (
	"fmt"
	"strings"

	"golang.org/x/net/context"
	"k8s.io/kubernetes/pkg/api"
	client "k8s.io/kubernetes/pkg/client/unversioned"

	"gitlab.com/gitlab-org/gitlab-ci-multi-runner/common"
	"gitlab.com/gitlab-org/gitlab-ci-multi-runner/executors"
)

const (
	defaultHelperImage = "munnerz/gitlab-runner-helper"
)

var (
	executorOptions = executors.ExecutorOptions{
		SharedBuildsDir: false,
		Shell: common.ShellScriptInfo{
			Shell:         "bash",
			Type:          common.NormalShell,
			RunnerCommand: "/usr/bin/gitlab-runner-helper",
		},
		ShowHostname:     true,
		SupportedOptions: []string{"image", "services", "artifacts", "cache"},
	}
)

type kubernetesOptions struct {
	Image    string   `json:"image"`
	Services []string `json:"services"`
}

type executor struct {
	executors.AbstractExecutor

	kubeClient   *client.Client
	prepod       *api.Pod
	pod          *api.Pod
	options      *kubernetesOptions
	extraOptions Options

	buildLimits   api.ResourceList
	serviceLimits api.ResourceList
}

func (s *executor) Prepare(globalConfig *common.Config, config *common.RunnerConfig, build *common.Build) error {
	err := s.AbstractExecutor.Prepare(globalConfig, config, build)
	if err != nil {
		return err
	}

	if s.BuildShell.PassFile {
		return fmt.Errorf("Kubernetes doesn't support shells that require script file")
	}

	s.extraOptions = DefaultOptions{s.Build.Variables}

	if !s.Config.Kubernetes.AllowPrivileged && s.extraOptions.Privileged() {
		return fmt.Errorf("Runner does not allow privileged containers")
	}

	err = build.Options.Decode(&s.options)
	if err != nil {
		return err
	}

	s.kubeClient, err = getKubeClient(config.Kubernetes)
	if err != nil {
		return fmt.Errorf("Error connecting to Kubernetes: %s", err.Error())
	}

	if s.serviceLimits, err = limits(s.Config.Kubernetes.ServiceCPUs, s.Config.Kubernetes.ServiceMemory); err != nil {
		return err
	}

	if s.buildLimits, err = limits(s.Config.Kubernetes.CPUs, s.Config.Kubernetes.Memory); err != nil {
		return err
	}

	if s.options.Image == "" {
		switch len(s.Config.Kubernetes.Image) {
		case 0:
			return fmt.Errorf("no image specified and no default set in config")
		default:
			s.options.Image = s.Config.Kubernetes.Image
		}
	}

	if len(s.Config.Kubernetes.HelperImage) == 0 {
		s.Config.Kubernetes.HelperImage = defaultHelperImage
	}

	if len(s.Config.Kubernetes.Namespace) == 0 {
		s.Config.Kubernetes.Namespace = "default"
	}

	s.Println("Using Kubernetes executor with image", s.options.Image, "...")

	return nil
}

func (s *executor) Run(cmd common.ExecutorCommand) error {
	s.Debugln("Starting Kubernetes command...")

	if s.pod == nil {
		services := make([]api.Container, len(s.options.Services))
		for i, image := range s.options.Services {
			resolvedImage := s.Build.GetAllVariables().ExpandValue(image)
			services[i] = s.buildContainer(fmt.Sprintf("svc-%d", i), resolvedImage, s.serviceLimits)
		}

		buildImage := s.Build.GetAllVariables().ExpandValue(s.options.Image)
		pod, err := s.kubeClient.Pods(s.Config.Kubernetes.Namespace).Create(&api.Pod{
			ObjectMeta: api.ObjectMeta{
				GenerateName: s.Build.ProjectUniqueName(),
				Namespace:    s.Config.Kubernetes.Namespace,
			},
			Spec: api.PodSpec{
				Volumes: []api.Volume{
					api.Volume{
						Name: "repo",
						VolumeSource: api.VolumeSource{
							EmptyDir: &api.EmptyDirVolumeSource{},
						},
					},
					api.Volume{
						Name: "etc-ssl-certs",
						VolumeSource: api.VolumeSource{
							HostPath: &api.HostPathVolumeSource{
								Path: "/etc/ssl/certs",
							},
						},
					},
					api.Volume{
						Name: "usr-share-ca-certificates",
						VolumeSource: api.VolumeSource{
							HostPath: &api.HostPathVolumeSource{
								Path: "/usr/share/ca-certificates",
							},
						},
					},
				},
				RestartPolicy: api.RestartPolicyNever,
				Containers: append([]api.Container{
					s.buildContainer("build", buildImage, s.buildLimits, s.BuildShell.DockerCommand...),
					s.buildContainer("pre", s.Config.Kubernetes.HelperImage, s.serviceLimits, s.BuildShell.DockerCommand...),
				}, services...),
			},
		})

		if err != nil {
			return err
		}

		s.pod = pod
	}

	var containerName string
	switch {
	case cmd.Predefined:
		containerName = "pre"
	default:
		containerName = "build"
	}

	ctx, cancel := context.WithCancel(context.Background())
	select {
	case err := <-s.runInContainer(ctx, containerName, cmd.Script):
		return err
	case <-cmd.Abort:
		cancel()
		return fmt.Errorf("build aborted")
	}
}

func (s *executor) Cleanup() {
	if s.pod != nil {
		err := s.kubeClient.Pods(s.pod.Namespace).Delete(s.pod.Name, nil)
		if err != nil {
			s.Errorln(fmt.Sprintf("Error cleaning up pod: %s", err.Error()))
		}
	}
	s.AbstractExecutor.Cleanup()
}

func (s *executor) buildContainer(name, image string, limits api.ResourceList, command ...string) api.Container {
	path := strings.Split(s.Build.BuildDir, "/")
	path = path[:len(path)-1]

	privileged := s.extraOptions.Privileged()

	return api.Container{
		Name:    name,
		Image:   image,
		Command: command,
		Env:     buildVariables(s.Build.GetAllVariables().PublicOrInternal()),
		Resources: api.ResourceRequirements{
			Limits: limits,
		},
		VolumeMounts: []api.VolumeMount{
			api.VolumeMount{
				Name:      "repo",
				MountPath: strings.Join(path, "/"),
			},
			api.VolumeMount{
				Name:      "etc-ssl-certs",
				MountPath: "/etc/ssl/certs",
			},
			api.VolumeMount{
				Name:      "usr-share-ca-certificates",
				MountPath: "/usr/share/ca-certificates",
			},
		},
		SecurityContext: &api.SecurityContext{
			Privileged: &privileged,
		},
		Stdin: true,
	}
}

func (s *executor) runInContainer(ctx context.Context, name, command string) <-chan error {
	errc := make(chan error, 1)
	go func() {
		defer close(errc)

		status, err := waitForPodRunning(ctx, s.kubeClient, s.pod, s.BuildTrace)

		if err != nil {
			errc <- err
			return
		}

		if status != api.PodRunning {
			errc <- fmt.Errorf("pod failed to enter running state: %s", status)
			return
		}

		config, err := getKubeClientConfig(s.Config.Kubernetes)

		if err != nil {
			errc <- err
			return
		}

		exec := ExecOptions{
			PodName:       s.pod.Name,
			Namespace:     s.pod.Namespace,
			ContainerName: name,
			Command:       s.BuildShell.DockerCommand,
			In:            strings.NewReader(command),
			Out:           s.BuildTrace,
			Err:           s.BuildTrace,
			Stdin:         true,
			Config:        config,
			Client:        s.kubeClient,
			Executor:      &DefaultRemoteExecutor{},
		}

		errc <- exec.Run()
	}()

	return errc
}

func createFn() common.Executor {
	return &executor{
		AbstractExecutor: executors.AbstractExecutor{
			ExecutorOptions: executorOptions,
		},
	}
}

func featuresFn(features *common.FeaturesInfo) {
	features.Variables = true
	features.Image = true
	features.Services = true
	features.Artifacts = true
	features.Cache = true
}

func init() {
	common.RegisterExecutor("kubernetes", executors.DefaultExecutorProvider{
		Creator:         createFn,
		FeaturesUpdater: featuresFn,
	})
}
<commit_msg>Create setupBuildPod function<commit_after>package kubernetes

import (
	"fmt"
	"strings"

	"golang.org/x/net/context"
	"k8s.io/kubernetes/pkg/api"
	client "k8s.io/kubernetes/pkg/client/unversioned"

	"gitlab.com/gitlab-org/gitlab-ci-multi-runner/common"
	"gitlab.com/gitlab-org/gitlab-ci-multi-runner/executors"
)

const (
	defaultHelperImage = "munnerz/gitlab-runner-helper"
)

var (
	executorOptions = executors.ExecutorOptions{
		SharedBuildsDir: false,
		Shell: common.ShellScriptInfo{
			Shell:         "bash",
			Type:          common.NormalShell,
			RunnerCommand: "/usr/bin/gitlab-runner-helper",
		},
		ShowHostname:     true,
		SupportedOptions: []string{"image", "services", "artifacts", "cache"},
	}
)

type kubernetesOptions struct {
	Image    string   `json:"image"`
	Services []string `json:"services"`
}

type executor struct {
	executors.AbstractExecutor

	kubeClient   *client.Client
	prepod       *api.Pod
	pod          *api.Pod
	options      *kubernetesOptions
	extraOptions Options

	buildLimits   api.ResourceList
	serviceLimits api.ResourceList
}

func (s *executor) Prepare(globalConfig *common.Config, config *common.RunnerConfig, build *common.Build) error {
	err := s.AbstractExecutor.Prepare(globalConfig, config, build)
	if err != nil {
		return err
	}

	if s.BuildShell.PassFile {
		return fmt.Errorf("Kubernetes doesn't support shells that require script file")
	}

	s.extraOptions = DefaultOptions{s.Build.Variables}

	if !s.Config.Kubernetes.AllowPrivileged && s.extraOptions.Privileged() {
		return fmt.Errorf("Runner does not allow privileged containers")
	}

	err = build.Options.Decode(&s.options)
	if err != nil {
		return err
	}

	s.kubeClient, err = getKubeClient(config.Kubernetes)
	if err != nil {
		return fmt.Errorf("Error connecting to Kubernetes: %s", err.Error())
	}

	if s.serviceLimits, err = limits(s.Config.Kubernetes.ServiceCPUs, s.Config.Kubernetes.ServiceMemory); err != nil {
		return err
	}

	if s.buildLimits, err = limits(s.Config.Kubernetes.CPUs, s.Config.Kubernetes.Memory); err != nil {
		return err
	}

	if s.options.Image == "" {
		switch len(s.Config.Kubernetes.Image) {
		case 0:
			return fmt.Errorf("no image specified and no default set in config")
		default:
			s.options.Image = s.Config.Kubernetes.Image
		}
	}

	if len(s.Config.Kubernetes.HelperImage) == 0 {
		s.Config.Kubernetes.HelperImage = defaultHelperImage
	}

	if len(s.Config.Kubernetes.Namespace) == 0 {
		s.Config.Kubernetes.Namespace = "default"
	}

	s.Println("Using Kubernetes executor with image", s.options.Image, "...")

	return nil
}

func (s *executor) Run(cmd common.ExecutorCommand) error {
	s.Debugln("Starting Kubernetes command...")

	if s.pod == nil {
		err := s.setupBuildPod()

		if err != nil {
			return err
		}
	}

	var containerName string
	switch {
	case cmd.Predefined:
		containerName = "pre"
	default:
		containerName = "build"
	}

	ctx, cancel := context.WithCancel(context.Background())
	select {
	case err := <-s.runInContainer(ctx, containerName, cmd.Script):
		return err
	case <-cmd.Abort:
		cancel()
		return fmt.Errorf("build aborted")
	}
}

func (s *executor) Cleanup() {
	if s.pod != nil {
		err := s.kubeClient.Pods(s.pod.Namespace).Delete(s.pod.Name, nil)
		if err != nil {
			s.Errorln(fmt.Sprintf("Error cleaning up pod: %s", err.Error()))
		}
	}
	s.AbstractExecutor.Cleanup()
}

func (s *executor) buildContainer(name, image string, limits api.ResourceList, command ...string) api.Container {
	path := strings.Split(s.Build.BuildDir, "/")
	path = path[:len(path)-1]

	privileged := s.extraOptions.Privileged()

	return api.Container{
		Name:    name,
		Image:   image,
		Command: command,
		Env:     buildVariables(s.Build.GetAllVariables().PublicOrInternal()),
		Resources: api.ResourceRequirements{
			Limits: limits,
		},
		VolumeMounts: []api.VolumeMount{
			api.VolumeMount{
				Name:      "repo",
				MountPath: strings.Join(path, "/"),
			},
			api.VolumeMount{
				Name:      "etc-ssl-certs",
				MountPath: "/etc/ssl/certs",
			},
			api.VolumeMount{
				Name:      "usr-share-ca-certificates",
				MountPath: "/usr/share/ca-certificates",
			},
		},
		SecurityContext: &api.SecurityContext{
			Privileged: &privileged,
		},
		Stdin: true,
	}
}

func (s *executor) setupBuildPod() error {
	services := make([]api.Container, len(s.options.Services))
	for i, image := range s.options.Services {
		resolvedImage := s.Build.GetAllVariables().ExpandValue(image)
		services[i] = s.buildContainer(fmt.Sprintf("svc-%d", i), resolvedImage, s.serviceLimits)
	}

	buildImage := s.Build.GetAllVariables().ExpandValue(s.options.Image)
	pod, err := s.kubeClient.Pods(s.Config.Kubernetes.Namespace).Create(&api.Pod{
		ObjectMeta: api.ObjectMeta{
			GenerateName: s.Build.ProjectUniqueName(),
			Namespace:    s.Config.Kubernetes.Namespace,
		},
		Spec: api.PodSpec{
			Volumes: []api.Volume{
				api.Volume{
					Name: "repo",
					VolumeSource: api.VolumeSource{
						EmptyDir: &api.EmptyDirVolumeSource{},
					},
				},
				api.Volume{
					Name: "etc-ssl-certs",
					VolumeSource: api.VolumeSource{
						HostPath: &api.HostPathVolumeSource{
							Path: "/etc/ssl/certs",
						},
					},
				},
				api.Volume{
					Name: "usr-share-ca-certificates",
					VolumeSource: api.VolumeSource{
						HostPath: &api.HostPathVolumeSource{
							Path: "/usr/share/ca-certificates",
						},
					},
				},
			},
			RestartPolicy: api.RestartPolicyNever,
			Containers: append([]api.Container{
				s.buildContainer("build", buildImage, s.buildLimits, s.BuildShell.DockerCommand...),
				s.buildContainer("pre", s.Config.Kubernetes.HelperImage, s.serviceLimits, s.BuildShell.DockerCommand...),
			}, services...),
		},
	})

	if err != nil {
		return err
	}

	s.pod = pod

	return nil
}

func (s *executor) runInContainer(ctx context.Context, name, command string) <-chan error {
	errc := make(chan error, 1)
	go func() {
		defer close(errc)

		status, err := waitForPodRunning(ctx, s.kubeClient, s.pod, s.BuildTrace)

		if err != nil {
			errc <- err
			return
		}

		if status != api.PodRunning {
			errc <- fmt.Errorf("pod failed to enter running state: %s", status)
			return
		}

		config, err := getKubeClientConfig(s.Config.Kubernetes)

		if err != nil {
			errc <- err
			return
		}

		exec := ExecOptions{
			PodName:       s.pod.Name,
			Namespace:     s.pod.Namespace,
			ContainerName: name,
			Command:       s.BuildShell.DockerCommand,
			In:            strings.NewReader(command),
			Out:           s.BuildTrace,
			Err:           s.BuildTrace,
			Stdin:         true,
			Config:        config,
			Client:        s.kubeClient,
			Executor:      &DefaultRemoteExecutor{},
		}

		errc <- exec.Run()
	}()

	return errc
}

func createFn() common.Executor {
	return &executor{
		AbstractExecutor: executors.AbstractExecutor{
			ExecutorOptions: executorOptions,
		},
	}
}

func featuresFn(features *common.FeaturesInfo) {
	features.Variables = true
	features.Image = true
	features.Services = true
	features.Artifacts = true
	features.Cache = true
}

func init() {
	common.RegisterExecutor("kubernetes", executors.DefaultExecutorProvider{
		Creator:         createFn,
		FeaturesUpdater: featuresFn,
	})
}
<|endoftext|>


<commit_before>package main

import (
	"encoding/json"
	"fmt"
	"math/rand"
	"path/filepath"
	"strconv"
	"strings"
	"time"

	tgbotapi "github.com/go-telegram-bot-api/telegram-bot-api"
	bt "github.com/ikool-cn/gobeanstalk-connection-pool"
)

var (
	telegramBot *TelegramBot
)

// TelegramBot ...
type TelegramBot struct {
	Name          string
	SelfChatID    int64
	ChannelChatID int64
	ComicPath     string
	DeleteDelay   time.Duration
	Client        *tgbotapi.BotAPI
	Queue         *bt.Pool
	Tube          string
}

// NewTelegramBot ...
func NewTelegramBot(cfg *TelegramConfig, btdAddr string) (t *TelegramBot) {
	bot, err := tgbotapi.NewBotAPI(cfg.Token)
	if err != nil {
		logger.Panicf("tg bot init failed: %+v", err)
	}
	delay, err := time.ParseDuration(cfg.DeleteDelay)
	if err != nil {
		logger.Panicf("delete delay error: %+v", err)
	}

	t = &TelegramBot{
		Name:          bot.Self.UserName,
		SelfChatID:    cfg.SelfChatID,
		ChannelChatID: cfg.ChannelChatID,
		ComicPath:     cfg.ComicPath,
		DeleteDelay:   delay,
		Client:        bot,
		Tube:          "tg",
	}
	t.Queue = &bt.Pool{
		Dial: func() (*bt.Conn, error) {
			return bt.Dial(btdAddr)
		},
		MaxIdle:     10,
		MaxActive:   100,
		IdleTimeout: 60 * time.Second,
		MaxLifetime: 180 * time.Second,
		Wait:        true,
	}
	return
}

func (t *TelegramBot) putQueue(msg []byte) {
	conn, err := t.Queue.Get()
	if err != nil {
		logger.Errorf("%+v: %s", err, string(msg))
		return
	}
	conn.Use(t.Tube)
	_, err = conn.Put(msg, 1, t.DeleteDelay, time.Minute)
	if err != nil {
		logger.Errorf("%+v", err)
		return
	}
}

func (t *TelegramBot) send(chat int64, msg string) (tgbotapi.Message, error) {
	logger.Debugf("[%d]%s", chat, msg)
	return t.Client.Send(tgbotapi.NewMessage(chat, msg))
}

func (t *TelegramBot) delMessage() {
	for {
		conn, err := t.Queue.Get()
		if err != nil {
			logger.Errorf("%+v", err)
			time.Sleep(3 * time.Second)
			continue
		}
		conn.Watch(t.Tube)
		job, err := conn.Reserve()
		if err != nil {
			logger.Warningf("%+v", err)
			time.Sleep(3 * time.Second)
			continue
		}
		msg := &tgbotapi.Message{}
		err = json.Unmarshal(job.Body, msg)
		if err != nil {
			logger.Errorf("%+v", err)
			err = conn.Bury(job.ID, 0)
			if err != nil {
				logger.Errorf("%+v", err)
			}
			time.Sleep(3 * time.Second)
			continue
		}
		delMsg := tgbotapi.DeleteMessageConfig{
			ChatID:    msg.Chat.ID,
			MessageID: msg.MessageID,
		}
		logger.Infof(":[%s]{%s}", getMsgTitle(msg), strconv.Quote(msg.Text))

		_, err = t.Client.DeleteMessage(delMsg)
		if err != nil {
			logger.Errorf("%+v", err)
			err = conn.Bury(job.ID, 0)
			if err != nil {
				logger.Errorf("%+v", err)
			}
			time.Sleep(3 * time.Second)
			continue
		}
		err = conn.Delete(job.ID)
		if err != nil {
			logger.Errorf("%+v", err)
			time.Sleep(3 * time.Second)
		}
		t.Queue.Release(conn, false)
	}
}

func (t *TelegramBot) tgBot() {
	u := tgbotapi.NewUpdate(0)
	u.Timeout = 60
	for {
		updates, err := t.Client.GetUpdatesChan(u)
		if err != nil {
			logger.Errorf("%+v", err)
			time.Sleep(3 * time.Second)
			continue
		}
		var message *tgbotapi.Message
		for update := range updates {
			if update.Message != nil {
				message = update.Message
			} else if update.EditedMessage != nil {
				message = update.EditedMessage
			} else if update.CallbackQuery != nil {
				logger.Infof(
					"recv:(%s)[%s]reaction:{%s}",
					update.CallbackQuery.ChatInstance,
					update.CallbackQuery.From.String(),
					update.CallbackQuery.Data,
				)
				data := strings.SplitN(update.CallbackQuery.Data, ":", 2)
				switch data[0] {
				case "comic", "pic":
					go onReaction(t, update.CallbackQuery)
				}
				continue
			} else {
				continue
			}
			if message.Chat.IsGroup() {
				logger.Infof(
					"recv:(%s)[%s]{%s}",
					message.Chat.Title,
					message.From.String(),
					strconv.Quote(message.Text))
			} else {
				logger.Infof(
					"recv:[%s]{%s}",
					message.From.String(),
					strconv.Quote(message.Text),
				)
			}

			if message.IsCommand() {
				switch message.Command() {
				case "start":
					go onStart(t, message)
				case "comic":
					go onComic(t, message)
				case "pic":
					go onPic(t, message)
				default:
					logger.Infof("ignore unkown cmd: %+v", message.Command())
					continue
				}
			} else {
				if message.Text == "" {
					continue
				}
				checkRepeat(t, message)
			}
		}
		logger.Warning("tg bot restarted.")
		time.Sleep(3 * time.Second)
	}
}

func checkRepeat(t *TelegramBot, message *tgbotapi.Message) {
	key := "tg_last_" + strconv.FormatInt(message.Chat.ID, 10)
	flattendMsg := strings.TrimSpace(message.Text)
	defer redisClient.LTrim(key, 0, 10)
	defer redisClient.LPush(key, flattendMsg)

	lastMsgs, err := redisClient.LRange(key, 0, 6).Result()
	if err != nil {
		logger.Errorf("%+v", err)
		return
	}
	i := 0
	for _, s := range lastMsgs {
		if s == flattendMsg {
			i++
		}
	}
	if i > 1 {
		redisClient.Del(key)
		logger.Infof("repeat: %s", strconv.Quote(message.Text))
		msg := tgbotapi.NewMessage(message.Chat.ID, message.Text)
		t.Client.Send(msg)
	}
}

func onStart(t *TelegramBot, message *tgbotapi.Message) {
	msg := tgbotapi.NewMessage(message.Chat.ID, "呀呀呀")
	msg.ReplyToMessageID = message.MessageID
	t.Client.Send(msg)
}

func onComic(t *TelegramBot, message *tgbotapi.Message) {
	files, err := filepath.Glob(filepath.Join(t.ComicPath, "*.epub"))
	if err != nil {
		logger.Errorf("%+v", err)
		return
	}
	rand.Seed(time.Now().UnixNano())
	file := files[rand.Intn(len(files))]
	number := strings.Split(strings.Split(file, "@")[1], ".")[0]
	msg := tgbotapi.NewMessage(message.Chat.ID, "🔞 https://nhentai.net/g/"+number)

	msg.ReplyMarkup = buildInlineKeyboardMarkup("comic", number)

	logger.Infof("send:[%s]{%s}", getMsgTitle(message), strconv.Quote(file))
	msgSent, err := t.Client.Send(msg)
	if err != nil {
		logger.Errorf("%+v", err)
		return
	}
	data, err := json.Marshal(msgSent)
	if err != nil {
		logger.Errorf("%+v", err)
		return
	}
	t.putQueue(data)
}

func onPic(t *TelegramBot, message *tgbotapi.Message) {
	files, err := filepath.Glob(filepath.Join(twitterBot.ImgPath, "*"))
	if err != nil {
		logger.Errorf("%+v", err)
		return
	}
	if files == nil {
		logger.Error("find no pics")
	}
	rand.Seed(time.Now().UnixNano())
	file := files[rand.Intn(len(files))]

	logger.Infof("send:[%s]{%s}", getMsgTitle(message), strconv.Quote(file))

	msg := tgbotapi.NewDocumentUpload(message.Chat.ID, file)
	msg.ReplyMarkup = buildInlineKeyboardMarkup("pic", filepath.Base(file))

	_, err = t.Client.Send(msg)
	if err != nil {
		logger.Errorf("%+v", err)
	}
}

func onReaction(t *TelegramBot, callbackQuery *tgbotapi.CallbackQuery) {
	var callbackText string

	_type, _id, reaction, err := saveReaction(callbackQuery.Data, callbackQuery.From.ID)
	if err == nil {
		diss := redisClient.SCard(buildReactionKey(_type, _id, "diss")).Val()
		if diss <= 1 {
			msg := tgbotapi.NewEditMessageReplyMarkup(
				callbackQuery.Message.Chat.ID,
				callbackQuery.Message.MessageID,
				buildInlineKeyboardMarkup(_type, _id),
			)
			_, err = t.Client.Send(msg)
		} else {
			delMsg := tgbotapi.DeleteMessageConfig{
				ChatID:    callbackQuery.Message.Chat.ID,
				MessageID: callbackQuery.Message.MessageID,
			}
			_, err = t.Client.DeleteMessage(delMsg)
			if err == nil {
				err = probate(_type, _id)
			}
		}
	}

	if err != nil {
		logger.Debugf("%+v", err)
		callbackText = err.Error()
	} else {
		callbackText = reaction + " " + _id + "!"
	}

	callbackMsg := tgbotapi.NewCallback(callbackQuery.ID, callbackText)
	_, err = t.Client.AnswerCallbackQuery(callbackMsg)
	if err != nil {
		logger.Errorf("%+v", err)
	}
}

func getMsgTitle(m *tgbotapi.Message) string {
	if m.Chat.IsGroup() {
		return m.Chat.Title
	}
	return m.From.String()
}

func buildReactionData(_type, _id, reaction string) string {
	return _type + ":" + _id + ":" + reaction
}
func buildReactionKey(_type, _id, reaction string) string {
	return "reaction_" + buildReactionData(_type, _id, reaction)
}

func buildInlineKeyboardMarkup(_type, _id string) tgbotapi.InlineKeyboardMarkup {

	likeCount, _ := redisClient.SCard(buildReactionKey(_type, _id, "like")).Result()
	dissCount, _ := redisClient.SCard(buildReactionKey(_type, _id, "diss")).Result()

	likeText := "❤"
	if likeCount > 0 {
		likeText = likeText + " " + strconv.FormatInt(likeCount, 10)
	}
	dissText := "💔"
	if dissCount > 0 {
		dissText = dissText + " " + strconv.FormatInt(dissCount, 10)
	}

	row := tgbotapi.NewInlineKeyboardRow(
		tgbotapi.NewInlineKeyboardButtonData(likeText, buildReactionData(_type, _id, "like")),
		tgbotapi.NewInlineKeyboardButtonData(dissText, buildReactionData(_type, _id, "diss")),
	)
	return tgbotapi.NewInlineKeyboardMarkup(row)
}

func saveReaction(key string, user int) (_type, _id, reaction string, err error) {
	token := strings.Split(key, ":")
	if len(token) != 3 {
		err = fmt.Errorf("react data error: %s", key)
		return
	}
	_type = token[0]
	_id = token[1]
	reaction = token[2]

	pipe := redisClient.Pipeline()
	switch reaction {
	case "like":
		likeCount := pipe.SAdd(buildReactionKey(_type, _id, "like"), strconv.Itoa(user))
		dissCount := pipe.SRem(buildReactionKey(_type, _id, "diss"), strconv.Itoa(user))
		_, err = pipe.Exec()
		if err == nil {
			if likeCount.Val()+dissCount.Val() == 0 {
				err = fmt.Errorf("not modified")
			}
		}
	case "diss":
		dissCount := pipe.SAdd(buildReactionKey(_type, _id, "diss"), strconv.Itoa(user))
		likeCount := pipe.SRem(buildReactionKey(_type, _id, "like"), strconv.Itoa(user))
		_, err = pipe.Exec()
		if err == nil {
			if likeCount.Val()+dissCount.Val() == 0 {
				err = fmt.Errorf("not modified")
			}
		}
	default:
		err = fmt.Errorf("react type error: %s", key)
	}
	return
}
<commit_msg>fix typo<commit_after>package main

import (
	"encoding/json"
	"fmt"
	"math/rand"
	"path/filepath"
	"strconv"
	"strings"
	"time"

	tgbotapi "github.com/go-telegram-bot-api/telegram-bot-api"
	bt "github.com/ikool-cn/gobeanstalk-connection-pool"
)

var (
	telegramBot *TelegramBot
)

// TelegramBot ...
type TelegramBot struct {
	Name          string
	SelfChatID    int64
	ChannelChatID int64
	ComicPath     string
	DeleteDelay   time.Duration
	Client        *tgbotapi.BotAPI
	Queue         *bt.Pool
	Tube          string
}

// NewTelegramBot ...
func NewTelegramBot(cfg *TelegramConfig, btdAddr string) (t *TelegramBot) {
	bot, err := tgbotapi.NewBotAPI(cfg.Token)
	if err != nil {
		logger.Panicf("tg bot init failed: %+v", err)
	}
	delay, err := time.ParseDuration(cfg.DeleteDelay)
	if err != nil {
		logger.Panicf("delete delay error: %+v", err)
	}

	t = &TelegramBot{
		Name:          bot.Self.UserName,
		SelfChatID:    cfg.SelfChatID,
		ChannelChatID: cfg.ChannelChatID,
		ComicPath:     cfg.ComicPath,
		DeleteDelay:   delay,
		Client:        bot,
		Tube:          "tg",
	}
	t.Queue = &bt.Pool{
		Dial: func() (*bt.Conn, error) {
			return bt.Dial(btdAddr)
		},
		MaxIdle:     10,
		MaxActive:   100,
		IdleTimeout: 60 * time.Second,
		MaxLifetime: 180 * time.Second,
		Wait:        true,
	}
	return
}

func (t *TelegramBot) putQueue(msg []byte) {
	conn, err := t.Queue.Get()
	if err != nil {
		logger.Errorf("%+v: %s", err, string(msg))
		return
	}
	conn.Use(t.Tube)
	_, err = conn.Put(msg, 1, t.DeleteDelay, time.Minute)
	if err != nil {
		logger.Errorf("%+v", err)
		return
	}
}

func (t *TelegramBot) send(chat int64, msg string) (tgbotapi.Message, error) {
	logger.Debugf("[%d]%s", chat, msg)
	return t.Client.Send(tgbotapi.NewMessage(chat, msg))
}

func (t *TelegramBot) delMessage() {
	for {
		conn, err := t.Queue.Get()
		if err != nil {
			logger.Errorf("%+v", err)
			time.Sleep(3 * time.Second)
			continue
		}
		conn.Watch(t.Tube)
		job, err := conn.Reserve()
		if err != nil {
			logger.Warningf("%+v", err)
			time.Sleep(3 * time.Second)
			continue
		}
		msg := &tgbotapi.Message{}
		err = json.Unmarshal(job.Body, msg)
		if err != nil {
			logger.Errorf("%+v", err)
			err = conn.Bury(job.ID, 0)
			if err != nil {
				logger.Errorf("%+v", err)
			}
			time.Sleep(3 * time.Second)
			continue
		}
		delMsg := tgbotapi.DeleteMessageConfig{
			ChatID:    msg.Chat.ID,
			MessageID: msg.MessageID,
		}
		logger.Infof(":[%s]{%s}", getMsgTitle(msg), strconv.Quote(msg.Text))

		_, err = t.Client.DeleteMessage(delMsg)
		if err != nil {
			logger.Errorf("%+v", err)
			err = conn.Bury(job.ID, 0)
			if err != nil {
				logger.Errorf("%+v", err)
			}
			time.Sleep(3 * time.Second)
			continue
		}
		err = conn.Delete(job.ID)
		if err != nil {
			logger.Errorf("%+v", err)
			time.Sleep(3 * time.Second)
		}
		t.Queue.Release(conn, false)
	}
}

func (t *TelegramBot) tgBot() {
	u := tgbotapi.NewUpdate(0)
	u.Timeout = 60
	for {
		updates, err := t.Client.GetUpdatesChan(u)
		if err != nil {
			logger.Errorf("%+v", err)
			time.Sleep(3 * time.Second)
			continue
		}
		var message *tgbotapi.Message
		for update := range updates {
			if update.Message != nil {
				message = update.Message
			} else if update.EditedMessage != nil {
				message = update.EditedMessage
			} else if update.CallbackQuery != nil {
				logger.Infof(
					"recv:(%s)[%s]reaction:{%s}",
					update.CallbackQuery.ChatInstance,
					update.CallbackQuery.From.String(),
					update.CallbackQuery.Data,
				)
				data := strings.SplitN(update.CallbackQuery.Data, ":", 2)
				switch data[0] {
				case "comic", "pic":
					go onReaction(t, update.CallbackQuery)
				}
				continue
			} else {
				continue
			}
			if message.Chat.IsGroup() {
				logger.Infof(
					"recv:(%s)[%s]{%s}",
					message.Chat.Title,
					message.From.String(),
					strconv.Quote(message.Text))
			} else {
				logger.Infof(
					"recv:[%s]{%s}",
					message.From.String(),
					strconv.Quote(message.Text),
				)
			}

			if message.IsCommand() {
				switch message.Command() {
				case "start":
					go onStart(t, message)
				case "comic":
					go onComic(t, message)
				case "pic":
					go onPic(t, message)
				default:
					logger.Infof("ignore unknown cmd: %+v", message.Command())
					continue
				}
			} else {
				if message.Text == "" {
					continue
				}
				checkRepeat(t, message)
			}
		}
		logger.Warning("tg bot restarted.")
		time.Sleep(3 * time.Second)
	}
}

func checkRepeat(t *TelegramBot, message *tgbotapi.Message) {
	key := "tg_last_" + strconv.FormatInt(message.Chat.ID, 10)
	flattendMsg := strings.TrimSpace(message.Text)
	defer redisClient.LTrim(key, 0, 10)
	defer redisClient.LPush(key, flattendMsg)

	lastMsgs, err := redisClient.LRange(key, 0, 6).Result()
	if err != nil {
		logger.Errorf("%+v", err)
		return
	}
	i := 0
	for _, s := range lastMsgs {
		if s == flattendMsg {
			i++
		}
	}
	if i > 1 {
		redisClient.Del(key)
		logger.Infof("repeat: %s", strconv.Quote(message.Text))
		msg := tgbotapi.NewMessage(message.Chat.ID, message.Text)
		t.Client.Send(msg)
	}
}

func onStart(t *TelegramBot, message *tgbotapi.Message) {
	msg := tgbotapi.NewMessage(message.Chat.ID, "呀呀呀")
	msg.ReplyToMessageID = message.MessageID
	t.Client.Send(msg)
}

func onComic(t *TelegramBot, message *tgbotapi.Message) {
	files, err := filepath.Glob(filepath.Join(t.ComicPath, "*.epub"))
	if err != nil {
		logger.Errorf("%+v", err)
		return
	}
	rand.Seed(time.Now().UnixNano())
	file := files[rand.Intn(len(files))]
	number := strings.Split(strings.Split(file, "@")[1], ".")[0]
	msg := tgbotapi.NewMessage(message.Chat.ID, "🔞 https://nhentai.net/g/"+number)

	msg.ReplyMarkup = buildInlineKeyboardMarkup("comic", number)

	logger.Infof("send:[%s]{%s}", getMsgTitle(message), strconv.Quote(file))
	msgSent, err := t.Client.Send(msg)
	if err != nil {
		logger.Errorf("%+v", err)
		return
	}
	data, err := json.Marshal(msgSent)
	if err != nil {
		logger.Errorf("%+v", err)
		return
	}
	t.putQueue(data)
}

func onPic(t *TelegramBot, message *tgbotapi.Message) {
	files, err := filepath.Glob(filepath.Join(twitterBot.ImgPath, "*"))
	if err != nil {
		logger.Errorf("%+v", err)
		return
	}
	if files == nil {
		logger.Error("find no pics")
	}
	rand.Seed(time.Now().UnixNano())
	file := files[rand.Intn(len(files))]

	logger.Infof("send:[%s]{%s}", getMsgTitle(message), strconv.Quote(file))

	msg := tgbotapi.NewDocumentUpload(message.Chat.ID, file)
	msg.ReplyMarkup = buildInlineKeyboardMarkup("pic", filepath.Base(file))

	_, err = t.Client.Send(msg)
	if err != nil {
		logger.Errorf("%+v", err)
	}
}

func onReaction(t *TelegramBot, callbackQuery *tgbotapi.CallbackQuery) {
	var callbackText string

	_type, _id, reaction, err := saveReaction(callbackQuery.Data, callbackQuery.From.ID)
	if err == nil {
		diss := redisClient.SCard(buildReactionKey(_type, _id, "diss")).Val()
		if diss <= 1 {
			msg := tgbotapi.NewEditMessageReplyMarkup(
				callbackQuery.Message.Chat.ID,
				callbackQuery.Message.MessageID,
				buildInlineKeyboardMarkup(_type, _id),
			)
			_, err = t.Client.Send(msg)
		} else {
			delMsg := tgbotapi.DeleteMessageConfig{
				ChatID:    callbackQuery.Message.Chat.ID,
				MessageID: callbackQuery.Message.MessageID,
			}
			_, err = t.Client.DeleteMessage(delMsg)
			if err == nil {
				err = probate(_type, _id)
			}
		}
	}

	if err != nil {
		logger.Debugf("%+v", err)
		callbackText = err.Error()
	} else {
		callbackText = reaction + " " + _id + "!"
	}

	callbackMsg := tgbotapi.NewCallback(callbackQuery.ID, callbackText)
	_, err = t.Client.AnswerCallbackQuery(callbackMsg)
	if err != nil {
		logger.Errorf("%+v", err)
	}
}

func getMsgTitle(m *tgbotapi.Message) string {
	if m.Chat.IsGroup() {
		return m.Chat.Title
	}
	return m.From.String()
}

func buildReactionData(_type, _id, reaction string) string {
	return _type + ":" + _id + ":" + reaction
}
func buildReactionKey(_type, _id, reaction string) string {
	return "reaction_" + buildReactionData(_type, _id, reaction)
}

func buildInlineKeyboardMarkup(_type, _id string) tgbotapi.InlineKeyboardMarkup {

	likeCount, _ := redisClient.SCard(buildReactionKey(_type, _id, "like")).Result()
	dissCount, _ := redisClient.SCard(buildReactionKey(_type, _id, "diss")).Result()

	likeText := "❤"
	if likeCount > 0 {
		likeText = likeText + " " + strconv.FormatInt(likeCount, 10)
	}
	dissText := "💔"
	if dissCount > 0 {
		dissText = dissText + " " + strconv.FormatInt(dissCount, 10)
	}

	row := tgbotapi.NewInlineKeyboardRow(
		tgbotapi.NewInlineKeyboardButtonData(likeText, buildReactionData(_type, _id, "like")),
		tgbotapi.NewInlineKeyboardButtonData(dissText, buildReactionData(_type, _id, "diss")),
	)
	return tgbotapi.NewInlineKeyboardMarkup(row)
}

func saveReaction(key string, user int) (_type, _id, reaction string, err error) {
	token := strings.Split(key, ":")
	if len(token) != 3 {
		err = fmt.Errorf("react data error: %s", key)
		return
	}
	_type = token[0]
	_id = token[1]
	reaction = token[2]

	pipe := redisClient.Pipeline()
	switch reaction {
	case "like":
		likeCount := pipe.SAdd(buildReactionKey(_type, _id, "like"), strconv.Itoa(user))
		dissCount := pipe.SRem(buildReactionKey(_type, _id, "diss"), strconv.Itoa(user))
		_, err = pipe.Exec()
		if err == nil {
			if likeCount.Val()+dissCount.Val() == 0 {
				err = fmt.Errorf("not modified")
			}
		}
	case "diss":
		dissCount := pipe.SAdd(buildReactionKey(_type, _id, "diss"), strconv.Itoa(user))
		likeCount := pipe.SRem(buildReactionKey(_type, _id, "like"), strconv.Itoa(user))
		_, err = pipe.Exec()
		if err == nil {
			if likeCount.Val()+dissCount.Val() == 0 {
				err = fmt.Errorf("not modified")
			}
		}
	default:
		err = fmt.Errorf("react type error: %s", key)
	}
	return
}
<|endoftext|>


<commit_before>package main

import (
  "flag"
	"fmt"
	"text/template"
	"os"
  "io/ioutil"
  "gopkg.in/yaml.v2"
  "github.com/winkapp/libclc"
  "path"
)

var template_dir string
var root string
var config_file string

func main() {
  // Figure out which directory we are going to prep files for
  flag.StringVar(&template_dir, "templates", "", "Your templates.")
  flag.StringVar(&root, "root", "./", "Directory for configs and output.")
  flag.StringVar(&config_file, "config", "clc.yaml", "Optional config file.")
  flag.Parse()
  command := flag.Arg(0)

  switch command {
  case "cc": cc(root)
  case "vagrant": vagrant(root)
  case "units": units(root)
  case "new": newCluster(root)
  }

}

func newCluster(path string) {
  cc(root)
  vagrant(root)
  units(root)
  return
}

func vagrant(path string) {
  ud(path)
  vf(path)
  crb(path)
  return
}

func units(path string) {
  unitsYaml := getFileBytes(path + "/units.yaml")
  var units libclc.UnitConfig
  err := yaml.Unmarshal(unitsYaml, &units)
  checkError(err)
  t := unitsTemplate()
  err = libclc.WriteUnits(&units, t, path + "/units")
  checkError(err)
}

func ud(p string) {
  // TODO: Open the files document for file includes
  file1 := libclc.File{
    Content: getFile("./etcd_env.sh"),
    Path: "/home/core/etcd_env.sh",
    Owner: "core:core",
    Permissions: "744",
  }

  data := libclc.CloudConfig{
    Files: []*libclc.File{&file1},
  }

  t := udTemplate(data)
  err := libclc.WriteUserData(&data, t, path.Join(p, "user-data"))
  checkError(err)
}

func vf(p string) {
  data := libclc.Vagrantfile{}
  t := vfTemplate(data)
  err := libclc.WriteVagrantfile(&data, t, path.Join(p, "Vagrantfile"))
  checkError(err)
}

func crb(p string) {
  data := libclc.Configrb{}
  t := crbTemplate(data)
  err := libclc.WriteConfigrb(&data, t, path.Join(p, "config.rb"))
  checkError(err)
}

func cc(p string) {
  data := ccData(p)
  t := ccTemplate(data)
  err := libclc.WriteCloudConfig(&data, t, path.Join(p, "cloud-config"))
	checkError(err)
}

func ccData(path string) (data libclc.CloudConfig) {
  // Open the files document for file includes
  file1 := libclc.File{
    Content: getFile("./etcd_env.sh"),
    Path: "/home/core/etcd_env.sh",
    Owner: "core:core",
    Permissions: "744",
  }

  // Open the document for the discovery url
  discovery_url := getFile(path + "/discovery_url")

  // Include files on cloud config
  data = libclc.CloudConfig{
    DiscoveryUrl: discovery_url,
    Files: []*libclc.File{&file1},
  }
  return
}

func unitsTemplate() (t *template.Template) {
  t = getTemplate("Service", "service.template")
  return
}

func ccTemplate(data libclc.CloudConfig) (t *template.Template) {
  t = getTemplate("Cloud Config", "cloud-config.template")
  return
}

func udTemplate(data libclc.CloudConfig) (t *template.Template) {
  t = getTemplate("User Data", "user-data.template")
  return
}

func vfTemplate(data libclc.Vagrantfile) (t *template.Template) {
  t = getTemplate("Vagrantfile", "Vagrantfile.template")
  return
}

func crbTemplate(data libclc.Configrb) (t *template.Template) {
  t = getTemplate("Config.rb", "config.rb.template")
  return
}

func getFileBytes(path string) []byte {
  dat, err := ioutil.ReadFile(path)
  checkError(err)
  return dat
}

func getFile(path string) string {
  return string(getFileBytes(path))
}

func getTemplate(name string, filename string) (t *template.Template) {
  // If a template directory not set, use default template.
  if template_dir == "" {
    return nil
  }
  // If a template file does not exist, use default template.
  _, err := os.Stat(path.Join(template_dir, filename))
  if err != nil {
    return nil
  }
  templ := getFile(path.Join(template_dir, filename))

  t = template.New(name)
  t, err = t.Parse(templ)
  checkError(err)
  return
}

func checkError(err error) {
	if err != nil {
		fmt.Println("Fatal error:", err.Error())
		os.Exit(1)
	}
}
<commit_msg>put templates directive in the config file.<commit_after>package main

import (
  "flag"
	"fmt"
	"text/template"
	"os"
  "io/ioutil"
  "gopkg.in/yaml.v2"
  "github.com/winkapp/libclc"
  "path"
)

type ClcConfig struct {
  Discovery string
  Templates string
  Files []*libclc.File
}

var root string
var config_file string
var config ClcConfig

func main() {
  // Figure out which directory we are going to prep files for
  flag.StringVar(&root, "root", "./", "Directory for configs and output.")
  flag.StringVar(&config_file, "config", "clc.yaml", "Optional config file.")
  flag.Parse()
  command := flag.Arg(0)

  // Initialize the config
  configYaml := getFileBytes(config_file)
  err := yaml.Unmarshal(configYaml, &config)
  checkError(err)

  switch command {
  case "cc": cc(root)
  case "vagrant": vagrant(root)
  case "units": units(root)
  case "new": newCluster(root)
  }

}

func newCluster(path string) {
  cc(root)
  vagrant(root)
  units(root)
  return
}

func vagrant(path string) {
  ud(path)
  vf(path)
  crb(path)
  return
}

func units(path string) {
  unitsYaml := getFileBytes(path + "/units.yaml")
  var units libclc.UnitConfig
  err := yaml.Unmarshal(unitsYaml, &units)
  checkError(err)
  t := unitsTemplate()
  err = libclc.WriteUnits(&units, t, path + "/units")
  checkError(err)
}

func ud(p string) {
  // TODO: Open the files document for file includes
  file1 := libclc.File{
    Content: getFile("./etcd_env.sh"),
    Path: "/home/core/etcd_env.sh",
    Owner: "core:core",
    Permissions: "744",
  }

  data := libclc.CloudConfig{
    Files: []*libclc.File{&file1},
  }

  t := udTemplate(data)
  err := libclc.WriteUserData(&data, t, path.Join(p, "user-data"))
  checkError(err)
}

func vf(p string) {
  data := libclc.Vagrantfile{}
  t := vfTemplate(data)
  err := libclc.WriteVagrantfile(&data, t, path.Join(p, "Vagrantfile"))
  checkError(err)
}

func crb(p string) {
  data := libclc.Configrb{}
  t := crbTemplate(data)
  err := libclc.WriteConfigrb(&data, t, path.Join(p, "config.rb"))
  checkError(err)
}

func cc(p string) {
  data := ccData(p)
  t := ccTemplate(data)
  err := libclc.WriteCloudConfig(&data, t, path.Join(p, "cloud-config"))
	checkError(err)
}

func ccData(path string) (data libclc.CloudConfig) {
  // Open the files document for file includes
  file1 := libclc.File{
    Content: getFile("./etcd_env.sh"),
    Path: "/home/core/etcd_env.sh",
    Owner: "core:core",
    Permissions: "744",
  }

  // Include files on cloud config
  data = libclc.CloudConfig{
    DiscoveryUrl: config.Discovery,
    Files: []*libclc.File{&file1},
  }
  return
}

func unitsTemplate() (t *template.Template) {
  t = getTemplate("Service", "service.template")
  return
}

func ccTemplate(data libclc.CloudConfig) (t *template.Template) {
  t = getTemplate("Cloud Config", "cloud-config.template")
  return
}

func udTemplate(data libclc.CloudConfig) (t *template.Template) {
  t = getTemplate("User Data", "user-data.template")
  return
}

func vfTemplate(data libclc.Vagrantfile) (t *template.Template) {
  t = getTemplate("Vagrantfile", "Vagrantfile.template")
  return
}

func crbTemplate(data libclc.Configrb) (t *template.Template) {
  t = getTemplate("Config.rb", "config.rb.template")
  return
}

func getFileBytes(path string) []byte {
  dat, err := ioutil.ReadFile(path)
  checkError(err)
  return dat
}

func getFile(path string) string {
  return string(getFileBytes(path))
}

func getTemplate(name string, filename string) (t *template.Template) {
  // If a template directory not set, use default template.
  if config.Templates == "" {
    return nil
  }
  // If a template file does not exist, use default template.
  _, err := os.Stat(path.Join(config.Templates, filename))
  if err != nil {
    return nil
  }
  templ := getFile(path.Join(config.Templates, filename))

  t = template.New(name)
  t, err = t.Parse(templ)
  checkError(err)
  return
}

func checkError(err error) {
	if err != nil {
		fmt.Println("Fatal error:", err.Error())
		os.Exit(1)
	}
}
<|endoftext|>


<commit_before>package main

import (
	"bytes"
	"crypto/sha1"
	"encoding/json"
	"errors"
	"fmt"
	"io"
	"io/ioutil"
	"log"
	"os"
	"path/filepath"
	"reflect"
	"strings"
	"syscall"
	"text/template"
)

func exists(path string) (bool, error) {
	_, err := os.Stat(path)
	if err == nil {
		return true, nil
	}
	if os.IsNotExist(err) {
		return false, nil
	}
	return false, err
}

func groupByMulti(entries []*RuntimeContainer, key, sep string) map[string][]*RuntimeContainer {
	groups := make(map[string][]*RuntimeContainer)
	for _, v := range entries {
		value := deepGet(*v, key)
		if value != nil {
			items := strings.Split(value.(string), sep)
			for _, item := range items {
				groups[item] = append(groups[item], v)
			}

		}
	}
	return groups
}

// groupBy groups a list of *RuntimeContainers by the path property key
func groupBy(entries []*RuntimeContainer, key string) map[string][]*RuntimeContainer {
	groups := make(map[string][]*RuntimeContainer)
	for _, v := range entries {
		value := deepGet(*v, key)
		if value != nil {
			groups[value.(string)] = append(groups[value.(string)], v)
		}
	}
	return groups
}

// groupByKeys is the same as groupBy but only returns a list of keys
func groupByKeys(entries []*RuntimeContainer, key string) []string {
	groups := groupBy(entries, key)
	ret := []string{}
	for k, _ := range groups {
		ret = append(ret, k)
	}
	return ret
}

// selects entries based on key
func where(entries []*RuntimeContainer, key string, cmp string) []*RuntimeContainer {
	selection := []*RuntimeContainer{}
	for _, v := range entries {
		value := deepGet(*v, key)
		if value == cmp {
			selection = append(selection, v)
		}
	}
	return selection
}

// selects entries where a key exists
func whereExist(entries []*RuntimeContainer, key string) []*RuntimeContainer {
	selection := []*RuntimeContainer{}
	for _, v := range entries {
		value := deepGet(*v, key)
		if value != nil {
			selection = append(selection, v)
		}
	}
	return selection
}

// selects entries where a key does not exist
func whereNotExist(entries []*RuntimeContainer, key string) []*RuntimeContainer {
	selection := []*RuntimeContainer{}
	for _, v := range entries {
		value := deepGet(*v, key)
		if value == nil {
			selection = append(selection, v)
		}
	}
	return selection
}

// selects entries based on key.  Assumes key is delimited and breaks it apart before comparing
func whereAny(entries []*RuntimeContainer, key, sep string, cmp []string) []*RuntimeContainer {
	selection := []*RuntimeContainer{}
	for _, v := range entries {
		value := deepGet(*v, key)
		if value != nil {
			items := strings.Split(value.(string), sep)
			if len(intersect(cmp, items)) > 0 {
				selection = append(selection, v)
			}
		}
	}
	return selection
}

// selects entries based on key.  Assumes key is delimited and breaks it apart before comparing
func whereAll(entries []*RuntimeContainer, key, sep string, cmp []string) []*RuntimeContainer {
	selection := []*RuntimeContainer{}
	req_count := len(cmp)
	for _, v := range entries {
		value := deepGet(*v, key)
		if value != nil {
			items := strings.Split(value.(string), sep)
			if len(intersect(cmp, items)) == req_count {
				selection = append(selection, v)
			}
		}
	}
	return selection
}

// hasPrefix returns whether a given string is a prefix of another string
func hasPrefix(prefix, s string) bool {
	return strings.HasPrefix(s, prefix)
}

// hasSuffix returns whether a given string is a suffix of another string
func hasSuffix(suffix, s string) bool {
	return strings.HasSuffix(s, suffix)
}

func keys(input interface{}) (interface{}, error) {
	if input == nil {
		return nil, nil
	}

	val := reflect.ValueOf(input)
	if val.Kind() != reflect.Map {
		return nil, fmt.Errorf("Cannot call keys on a non-map value: %v", input)
	}

	vk := val.MapKeys()
	k := make([]interface{}, val.Len())
	for i, _ := range k {
		k[i] = vk[i].Interface()
	}

	return k, nil
}

func intersect(l1, l2 []string) []string {
	m := make(map[string]bool)
	m2 := make(map[string]bool)
	for _, v := range l2 {
		m2[v] = true
	}
	for _, v := range l1 {
		if m2[v] {
			m[v] = true
		}
	}
	keys := make([]string, 0, len(m))
	for k := range m {
		keys = append(keys, k)
	}
	return keys
}

func contains(item map[string]string, key string) bool {
	if _, ok := item[key]; ok {
		return true
	}
	return false
}

func dict(values ...interface{}) (map[string]interface{}, error) {
	if len(values)%2 != 0 {
		return nil, errors.New("invalid dict call")
	}
	dict := make(map[string]interface{}, len(values)/2)
	for i := 0; i < len(values); i += 2 {
		key, ok := values[i].(string)
		if !ok {
			return nil, errors.New("dict keys must be strings")
		}
		dict[key] = values[i+1]
	}
	return dict, nil
}

func hashSha1(input string) string {
	h := sha1.New()
	io.WriteString(h, input)
	return fmt.Sprintf("%x", h.Sum(nil))
}

func marshalJson(input interface{}) (string, error) {
	var buf bytes.Buffer
	enc := json.NewEncoder(&buf)
	if err := enc.Encode(input); err != nil {
		return "", err
	}
	return strings.TrimSuffix(buf.String(), "\n"), nil
}

// arrayFirst returns first item in the array or nil if the
// input is nil or empty
func arrayFirst(input interface{}) interface{} {
	if input == nil {
		return nil
	}

	arr := reflect.ValueOf(input)

	if arr.Len() == 0 {
		return nil
	}

	return arr.Index(0).Interface()
}

// arrayLast returns last item in the array
func arrayLast(input interface{}) interface{} {
	arr := reflect.ValueOf(input)
	return arr.Index(arr.Len() - 1).Interface()
}

// arrayClosest find the longest matching substring in values
// that matches input
func arrayClosest(values []string, input string) string {
	best := ""
	for _, v := range values {
		if strings.Contains(input, v) && len(v) > len(best) {
			best = v
		}
	}
	return best
}

// dirList returns a list of files in the specified path
func dirList(path string) ([]string, error) {
	names := []string{}
	files, err := ioutil.ReadDir(path)
	if err != nil {
		return names, err
	}
	for _, f := range files {
		names = append(names, f.Name())
	}
	return names, nil
}

// coalesce returns the first non nil argument
func coalesce(input ...interface{}) interface{} {
	for _, v := range input {
		if v != nil {
			return v
		}
	}
	return nil
}

// trimPrefix returns whether a given string is a prefix of another string
func trimPrefix(prefix, s string) string {
	return strings.TrimPrefix(s, prefix)
}

// trimSuffix returns whether a given string is a suffix of another string
func trimSuffix(suffix, s string) string {
	return strings.TrimSuffix(s, suffix)
}

func generateFile(config Config, containers Context) bool {
	templatePath := config.Template
	tmpl, err := template.New(filepath.Base(templatePath)).Funcs(template.FuncMap{
		"closest":       arrayClosest,
		"coalesce":      coalesce,
		"contains":      contains,
		"dict":          dict,
		"dir":           dirList,
		"exists":        exists,
		"first":         arrayFirst,
		"groupBy":       groupBy,
		"groupByKeys":   groupByKeys,
		"groupByMulti":  groupByMulti,
		"hasPrefix":     hasPrefix,
		"hasSuffix":     hasSuffix,
		"json":          marshalJson,
		"intersect":     intersect,
		"keys":          keys,
		"last":          arrayLast,
		"replace":       strings.Replace,
		"sha1":          hashSha1,
		"split":         strings.Split,
		"trimPrefix":    trimPrefix,
		"trimSuffix":    trimSuffix,
		"where":         where,
		"whereExist":    whereExist,
		"whereNotExist": whereNotExist,
		"whereAny":      whereAny,
		"whereAll":      whereAll,
	}).ParseFiles(templatePath)
	if err != nil {
		log.Fatalf("unable to parse template: %s", err)
	}

	filteredContainers := Context{}
	if config.OnlyPublished {
		for _, container := range containers {
			if len(container.PublishedAddresses()) > 0 {
				filteredContainers = append(filteredContainers, container)
			}
		}
	} else if config.OnlyExposed {
		for _, container := range containers {
			if len(container.Addresses) > 0 {
				filteredContainers = append(filteredContainers, container)
			}
		}
	} else {
		filteredContainers = containers
	}

	dest := os.Stdout
	if config.Dest != "" {
		dest, err = ioutil.TempFile(filepath.Dir(config.Dest), "docker-gen")
		defer func() {
			dest.Close()
			os.Remove(dest.Name())
		}()
		if err != nil {
			log.Fatalf("unable to create temp file: %s\n", err)
		}
	}

	var buf bytes.Buffer
	multiwriter := io.MultiWriter(dest, &buf)
	err = tmpl.ExecuteTemplate(multiwriter, filepath.Base(templatePath), &filteredContainers)
	if err != nil {
		log.Fatalf("template error: %s\n", err)
	}

	if config.Dest != "" {

		contents := []byte{}
		if fi, err := os.Stat(config.Dest); err == nil {
			if err := dest.Chmod(fi.Mode()); err != nil {
				log.Fatalf("unable to chmod temp file: %s\n", err)
			}
			if err := dest.Chown(int(fi.Sys().(*syscall.Stat_t).Uid), int(fi.Sys().(*syscall.Stat_t).Gid)); err != nil {
				log.Fatalf("unable to chown temp file: %s\n", err)
			}
			contents, err = ioutil.ReadFile(config.Dest)
			if err != nil {
				log.Fatalf("unable to compare current file contents: %s: %s\n", config.Dest, err)
			}
		}

		if bytes.Compare(contents, buf.Bytes()) != 0 {
			err = os.Rename(dest.Name(), config.Dest)
			if err != nil {
				log.Fatalf("unable to create dest file %s: %s\n", config.Dest, err)
			}
			log.Printf("Generated '%s' from %d containers", config.Dest, len(filteredContainers))
			return true
		}
		return false
	}
	return true
}
<commit_msg>Add newTemplate function<commit_after>package main

import (
	"bytes"
	"crypto/sha1"
	"encoding/json"
	"errors"
	"fmt"
	"io"
	"io/ioutil"
	"log"
	"os"
	"path/filepath"
	"reflect"
	"strings"
	"syscall"
	"text/template"
)

func exists(path string) (bool, error) {
	_, err := os.Stat(path)
	if err == nil {
		return true, nil
	}
	if os.IsNotExist(err) {
		return false, nil
	}
	return false, err
}

func groupByMulti(entries []*RuntimeContainer, key, sep string) map[string][]*RuntimeContainer {
	groups := make(map[string][]*RuntimeContainer)
	for _, v := range entries {
		value := deepGet(*v, key)
		if value != nil {
			items := strings.Split(value.(string), sep)
			for _, item := range items {
				groups[item] = append(groups[item], v)
			}

		}
	}
	return groups
}

// groupBy groups a list of *RuntimeContainers by the path property key
func groupBy(entries []*RuntimeContainer, key string) map[string][]*RuntimeContainer {
	groups := make(map[string][]*RuntimeContainer)
	for _, v := range entries {
		value := deepGet(*v, key)
		if value != nil {
			groups[value.(string)] = append(groups[value.(string)], v)
		}
	}
	return groups
}

// groupByKeys is the same as groupBy but only returns a list of keys
func groupByKeys(entries []*RuntimeContainer, key string) []string {
	groups := groupBy(entries, key)
	ret := []string{}
	for k, _ := range groups {
		ret = append(ret, k)
	}
	return ret
}

// selects entries based on key
func where(entries []*RuntimeContainer, key string, cmp string) []*RuntimeContainer {
	selection := []*RuntimeContainer{}
	for _, v := range entries {
		value := deepGet(*v, key)
		if value == cmp {
			selection = append(selection, v)
		}
	}
	return selection
}

// selects entries where a key exists
func whereExist(entries []*RuntimeContainer, key string) []*RuntimeContainer {
	selection := []*RuntimeContainer{}
	for _, v := range entries {
		value := deepGet(*v, key)
		if value != nil {
			selection = append(selection, v)
		}
	}
	return selection
}

// selects entries where a key does not exist
func whereNotExist(entries []*RuntimeContainer, key string) []*RuntimeContainer {
	selection := []*RuntimeContainer{}
	for _, v := range entries {
		value := deepGet(*v, key)
		if value == nil {
			selection = append(selection, v)
		}
	}
	return selection
}

// selects entries based on key.  Assumes key is delimited and breaks it apart before comparing
func whereAny(entries []*RuntimeContainer, key, sep string, cmp []string) []*RuntimeContainer {
	selection := []*RuntimeContainer{}
	for _, v := range entries {
		value := deepGet(*v, key)
		if value != nil {
			items := strings.Split(value.(string), sep)
			if len(intersect(cmp, items)) > 0 {
				selection = append(selection, v)
			}
		}
	}
	return selection
}

// selects entries based on key.  Assumes key is delimited and breaks it apart before comparing
func whereAll(entries []*RuntimeContainer, key, sep string, cmp []string) []*RuntimeContainer {
	selection := []*RuntimeContainer{}
	req_count := len(cmp)
	for _, v := range entries {
		value := deepGet(*v, key)
		if value != nil {
			items := strings.Split(value.(string), sep)
			if len(intersect(cmp, items)) == req_count {
				selection = append(selection, v)
			}
		}
	}
	return selection
}

// hasPrefix returns whether a given string is a prefix of another string
func hasPrefix(prefix, s string) bool {
	return strings.HasPrefix(s, prefix)
}

// hasSuffix returns whether a given string is a suffix of another string
func hasSuffix(suffix, s string) bool {
	return strings.HasSuffix(s, suffix)
}

func keys(input interface{}) (interface{}, error) {
	if input == nil {
		return nil, nil
	}

	val := reflect.ValueOf(input)
	if val.Kind() != reflect.Map {
		return nil, fmt.Errorf("Cannot call keys on a non-map value: %v", input)
	}

	vk := val.MapKeys()
	k := make([]interface{}, val.Len())
	for i, _ := range k {
		k[i] = vk[i].Interface()
	}

	return k, nil
}

func intersect(l1, l2 []string) []string {
	m := make(map[string]bool)
	m2 := make(map[string]bool)
	for _, v := range l2 {
		m2[v] = true
	}
	for _, v := range l1 {
		if m2[v] {
			m[v] = true
		}
	}
	keys := make([]string, 0, len(m))
	for k := range m {
		keys = append(keys, k)
	}
	return keys
}

func contains(item map[string]string, key string) bool {
	if _, ok := item[key]; ok {
		return true
	}
	return false
}

func dict(values ...interface{}) (map[string]interface{}, error) {
	if len(values)%2 != 0 {
		return nil, errors.New("invalid dict call")
	}
	dict := make(map[string]interface{}, len(values)/2)
	for i := 0; i < len(values); i += 2 {
		key, ok := values[i].(string)
		if !ok {
			return nil, errors.New("dict keys must be strings")
		}
		dict[key] = values[i+1]
	}
	return dict, nil
}

func hashSha1(input string) string {
	h := sha1.New()
	io.WriteString(h, input)
	return fmt.Sprintf("%x", h.Sum(nil))
}

func marshalJson(input interface{}) (string, error) {
	var buf bytes.Buffer
	enc := json.NewEncoder(&buf)
	if err := enc.Encode(input); err != nil {
		return "", err
	}
	return strings.TrimSuffix(buf.String(), "\n"), nil
}

// arrayFirst returns first item in the array or nil if the
// input is nil or empty
func arrayFirst(input interface{}) interface{} {
	if input == nil {
		return nil
	}

	arr := reflect.ValueOf(input)

	if arr.Len() == 0 {
		return nil
	}

	return arr.Index(0).Interface()
}

// arrayLast returns last item in the array
func arrayLast(input interface{}) interface{} {
	arr := reflect.ValueOf(input)
	return arr.Index(arr.Len() - 1).Interface()
}

// arrayClosest find the longest matching substring in values
// that matches input
func arrayClosest(values []string, input string) string {
	best := ""
	for _, v := range values {
		if strings.Contains(input, v) && len(v) > len(best) {
			best = v
		}
	}
	return best
}

// dirList returns a list of files in the specified path
func dirList(path string) ([]string, error) {
	names := []string{}
	files, err := ioutil.ReadDir(path)
	if err != nil {
		return names, err
	}
	for _, f := range files {
		names = append(names, f.Name())
	}
	return names, nil
}

// coalesce returns the first non nil argument
func coalesce(input ...interface{}) interface{} {
	for _, v := range input {
		if v != nil {
			return v
		}
	}
	return nil
}

// trimPrefix returns whether a given string is a prefix of another string
func trimPrefix(prefix, s string) string {
	return strings.TrimPrefix(s, prefix)
}

// trimSuffix returns whether a given string is a suffix of another string
func trimSuffix(suffix, s string) string {
	return strings.TrimSuffix(s, suffix)
}

func newTemplate(name string) *template.Template {
	tmpl := template.New(name).Funcs(template.FuncMap{
		"closest":       arrayClosest,
		"coalesce":      coalesce,
		"contains":      contains,
		"dict":          dict,
		"dir":           dirList,
		"exists":        exists,
		"first":         arrayFirst,
		"groupBy":       groupBy,
		"groupByKeys":   groupByKeys,
		"groupByMulti":  groupByMulti,
		"hasPrefix":     hasPrefix,
		"hasSuffix":     hasSuffix,
		"json":          marshalJson,
		"intersect":     intersect,
		"keys":          keys,
		"last":          arrayLast,
		"replace":       strings.Replace,
		"sha1":          hashSha1,
		"split":         strings.Split,
		"trimPrefix":    trimPrefix,
		"trimSuffix":    trimSuffix,
		"where":         where,
		"whereExist":    whereExist,
		"whereNotExist": whereNotExist,
		"whereAny":      whereAny,
		"whereAll":      whereAll,
	})
	return tmpl
}

func generateFile(config Config, containers Context) bool {
	templatePath := config.Template
	tmpl, err := newTemplate(filepath.Base(templatePath)).ParseFiles(templatePath)
	if err != nil {
		log.Fatalf("unable to parse template: %s", err)
	}

	filteredContainers := Context{}
	if config.OnlyPublished {
		for _, container := range containers {
			if len(container.PublishedAddresses()) > 0 {
				filteredContainers = append(filteredContainers, container)
			}
		}
	} else if config.OnlyExposed {
		for _, container := range containers {
			if len(container.Addresses) > 0 {
				filteredContainers = append(filteredContainers, container)
			}
		}
	} else {
		filteredContainers = containers
	}

	dest := os.Stdout
	if config.Dest != "" {
		dest, err = ioutil.TempFile(filepath.Dir(config.Dest), "docker-gen")
		defer func() {
			dest.Close()
			os.Remove(dest.Name())
		}()
		if err != nil {
			log.Fatalf("unable to create temp file: %s\n", err)
		}
	}

	var buf bytes.Buffer
	multiwriter := io.MultiWriter(dest, &buf)
	err = tmpl.ExecuteTemplate(multiwriter, filepath.Base(templatePath), &filteredContainers)
	if err != nil {
		log.Fatalf("template error: %s\n", err)
	}

	if config.Dest != "" {

		contents := []byte{}
		if fi, err := os.Stat(config.Dest); err == nil {
			if err := dest.Chmod(fi.Mode()); err != nil {
				log.Fatalf("unable to chmod temp file: %s\n", err)
			}
			if err := dest.Chown(int(fi.Sys().(*syscall.Stat_t).Uid), int(fi.Sys().(*syscall.Stat_t).Gid)); err != nil {
				log.Fatalf("unable to chown temp file: %s\n", err)
			}
			contents, err = ioutil.ReadFile(config.Dest)
			if err != nil {
				log.Fatalf("unable to compare current file contents: %s: %s\n", config.Dest, err)
			}
		}

		if bytes.Compare(contents, buf.Bytes()) != 0 {
			err = os.Rename(dest.Name(), config.Dest)
			if err != nil {
				log.Fatalf("unable to create dest file %s: %s\n", config.Dest, err)
			}
			log.Printf("Generated '%s' from %d containers", config.Dest, len(filteredContainers))
			return true
		}
		return false
	}
	return true
}
<|endoftext|>


<commit_before>package main

import (
	"bytes"
	"fmt"
	"io/ioutil"
	"path/filepath"
	"text/template"

	dep "github.com/hashicorp/consul-template/dependency"
)

type Template struct {
	// Path is the path to this template on disk.
	Path string

	// contents is string contents for this file when read from disk.
	contents string
}

// NewTemplate creates and parses a new Consul Template template at the given
// path. If the template does not exist, an error is returned. During
// initialization, the template is read and is parsed for dependencies. Any
// errors that occur are returned.
func NewTemplate(path string) (*Template, error) {
	template := &Template{Path: path}
	if err := template.init(); err != nil {
		return nil, err
	}

	return template, nil
}

// Execute evaluates this template in the context of the given brain.
//
// The first return value is the list of used dependencies.
// The second return value is the list of missing dependencies.
// The third return value is the rendered text.
// The fourth return value any error that occurs.
func (t *Template) Execute(brain *Brain) ([]dep.Dependency, []dep.Dependency, []byte, error) {
	usedMap := make(map[string]dep.Dependency)
	missingMap := make(map[string]dep.Dependency)
	name := filepath.Base(t.Path)
	funcs := funcMap(brain, usedMap, missingMap)

	tmpl, err := template.New(name).Funcs(funcs).Parse(t.contents)
	if err != nil {
		return nil, nil, nil, fmt.Errorf("template: %s", err)
	}

	// TODO: accept an io.Writer instead
	buff := new(bytes.Buffer)
	if err := tmpl.Execute(buff, nil); err != nil {
		return nil, nil, nil, fmt.Errorf("template: %s", err)
	}

	// Update this list of this template's dependencies
	var used []dep.Dependency
	for _, dep := range usedMap {
		used = append(used, dep)
	}

	// Compile the list of missing dependencies
	var missing []dep.Dependency
	for _, dep := range missingMap {
		missing = append(missing, dep)
	}

	return used, missing, buff.Bytes(), nil
}

// init reads the template file and initializes required variables.
func (t *Template) init() error {
	// Render the template
	contents, err := ioutil.ReadFile(t.Path)
	if err != nil {
		return err
	}
	t.contents = string(contents)

	return nil
}

// funcMap is the map of template functions to their respective functions.
func funcMap(brain *Brain, used, missing map[string]dep.Dependency) template.FuncMap {
	return template.FuncMap{
		// API functions
		"datacenters": datacentersFunc(brain, used, missing),
		"file":        fileFunc(brain, used, missing),
		"key":         keyFunc(brain, used, missing),
		"ls":          lsFunc(brain, used, missing),
		"node":        nodeFunc(brain, used, missing),
		"nodes":       nodesFunc(brain, used, missing),
		"service":     serviceFunc(brain, used, missing),
		"services":    servicesFunc(brain, used, missing),
		"tree":        treeFunc(brain, used, missing),
		"vault":       vaultFunc(brain, used, missing),

		// Helper functions
		"byKey":           byKey,
		"byTag":           byTag,
		"env":             env,
		"explode":         explode,
		"loop":            loop,
		"join":            join,
		"parseBool":       parseBool,
		"parseFloat":      parseFloat,
		"parseInt":        parseInt,
		"parseJSON":       parseJSON,
		"parseUint":       parseUint,
		"plugin":          plugin,
		"regexReplaceAll": regexReplaceAll,
		"regexMatch":      regexMatch,
		"replaceAll":      replaceAll,
		"timestamp":       timestamp,
		"toLower":         toLower,
		"toJSON":          toJSON,
		"toJSONPretty":    toJSONPretty,
		"toTitle":         toTitle,
		"toUpper":         toUpper,
		"toYAML":          toYAML,
		"split":           split,
	}
}
<commit_msg>Add 'senv' : set env variable<commit_after>package main

import (
	"bytes"
	"fmt"
	"io/ioutil"
	"path/filepath"
	"text/template"

	dep "github.com/hashicorp/consul-template/dependency"
)

type Template struct {
	// Path is the path to this template on disk.
	Path string

	// contents is string contents for this file when read from disk.
	contents string
}

// NewTemplate creates and parses a new Consul Template template at the given
// path. If the template does not exist, an error is returned. During
// initialization, the template is read and is parsed for dependencies. Any
// errors that occur are returned.
func NewTemplate(path string) (*Template, error) {
	template := &Template{Path: path}
	if err := template.init(); err != nil {
		return nil, err
	}

	return template, nil
}

// Execute evaluates this template in the context of the given brain.
//
// The first return value is the list of used dependencies.
// The second return value is the list of missing dependencies.
// The third return value is the rendered text.
// The fourth return value any error that occurs.
func (t *Template) Execute(brain *Brain) ([]dep.Dependency, []dep.Dependency, []byte, error) {
	usedMap := make(map[string]dep.Dependency)
	missingMap := make(map[string]dep.Dependency)
	name := filepath.Base(t.Path)
	funcs := funcMap(brain, usedMap, missingMap)

	tmpl, err := template.New(name).Funcs(funcs).Parse(t.contents)
	if err != nil {
		return nil, nil, nil, fmt.Errorf("template: %s", err)
	}

	// TODO: accept an io.Writer instead
	buff := new(bytes.Buffer)
	if err := tmpl.Execute(buff, nil); err != nil {
		return nil, nil, nil, fmt.Errorf("template: %s", err)
	}

	// Update this list of this template's dependencies
	var used []dep.Dependency
	for _, dep := range usedMap {
		used = append(used, dep)
	}

	// Compile the list of missing dependencies
	var missing []dep.Dependency
	for _, dep := range missingMap {
		missing = append(missing, dep)
	}

	return used, missing, buff.Bytes(), nil
}

// init reads the template file and initializes required variables.
func (t *Template) init() error {
	// Render the template
	contents, err := ioutil.ReadFile(t.Path)
	if err != nil {
		return err
	}
	t.contents = string(contents)

	return nil
}

// funcMap is the map of template functions to their respective functions.
func funcMap(brain *Brain, used, missing map[string]dep.Dependency) template.FuncMap {
	return template.FuncMap{
		// API functions
		"datacenters": datacentersFunc(brain, used, missing),
		"file":        fileFunc(brain, used, missing),
		"key":         keyFunc(brain, used, missing),
		"ls":          lsFunc(brain, used, missing),
		"node":        nodeFunc(brain, used, missing),
		"nodes":       nodesFunc(brain, used, missing),
		"service":     serviceFunc(brain, used, missing),
		"services":    servicesFunc(brain, used, missing),
		"tree":        treeFunc(brain, used, missing),
		"vault":       vaultFunc(brain, used, missing),

		// Helper functions
		"byKey":           byKey,
		"byTag":           byTag,
		"env":             env,
		"explode":         explode,
		"loop":            loop,
		"join":            join,
		"parseBool":       parseBool,
		"parseFloat":      parseFloat,
		"parseInt":        parseInt,
		"parseJSON":       parseJSON,
		"parseUint":       parseUint,
		"plugin":          plugin,
		"regexReplaceAll": regexReplaceAll,
		"regexMatch":      regexMatch,
		"replaceAll":      replaceAll,
		"timestamp":       timestamp,
		"toLower":         toLower,
		"toJSON":          toJSON,
		"toJSONPretty":    toJSONPretty,
		"toTitle":         toTitle,
		"toUpper":         toUpper,
		"toYAML":          toYAML,
		"split":           split,
		"senv":	   	   senv,
	}
}
<|endoftext|>


<commit_before>// Copyright 2009 The Go Authors.  All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package net

import (
	"io"
	"os"
	"syscall"
	"time"
)

// TCPConn is an implementation of the Conn interface for TCP network
// connections.
type TCPConn struct {
	conn
}

func newTCPConn(fd *netFD) *TCPConn {
	return &TCPConn{conn{fd}}
}

// ReadFrom implements the io.ReaderFrom ReadFrom method.
func (c *TCPConn) ReadFrom(r io.Reader) (int64, error) {
	return genericReadFrom(c, r)
}

// CloseRead shuts down the reading side of the TCP connection.
// Most callers should just use Close.
func (c *TCPConn) CloseRead() error {
	if !c.ok() {
		return syscall.EINVAL
	}
	return c.fd.CloseRead()
}

// CloseWrite shuts down the writing side of the TCP connection.
// Most callers should just use Close.
func (c *TCPConn) CloseWrite() error {
	if !c.ok() {
		return syscall.EINVAL
	}
	return c.fd.CloseWrite()
}

// SetLinger sets the behavior of Close() on a connection which still
// has data waiting to be sent or to be acknowledged.
//
// If sec < 0 (the default), Close returns immediately and the
// operating system finishes sending the data in the background.
//
// If sec == 0, Close returns immediately and the operating system
// discards any unsent or unacknowledged data.
//
// If sec > 0, Close blocks for at most sec seconds waiting for data
// to be sent and acknowledged.
func (c *TCPConn) SetLinger(sec int) error {
	return syscall.EPLAN9
}

// SetKeepAlive sets whether the operating system should send
// keepalive messages on the connection.
func (c *TCPConn) SetKeepAlive(keepalive bool) error {
	return syscall.EPLAN9
}

// SetNoDelay controls whether the operating system should delay
// packet transmission in hopes of sending fewer packets (Nagle's
// algorithm).  The default is true (no delay), meaning that data is
// sent as soon as possible after a Write.
func (c *TCPConn) SetNoDelay(noDelay bool) error {
	return syscall.EPLAN9
}

// DialTCP connects to the remote address raddr on the network net,
// which must be "tcp", "tcp4", or "tcp6".  If laddr is not nil, it is
// used as the local address for the connection.
func DialTCP(net string, laddr, raddr *TCPAddr) (*TCPConn, error) {
	return dialTCP(net, laddr, raddr, noDeadline)
}

func dialTCP(net string, laddr, raddr *TCPAddr, deadline time.Time) (*TCPConn, error) {
	if !deadline.IsZero() {
		panic("net.dialTCP: deadline not implemented on Plan 9")
	}
	switch net {
	case "tcp", "tcp4", "tcp6":
	default:
		return nil, &OpError{"dial", net, raddr, UnknownNetworkError(net)}
	}
	if raddr == nil {
		return nil, &OpError{"dial", net, nil, errMissingAddress}
	}
	fd, err := dialPlan9(net, laddr, raddr)
	if err != nil {
		return nil, err
	}
	return newTCPConn(fd), nil
}

// TCPListener is a TCP network listener.  Clients should typically
// use variables of type Listener instead of assuming TCP.
type TCPListener struct {
	fd *netFD
}

// AcceptTCP accepts the next incoming call and returns the new
// connection and the remote address.
func (l *TCPListener) AcceptTCP() (*TCPConn, error) {
	if l == nil || l.fd == nil || l.fd.ctl == nil {
		return nil, syscall.EINVAL
	}
	fd, err := l.fd.acceptPlan9()
	if err != nil {
		return nil, err
	}
	return newTCPConn(fd), nil
}

// Accept implements the Accept method in the Listener interface; it
// waits for the next call and returns a generic Conn.
func (l *TCPListener) Accept() (Conn, error) {
	if l == nil || l.fd == nil || l.fd.ctl == nil {
		return nil, syscall.EINVAL
	}
	c, err := l.AcceptTCP()
	if err != nil {
		return nil, err
	}
	return c, nil
}

// Close stops listening on the TCP address.
// Already Accepted connections are not closed.
func (l *TCPListener) Close() error {
	if l == nil || l.fd == nil || l.fd.ctl == nil {
		return syscall.EINVAL
	}
	if _, err := l.fd.ctl.WriteString("hangup"); err != nil {
		l.fd.ctl.Close()
		return &OpError{"close", l.fd.ctl.Name(), l.fd.laddr, err}
	}
	return l.fd.ctl.Close()
}

// Addr returns the listener's network address, a *TCPAddr.
func (l *TCPListener) Addr() Addr { return l.fd.laddr }

// SetDeadline sets the deadline associated with the listener.
// A zero time value disables the deadline.
func (l *TCPListener) SetDeadline(t time.Time) error {
	if l == nil || l.fd == nil || l.fd.ctl == nil {
		return syscall.EINVAL
	}
	return setDeadline(l.fd, t)
}

// File returns a copy of the underlying os.File, set to blocking
// mode.  It is the caller's responsibility to close f when finished.
// Closing l does not affect f, and closing f does not affect l.
//
// The returned os.File's file descriptor is different from the
// connection's.  Attempting to change properties of the original
// using this duplicate may or may not have the desired effect.
func (l *TCPListener) File() (f *os.File, err error) { return l.dup() }

// ListenTCP announces on the TCP address laddr and returns a TCP
// listener.  Net must be "tcp", "tcp4", or "tcp6".  If laddr has a
// port of 0, ListenTCP will choose an available port.  The caller can
// use the Addr method of TCPListener to retrieve the chosen address.
func ListenTCP(net string, laddr *TCPAddr) (*TCPListener, error) {
	switch net {
	case "tcp", "tcp4", "tcp6":
	default:
		return nil, &OpError{"listen", net, laddr, UnknownNetworkError(net)}
	}
	if laddr == nil {
		laddr = &TCPAddr{}
	}
	fd, err := listenPlan9(net, laddr)
	if err != nil {
		return nil, err
	}
	return &TCPListener{fd}, nil
}
<commit_msg>net: add missing SetKeepAlivePeriod for Plan 9<commit_after>// Copyright 2009 The Go Authors.  All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package net

import (
	"io"
	"os"
	"syscall"
	"time"
)

// TCPConn is an implementation of the Conn interface for TCP network
// connections.
type TCPConn struct {
	conn
}

func newTCPConn(fd *netFD) *TCPConn {
	return &TCPConn{conn{fd}}
}

// ReadFrom implements the io.ReaderFrom ReadFrom method.
func (c *TCPConn) ReadFrom(r io.Reader) (int64, error) {
	return genericReadFrom(c, r)
}

// CloseRead shuts down the reading side of the TCP connection.
// Most callers should just use Close.
func (c *TCPConn) CloseRead() error {
	if !c.ok() {
		return syscall.EINVAL
	}
	return c.fd.CloseRead()
}

// CloseWrite shuts down the writing side of the TCP connection.
// Most callers should just use Close.
func (c *TCPConn) CloseWrite() error {
	if !c.ok() {
		return syscall.EINVAL
	}
	return c.fd.CloseWrite()
}

// SetLinger sets the behavior of Close() on a connection which still
// has data waiting to be sent or to be acknowledged.
//
// If sec < 0 (the default), Close returns immediately and the
// operating system finishes sending the data in the background.
//
// If sec == 0, Close returns immediately and the operating system
// discards any unsent or unacknowledged data.
//
// If sec > 0, Close blocks for at most sec seconds waiting for data
// to be sent and acknowledged.
func (c *TCPConn) SetLinger(sec int) error {
	return syscall.EPLAN9
}

// SetKeepAlive sets whether the operating system should send
// keepalive messages on the connection.
func (c *TCPConn) SetKeepAlive(keepalive bool) error {
	return syscall.EPLAN9
}

// SetKeepAlivePeriod sets period between keep alives.
func (c *TCPConn) SetKeepAlivePeriod(d time.Duration) error {
	return syscall.EPLAN9
}

// SetNoDelay controls whether the operating system should delay
// packet transmission in hopes of sending fewer packets (Nagle's
// algorithm).  The default is true (no delay), meaning that data is
// sent as soon as possible after a Write.
func (c *TCPConn) SetNoDelay(noDelay bool) error {
	return syscall.EPLAN9
}

// DialTCP connects to the remote address raddr on the network net,
// which must be "tcp", "tcp4", or "tcp6".  If laddr is not nil, it is
// used as the local address for the connection.
func DialTCP(net string, laddr, raddr *TCPAddr) (*TCPConn, error) {
	return dialTCP(net, laddr, raddr, noDeadline)
}

func dialTCP(net string, laddr, raddr *TCPAddr, deadline time.Time) (*TCPConn, error) {
	if !deadline.IsZero() {
		panic("net.dialTCP: deadline not implemented on Plan 9")
	}
	switch net {
	case "tcp", "tcp4", "tcp6":
	default:
		return nil, &OpError{"dial", net, raddr, UnknownNetworkError(net)}
	}
	if raddr == nil {
		return nil, &OpError{"dial", net, nil, errMissingAddress}
	}
	fd, err := dialPlan9(net, laddr, raddr)
	if err != nil {
		return nil, err
	}
	return newTCPConn(fd), nil
}

// TCPListener is a TCP network listener.  Clients should typically
// use variables of type Listener instead of assuming TCP.
type TCPListener struct {
	fd *netFD
}

// AcceptTCP accepts the next incoming call and returns the new
// connection and the remote address.
func (l *TCPListener) AcceptTCP() (*TCPConn, error) {
	if l == nil || l.fd == nil || l.fd.ctl == nil {
		return nil, syscall.EINVAL
	}
	fd, err := l.fd.acceptPlan9()
	if err != nil {
		return nil, err
	}
	return newTCPConn(fd), nil
}

// Accept implements the Accept method in the Listener interface; it
// waits for the next call and returns a generic Conn.
func (l *TCPListener) Accept() (Conn, error) {
	if l == nil || l.fd == nil || l.fd.ctl == nil {
		return nil, syscall.EINVAL
	}
	c, err := l.AcceptTCP()
	if err != nil {
		return nil, err
	}
	return c, nil
}

// Close stops listening on the TCP address.
// Already Accepted connections are not closed.
func (l *TCPListener) Close() error {
	if l == nil || l.fd == nil || l.fd.ctl == nil {
		return syscall.EINVAL
	}
	if _, err := l.fd.ctl.WriteString("hangup"); err != nil {
		l.fd.ctl.Close()
		return &OpError{"close", l.fd.ctl.Name(), l.fd.laddr, err}
	}
	return l.fd.ctl.Close()
}

// Addr returns the listener's network address, a *TCPAddr.
func (l *TCPListener) Addr() Addr { return l.fd.laddr }

// SetDeadline sets the deadline associated with the listener.
// A zero time value disables the deadline.
func (l *TCPListener) SetDeadline(t time.Time) error {
	if l == nil || l.fd == nil || l.fd.ctl == nil {
		return syscall.EINVAL
	}
	return setDeadline(l.fd, t)
}

// File returns a copy of the underlying os.File, set to blocking
// mode.  It is the caller's responsibility to close f when finished.
// Closing l does not affect f, and closing f does not affect l.
//
// The returned os.File's file descriptor is different from the
// connection's.  Attempting to change properties of the original
// using this duplicate may or may not have the desired effect.
func (l *TCPListener) File() (f *os.File, err error) { return l.dup() }

// ListenTCP announces on the TCP address laddr and returns a TCP
// listener.  Net must be "tcp", "tcp4", or "tcp6".  If laddr has a
// port of 0, ListenTCP will choose an available port.  The caller can
// use the Addr method of TCPListener to retrieve the chosen address.
func ListenTCP(net string, laddr *TCPAddr) (*TCPListener, error) {
	switch net {
	case "tcp", "tcp4", "tcp6":
	default:
		return nil, &OpError{"listen", net, laddr, UnknownNetworkError(net)}
	}
	if laddr == nil {
		laddr = &TCPAddr{}
	}
	fd, err := listenPlan9(net, laddr)
	if err != nil {
		return nil, err
	}
	return &TCPListener{fd}, nil
}
<|endoftext|>