<commit_before>package main

import (
	"bufio"
	"bytes"
	"encoding/binary"
	"errors"
	"io"
	"log"
	"net"
	"os"
	"path"
	"strings"
	"time"

	_ "github.com/lib/pq"
)

var startupPacketSizeInvalid = errors.New("Terminating connection that provided an abnormally sized startup message packet")
var unsupportedProtocolVersion = errors.New("Unexpected protocol version number; expected 196608")
var incorrectlyFormattedPacket = errors.New("Incorrectly formatted protocol packet")

type startupMessage map[string]string

type backendKeyData struct {
	processId int32
	secretKey int32
}

func sendError(conn net.Conn, errorMessage string) {
	errorMessageExcludingSize := &bytes.Buffer{}
	errorMessageExcludingSize.Grow(1024)

	errorMessageExcludingSize.Write([]byte("S"))
	errorMessageExcludingSize.Write([]byte("ERROR"))
	errorMessageExcludingSize.Write([]byte{0})

	errorMessageExcludingSize.Write([]byte("C"))
	errorMessageExcludingSize.Write([]byte("08000")) // connection exception
	errorMessageExcludingSize.Write([]byte{0})

	errorMessageExcludingSize.Write([]byte("M"))
	errorMessageExcludingSize.Write([]byte(errorMessage))
	errorMessageExcludingSize.Write([]byte{0})

	// Terminating ErrorResponse byte
	errorMessageExcludingSize.Write([]byte{0})

	// Send the error message on the connection.  No error handling here; the connection
	// isn't likely to live for long now anyways. :-)
	conn.Write([]byte{'E'})
	binary.Write(conn, binary.BigEndian, int32(errorMessageExcludingSize.Len()+4))
	conn.Write(errorMessageExcludingSize.Bytes())
}

func readStartupMessage(conn net.Conn) (*startupMessage, error) {
	return readStartupMessageInternal(conn, true)
}

func readStartupMessageInternal(conn net.Conn, allowRecursion bool) (*startupMessage, error) {
	var startupMessageSize int32
	err := binary.Read(conn, binary.BigEndian, &startupMessageSize)
	if err != nil {
		return nil, err
	}

	if startupMessageSize < 0 || startupMessageSize > 8096 {
		sendError(conn, "Startup packet size invalid")
		return nil, startupPacketSizeInvalid
	}

	log.Printf("startup packet was %v bytes", startupMessageSize)

	startupMessageData := make([]byte, startupMessageSize-4)
	_, err = io.ReadFull(conn, startupMessageData)
	if err != nil {
		sendError(conn, "Socket read error")
		return nil, err
	}

	log.Printf("startup packet read")

	var protocolVersionNumber int32
	buf := bytes.NewBuffer(startupMessageData)
	err = binary.Read(buf, binary.BigEndian, &protocolVersionNumber)
	if err != nil {
		sendError(conn, "Socket read error")
		return nil, err
	}

	if protocolVersionNumber == 80877103 && allowRecursion {
		log.Printf("SSLRequest received; returning N")
		conn.Write([]byte{'N'})
		return readStartupMessageInternal(conn, false)
	} else if protocolVersionNumber == 80877102 {
		log.Printf("CancelRequest packet received; not supported yet!")
		return nil, unsupportedProtocolVersion
	} else if protocolVersionNumber != 196608 {
		sendError(conn, "Unsupported protocol version")
		return nil, unsupportedProtocolVersion
	}

	startupMessageData = startupMessageData[4:]
	startupParameters := make(startupMessage)
	for {
		nextZero := bytes.IndexByte(startupMessageData, 0)
		if nextZero == -1 {
			sendError(conn, "Malformed startup packet")
			return nil, incorrectlyFormattedPacket
		} else if nextZero == 0 {
			break
		}

		key := string(startupMessageData[:nextZero])
		startupMessageData = startupMessageData[nextZero+1:]

		nextZero = bytes.IndexByte(startupMessageData, 0)
		if nextZero == -1 {
			sendError(conn, "Malformed startup packet")
			return nil, incorrectlyFormattedPacket
		}
		value := string(startupMessageData[:nextZero])
		startupMessageData = startupMessageData[nextZero+1:]

		log.Printf("key = %v, value = %v", key, value)
		startupParameters[key] = value
	}

	return &startupParameters, nil
}

func handleIncomingConnection(conn net.Conn, masterRequestChannel, replicaRequestChannel chan<- serverRequest) {
	defer conn.Close()

	// One-minute timeout to read the startup message
	conn.SetReadDeadline(time.Now().Add(time.Minute))

	startupMessage, err := readStartupMessage(conn)
	if err != nil {
		log.Print(err)
		return
	}
	startupParameters := *startupMessage

	// Reset read deadline to no timeout
	conn.SetReadDeadline(time.Time{})

	// Check if we're going to connect to a replica or to the master
	dbName, ok := startupParameters["database"]
	wantReplica := false
	if !ok {
		dbName, ok = startupParameters["user"]
		if !ok {
			sendError(conn, "Missing database or user parameter")
			log.Printf("Expected database or user parameter, neither found")
			return
		}
	}
	if strings.HasSuffix(dbName, "_replica") {
		wantReplica = true
		startupParameters["database"] = dbName[:len(dbName)-8]
		log.Printf("Rewriting database name from %v to %v", dbName, startupParameters["database"])
	}

	// Fetch a backend server, either a master or a replica
	responseChannel := make(chan *string)
	request := serverRequest{responseChannel}
	if wantReplica {
		replicaRequestChannel <- request
	} else {
		masterRequestChannel <- request
	}
	backend := <-responseChannel
	if backend == nil {
		sendError(conn, "Unable to find satisfactory backend server")
		log.Println("Unable to find satisfactory backend server")
		return
	}

	// Create the new startup message w/ the possibly different startupParameters
	var protocolVersion int32 = 196608
	newStartupMessageExcludingSize := &bytes.Buffer{}
	newStartupMessageExcludingSize.Grow(1024)
	binary.Write(newStartupMessageExcludingSize, binary.BigEndian, protocolVersion)
	for key, value := range startupParameters {
		newStartupMessageExcludingSize.Write([]byte(key))
		newStartupMessageExcludingSize.Write([]byte{0})
		newStartupMessageExcludingSize.Write([]byte(value))
		newStartupMessageExcludingSize.Write([]byte{0})
	}
	// Terminating startup packet byte
	newStartupMessageExcludingSize.Write([]byte{0})

	// Send the new connection our startup packet
	log.Printf("backend to connect to: %v", *backend)
	upstream, err := net.Dial(network(*backend))
	if err != nil {
		sendError(conn, "Unable to connect to backend server")
		log.Print(err)
		return
	}
	err = binary.Write(upstream, binary.BigEndian, int32(newStartupMessageExcludingSize.Len()+4))
	if err != nil {
		sendError(conn, "Backend network error")
		log.Print(err)
		return
	}
	_, err = upstream.Write(newStartupMessageExcludingSize.Bytes())
	if err != nil {
		sendError(conn, "Backend network error")
		log.Print(err)
		return
	}

	// Begin copying all input from the client to the upstream connection.
	go func() {
		numCopied, err := io.Copy(conn, upstream)
		log.Printf("Copy(conn, upstream) -> %v, %v", numCopied, err)
	}()

	// Proxy upstream -> conn, but attempting to extract the BackendKeyData packet
	_, err = proxyPacketsUntilBackendKeyDataReceived(conn, upstream)
	if err != nil {
		log.Printf("XXX %v", err)
		sendError(conn, err.Error())
		log.Print(err)
		return
	}

	// Stream data between the two network connections
	// Also begin copying all input from the upstream connection to the client.
	numCopied, err := io.Copy(upstream, conn)
	log.Printf("Copy(upstream, conn) -> %v, %v", numCopied, err)
	if err != nil {
		log.Print(err)
		return
	}

	log.Printf("Connection closed softly")
}

// Proxy backend -> client, but attempting to extract the BackendKeyData packet
func proxyPacketsUntilBackendKeyDataReceived(client, backend net.Conn) (*backendKeyData, error) {

	typeBuffer := make([]byte, 1)
	bufferedClient := bufio.NewWriter(client)
	var messageSize int32

	for {
		_, err := io.ReadFull(backend, typeBuffer)
		if err != nil {
			return nil, err
		}
		_, err = bufferedClient.Write(typeBuffer)
		if err != nil {
			return nil, err
		}

		err = binary.Read(backend, binary.BigEndian, &messageSize)
		if err != nil {
			return nil, err
		} else if messageSize < 0 || messageSize > 8096 {
			return nil, startupPacketSizeInvalid
		}
		err = binary.Write(bufferedClient, binary.BigEndian, &messageSize)
		if err != nil {
			return nil, err
		}

		// BackendKeyData message
		if typeBuffer[0] == 'K' {
			retval := backendKeyData{}
			err = binary.Read(backend, binary.BigEndian, &retval.processId)
			if err != nil {
				return nil, err
			}
			err = binary.Read(backend, binary.BigEndian, &retval.secretKey)
			if err != nil {
				return nil, err
			}

			err = binary.Write(bufferedClient, binary.BigEndian, &retval.processId)
			if err != nil {
				return nil, err
			}
			err = binary.Write(bufferedClient, binary.BigEndian, &retval.secretKey)
			if err != nil {
				return nil, err
			}

			log.Printf("backendKeyData %v %v", retval.processId, retval.secretKey)

			err = bufferedClient.Flush()
			if err != nil {
				return nil, err
			}

			return &retval, nil
		} else {
			messageBuffer := make([]byte, messageSize-4) // size includes the Int32 containing the size
			_, err = io.ReadFull(backend, messageBuffer)
			if err != nil {
				return nil, err
			}
			_, err = bufferedClient.Write(messageBuffer)
			if err != nil {
				return nil, err
			}

			err = bufferedClient.Flush()
			if err != nil {
				return nil, err
			}
		}
	}
}

//// BEGIN: Copy/hacked from lib/pq
func network(name string) (string, string) {
	o := make(Values)

	// A number of defaults are applied here, in this order:
	//
	// * Very low precedence defaults applied in every situation
	// * Environment variables
	// * Explicitly passed connection information
	o.Set("host", "localhost")
	o.Set("port", "5432")

	for k, v := range parseEnviron(os.Environ()) {
		o.Set(k, v)
	}

	parseOpts(name, o)

	// Don't care about user, just host & port
	//// If a user is not provided by any other means, the last
	//// resort is to use the current operating system provided user
	//// name.
	//if o.Get("user") == "" {
	//  u, err := userCurrent()
	//  if err != nil {
	//      return nil, err
	//  } else {
	//      o.Set("user", u)
	//  }
	//}

	host := o.Get("host")

	if strings.HasPrefix(host, "/") {
		sockPath := path.Join(host, ".s.PGSQL."+o.Get("port"))
		return "unix", sockPath
	}

	return "tcp", host + ":" + o.Get("port")
}

type Values map[string]string

func (vs Values) Set(k, v string) {
	vs[k] = v
}

func (vs Values) Get(k string) (v string) {
	return vs[k]
}

func parseOpts(name string, o Values) {
	if len(name) == 0 {
		return
	}

	name = strings.TrimSpace(name)

	ps := strings.Split(name, " ")
	for _, p := range ps {
		kv := strings.Split(p, "=")
		if len(kv) < 2 {
			log.Fatalf("invalid option: %q", p)
		}
		o.Set(kv[0], kv[1])
	}
}

// parseEnviron tries to mimic some of libpq's environment handling
//
// To ease testing, it does not directly reference os.Environ, but is
// designed to accept its output.
//
// Environment-set connection information is intended to have a higher
// precedence than a library default but lower than any explicitly
// passed information (such as in the URL or connection string).
func parseEnviron(env []string) (out map[string]string) {
	out = make(map[string]string)

	for _, v := range env {
		parts := strings.SplitN(v, "=", 2)

		accrue := func(keyname string) {
			out[keyname] = parts[1]
		}

		// The order of these is the same as is seen in the
		// PostgreSQL 9.1 manual, with omissions briefly
		// noted.
		switch parts[0] {
		case "PGHOST":
			accrue("host")
		case "PGHOSTADDR":
			accrue("hostaddr")
		case "PGPORT":
			accrue("port")
		case "PGDATABASE":
			accrue("dbname")
		case "PGUSER":
			accrue("user")
		case "PGPASSWORD":
			accrue("password")
		// skip PGPASSFILE, PGSERVICE, PGSERVICEFILE,
		// PGREALM
		case "PGOPTIONS":
			accrue("options")
		case "PGAPPNAME":
			accrue("application_name")
		case "PGSSLMODE":
			accrue("sslmode")
		case "PGREQUIRESSL":
			accrue("requiressl")
		case "PGSSLCERT":
			accrue("sslcert")
		case "PGSSLKEY":
			accrue("sslkey")
		case "PGSSLROOTCERT":
			accrue("sslrootcert")
		case "PGSSLCRL":
			accrue("sslcrl")
		case "PGREQUIREPEER":
			accrue("requirepeer")
		case "PGKRBSRVNAME":
			accrue("krbsrvname")
		case "PGGSSLIB":
			accrue("gsslib")
		case "PGCONNECT_TIMEOUT":
			accrue("connect_timeout")
		case "PGCLIENTENCODING":
			accrue("client_encoding")
			// skip PGDATESTYLE, PGTZ, PGGEQO, PGSYSCONFDIR,
			// PGLOCALEDIR
		}
	}

	return out
}

//// END: Copy/hacked from lib/pq
<commit_msg>Remove extra debug stmt<commit_after>package main

import (
	"bufio"
	"bytes"
	"encoding/binary"
	"errors"
	"io"
	"log"
	"net"
	"os"
	"path"
	"strings"
	"time"

	_ "github.com/lib/pq"
)

var startupPacketSizeInvalid = errors.New("Terminating connection that provided an abnormally sized startup message packet")
var unsupportedProtocolVersion = errors.New("Unexpected protocol version number; expected 196608")
var incorrectlyFormattedPacket = errors.New("Incorrectly formatted protocol packet")

type startupMessage map[string]string

type backendKeyData struct {
	processId int32
	secretKey int32
}

func sendError(conn net.Conn, errorMessage string) {
	errorMessageExcludingSize := &bytes.Buffer{}
	errorMessageExcludingSize.Grow(1024)

	errorMessageExcludingSize.Write([]byte("S"))
	errorMessageExcludingSize.Write([]byte("ERROR"))
	errorMessageExcludingSize.Write([]byte{0})

	errorMessageExcludingSize.Write([]byte("C"))
	errorMessageExcludingSize.Write([]byte("08000")) // connection exception
	errorMessageExcludingSize.Write([]byte{0})

	errorMessageExcludingSize.Write([]byte("M"))
	errorMessageExcludingSize.Write([]byte(errorMessage))
	errorMessageExcludingSize.Write([]byte{0})

	// Terminating ErrorResponse byte
	errorMessageExcludingSize.Write([]byte{0})

	// Send the error message on the connection.  No error handling here; the connection
	// isn't likely to live for long now anyways. :-)
	conn.Write([]byte{'E'})
	binary.Write(conn, binary.BigEndian, int32(errorMessageExcludingSize.Len()+4))
	conn.Write(errorMessageExcludingSize.Bytes())
}

func readStartupMessage(conn net.Conn) (*startupMessage, error) {
	return readStartupMessageInternal(conn, true)
}

func readStartupMessageInternal(conn net.Conn, allowRecursion bool) (*startupMessage, error) {
	var startupMessageSize int32
	err := binary.Read(conn, binary.BigEndian, &startupMessageSize)
	if err != nil {
		return nil, err
	}

	if startupMessageSize < 0 || startupMessageSize > 8096 {
		sendError(conn, "Startup packet size invalid")
		return nil, startupPacketSizeInvalid
	}

	log.Printf("startup packet was %v bytes", startupMessageSize)

	startupMessageData := make([]byte, startupMessageSize-4)
	_, err = io.ReadFull(conn, startupMessageData)
	if err != nil {
		sendError(conn, "Socket read error")
		return nil, err
	}

	log.Printf("startup packet read")

	var protocolVersionNumber int32
	buf := bytes.NewBuffer(startupMessageData)
	err = binary.Read(buf, binary.BigEndian, &protocolVersionNumber)
	if err != nil {
		sendError(conn, "Socket read error")
		return nil, err
	}

	if protocolVersionNumber == 80877103 && allowRecursion {
		log.Printf("SSLRequest received; returning N")
		conn.Write([]byte{'N'})
		return readStartupMessageInternal(conn, false)
	} else if protocolVersionNumber == 80877102 {
		log.Printf("CancelRequest packet received; not supported yet!")
		return nil, unsupportedProtocolVersion
	} else if protocolVersionNumber != 196608 {
		sendError(conn, "Unsupported protocol version")
		return nil, unsupportedProtocolVersion
	}

	startupMessageData = startupMessageData[4:]
	startupParameters := make(startupMessage)
	for {
		nextZero := bytes.IndexByte(startupMessageData, 0)
		if nextZero == -1 {
			sendError(conn, "Malformed startup packet")
			return nil, incorrectlyFormattedPacket
		} else if nextZero == 0 {
			break
		}

		key := string(startupMessageData[:nextZero])
		startupMessageData = startupMessageData[nextZero+1:]

		nextZero = bytes.IndexByte(startupMessageData, 0)
		if nextZero == -1 {
			sendError(conn, "Malformed startup packet")
			return nil, incorrectlyFormattedPacket
		}
		value := string(startupMessageData[:nextZero])
		startupMessageData = startupMessageData[nextZero+1:]

		log.Printf("key = %v, value = %v", key, value)
		startupParameters[key] = value
	}

	return &startupParameters, nil
}

func handleIncomingConnection(conn net.Conn, masterRequestChannel, replicaRequestChannel chan<- serverRequest) {
	defer conn.Close()

	// One-minute timeout to read the startup message
	conn.SetReadDeadline(time.Now().Add(time.Minute))

	startupMessage, err := readStartupMessage(conn)
	if err != nil {
		log.Print(err)
		return
	}
	startupParameters := *startupMessage

	// Reset read deadline to no timeout
	conn.SetReadDeadline(time.Time{})

	// Check if we're going to connect to a replica or to the master
	dbName, ok := startupParameters["database"]
	wantReplica := false
	if !ok {
		dbName, ok = startupParameters["user"]
		if !ok {
			sendError(conn, "Missing database or user parameter")
			log.Printf("Expected database or user parameter, neither found")
			return
		}
	}
	if strings.HasSuffix(dbName, "_replica") {
		wantReplica = true
		startupParameters["database"] = dbName[:len(dbName)-8]
		log.Printf("Rewriting database name from %v to %v", dbName, startupParameters["database"])
	}

	// Fetch a backend server, either a master or a replica
	responseChannel := make(chan *string)
	request := serverRequest{responseChannel}
	if wantReplica {
		replicaRequestChannel <- request
	} else {
		masterRequestChannel <- request
	}
	backend := <-responseChannel
	if backend == nil {
		sendError(conn, "Unable to find satisfactory backend server")
		log.Println("Unable to find satisfactory backend server")
		return
	}

	// Create the new startup message w/ the possibly different startupParameters
	var protocolVersion int32 = 196608
	newStartupMessageExcludingSize := &bytes.Buffer{}
	newStartupMessageExcludingSize.Grow(1024)
	binary.Write(newStartupMessageExcludingSize, binary.BigEndian, protocolVersion)
	for key, value := range startupParameters {
		newStartupMessageExcludingSize.Write([]byte(key))
		newStartupMessageExcludingSize.Write([]byte{0})
		newStartupMessageExcludingSize.Write([]byte(value))
		newStartupMessageExcludingSize.Write([]byte{0})
	}
	// Terminating startup packet byte
	newStartupMessageExcludingSize.Write([]byte{0})

	// Send the new connection our startup packet
	log.Printf("backend to connect to: %v", *backend)
	upstream, err := net.Dial(network(*backend))
	if err != nil {
		sendError(conn, "Unable to connect to backend server")
		log.Print(err)
		return
	}
	err = binary.Write(upstream, binary.BigEndian, int32(newStartupMessageExcludingSize.Len()+4))
	if err != nil {
		sendError(conn, "Backend network error")
		log.Print(err)
		return
	}
	_, err = upstream.Write(newStartupMessageExcludingSize.Bytes())
	if err != nil {
		sendError(conn, "Backend network error")
		log.Print(err)
		return
	}

	// Begin copying all input from the client to the upstream connection.
	go func() {
		numCopied, err := io.Copy(conn, upstream)
		log.Printf("Copy(conn, upstream) -> %v, %v", numCopied, err)
	}()

	// Proxy upstream -> conn, but attempting to extract the BackendKeyData packet
	_, err = proxyPacketsUntilBackendKeyDataReceived(conn, upstream)
	if err != nil {
		sendError(conn, err.Error())
		log.Print(err)
		return
	}

	// Stream data between the two network connections
	// Also begin copying all input from the upstream connection to the client.
	numCopied, err := io.Copy(upstream, conn)
	log.Printf("Copy(upstream, conn) -> %v, %v", numCopied, err)
	if err != nil {
		log.Print(err)
		return
	}

	log.Printf("Connection closed softly")
}

// Proxy backend -> client, but attempting to extract the BackendKeyData packet
func proxyPacketsUntilBackendKeyDataReceived(client, backend net.Conn) (*backendKeyData, error) {

	typeBuffer := make([]byte, 1)
	bufferedClient := bufio.NewWriter(client)
	var messageSize int32

	for {
		_, err := io.ReadFull(backend, typeBuffer)
		if err != nil {
			return nil, err
		}
		_, err = bufferedClient.Write(typeBuffer)
		if err != nil {
			return nil, err
		}

		err = binary.Read(backend, binary.BigEndian, &messageSize)
		if err != nil {
			return nil, err
		} else if messageSize < 0 || messageSize > 8096 {
			return nil, startupPacketSizeInvalid
		}
		err = binary.Write(bufferedClient, binary.BigEndian, &messageSize)
		if err != nil {
			return nil, err
		}

		// BackendKeyData message
		if typeBuffer[0] == 'K' {
			retval := backendKeyData{}
			err = binary.Read(backend, binary.BigEndian, &retval.processId)
			if err != nil {
				return nil, err
			}
			err = binary.Read(backend, binary.BigEndian, &retval.secretKey)
			if err != nil {
				return nil, err
			}

			err = binary.Write(bufferedClient, binary.BigEndian, &retval.processId)
			if err != nil {
				return nil, err
			}
			err = binary.Write(bufferedClient, binary.BigEndian, &retval.secretKey)
			if err != nil {
				return nil, err
			}

			log.Printf("backendKeyData %v %v", retval.processId, retval.secretKey)

			err = bufferedClient.Flush()
			if err != nil {
				return nil, err
			}

			return &retval, nil
		} else {
			messageBuffer := make([]byte, messageSize-4) // size includes the Int32 containing the size
			_, err = io.ReadFull(backend, messageBuffer)
			if err != nil {
				return nil, err
			}
			_, err = bufferedClient.Write(messageBuffer)
			if err != nil {
				return nil, err
			}

			err = bufferedClient.Flush()
			if err != nil {
				return nil, err
			}
		}
	}
}

//// BEGIN: Copy/hacked from lib/pq
func network(name string) (string, string) {
	o := make(Values)

	// A number of defaults are applied here, in this order:
	//
	// * Very low precedence defaults applied in every situation
	// * Environment variables
	// * Explicitly passed connection information
	o.Set("host", "localhost")
	o.Set("port", "5432")

	for k, v := range parseEnviron(os.Environ()) {
		o.Set(k, v)
	}

	parseOpts(name, o)

	// Don't care about user, just host & port
	//// If a user is not provided by any other means, the last
	//// resort is to use the current operating system provided user
	//// name.
	//if o.Get("user") == "" {
	//  u, err := userCurrent()
	//  if err != nil {
	//      return nil, err
	//  } else {
	//      o.Set("user", u)
	//  }
	//}

	host := o.Get("host")

	if strings.HasPrefix(host, "/") {
		sockPath := path.Join(host, ".s.PGSQL."+o.Get("port"))
		return "unix", sockPath
	}

	return "tcp", host + ":" + o.Get("port")
}

type Values map[string]string

func (vs Values) Set(k, v string) {
	vs[k] = v
}

func (vs Values) Get(k string) (v string) {
	return vs[k]
}

func parseOpts(name string, o Values) {
	if len(name) == 0 {
		return
	}

	name = strings.TrimSpace(name)

	ps := strings.Split(name, " ")
	for _, p := range ps {
		kv := strings.Split(p, "=")
		if len(kv) < 2 {
			log.Fatalf("invalid option: %q", p)
		}
		o.Set(kv[0], kv[1])
	}
}

// parseEnviron tries to mimic some of libpq's environment handling
//
// To ease testing, it does not directly reference os.Environ, but is
// designed to accept its output.
//
// Environment-set connection information is intended to have a higher
// precedence than a library default but lower than any explicitly
// passed information (such as in the URL or connection string).
func parseEnviron(env []string) (out map[string]string) {
	out = make(map[string]string)

	for _, v := range env {
		parts := strings.SplitN(v, "=", 2)

		accrue := func(keyname string) {
			out[keyname] = parts[1]
		}

		// The order of these is the same as is seen in the
		// PostgreSQL 9.1 manual, with omissions briefly
		// noted.
		switch parts[0] {
		case "PGHOST":
			accrue("host")
		case "PGHOSTADDR":
			accrue("hostaddr")
		case "PGPORT":
			accrue("port")
		case "PGDATABASE":
			accrue("dbname")
		case "PGUSER":
			accrue("user")
		case "PGPASSWORD":
			accrue("password")
		// skip PGPASSFILE, PGSERVICE, PGSERVICEFILE,
		// PGREALM
		case "PGOPTIONS":
			accrue("options")
		case "PGAPPNAME":
			accrue("application_name")
		case "PGSSLMODE":
			accrue("sslmode")
		case "PGREQUIRESSL":
			accrue("requiressl")
		case "PGSSLCERT":
			accrue("sslcert")
		case "PGSSLKEY":
			accrue("sslkey")
		case "PGSSLROOTCERT":
			accrue("sslrootcert")
		case "PGSSLCRL":
			accrue("sslcrl")
		case "PGREQUIREPEER":
			accrue("requirepeer")
		case "PGKRBSRVNAME":
			accrue("krbsrvname")
		case "PGGSSLIB":
			accrue("gsslib")
		case "PGCONNECT_TIMEOUT":
			accrue("connect_timeout")
		case "PGCLIENTENCODING":
			accrue("client_encoding")
			// skip PGDATESTYLE, PGTZ, PGGEQO, PGSYSCONFDIR,
			// PGLOCALEDIR
		}
	}

	return out
}

//// END: Copy/hacked from lib/pq
<|endoftext|>


<commit_before>package stripe

import "encoding/json"

// CouponDuration is the list of allowed values for the coupon's duration.
type CouponDuration string

const (
	CouponDurationForever   CouponDuration = "forever"
	CouponDurationOnce      CouponDuration = "once"
	CouponDurationRepeating CouponDuration = "repeating"
)

// CouponParams is the set of parameters that can be used when creating a coupon.
// For more details see https://stripe.com/docs/api#create_coupon.
type CouponParams struct {
	Params           `form:"*"`
	AmountOff        *int64  `form:"amount_off"`
	Currency         *string `form:"currency"`
	Duration         *string `form:"duration"`
	DurationInMonths *int64  `form:"duration_in_months"`
	ID               *string `form:"id"`
	MaxRedemptions   *int64  `form:"max_redemptions"`
	PercentOff       *int64  `form:"percent_off"`
	RedeemBy         *int64  `form:"redeem_by"`
}

// CouponListParams is the set of parameters that can be used when listing coupons.
// For more detail see https://stripe.com/docs/api#list_coupons.
type CouponListParams struct {
	ListParams   `form:"*"`
	Created      *int64            `form:"created"`
	CreatedRange *RangeQueryParams `form:"created"`
}

// Coupon is the resource representing a Stripe coupon.
// For more details see https://stripe.com/docs/api#coupons.
type Coupon struct {
	AmountOff        int64             `json:"amount_off"`
	Created          int64             `json:"created"`
	Currency         Currency          `json:"currency"`
	Deleted          bool              `json:"deleted"`
	Duration         CouponDuration    `json:"duration"`
	DurationInMonths int64             `json:"duration_in_months"`
	ID               string            `json:"id"`
	Livemode         bool              `json:"livemode"`
	MaxRedemptions   int64             `json:"max_redemptions"`
	Metadata         map[string]string `json:"metadata"`
	PercentOff       int64             `json:"percent_off"`
	RedeemBy         int64             `json:"redeem_by"`
	TimesRedeemed    int64             `json:"times_redeemed"`
	Valid            bool              `json:"valid"`
}

// CouponList is a list of coupons as retrieved from a list endpoint.
type CouponList struct {
	ListMeta
	Data []*Coupon `json:"data"`
}

// UnmarshalJSON handles deserialization of a Coupon.
// This custom unmarshaling is needed because the resulting
// property may be an id or the full struct if it was expanded.
func (c *Coupon) UnmarshalJSON(data []byte) error {
	if id, ok := ParseID(data); ok {
		c.ID = id
		return nil
	}

	type coupon Coupon
	var v coupon
	if err := json.Unmarshal(data, &v); err != nil {
		return err
	}

	*c = Coupon(v)
	return nil
}
<commit_msg>Add name to the Coupon resource<commit_after>package stripe

import "encoding/json"

// CouponDuration is the list of allowed values for the coupon's duration.
type CouponDuration string

const (
	CouponDurationForever   CouponDuration = "forever"
	CouponDurationOnce      CouponDuration = "once"
	CouponDurationRepeating CouponDuration = "repeating"
)

// CouponParams is the set of parameters that can be used when creating a coupon.
// For more details see https://stripe.com/docs/api#create_coupon.
type CouponParams struct {
	Params           `form:"*"`
	AmountOff        *int64  `form:"amount_off"`
	Currency         *string `form:"currency"`
	Duration         *string `form:"duration"`
	DurationInMonths *int64  `form:"duration_in_months"`
	ID               *string `form:"id"`
	MaxRedemptions   *int64  `form:"max_redemptions"`
	Name             *string `form:"name"`
	PercentOff       *int64  `form:"percent_off"`
	RedeemBy         *int64  `form:"redeem_by"`
}

// CouponListParams is the set of parameters that can be used when listing coupons.
// For more detail see https://stripe.com/docs/api#list_coupons.
type CouponListParams struct {
	ListParams   `form:"*"`
	Created      *int64            `form:"created"`
	CreatedRange *RangeQueryParams `form:"created"`
}

// Coupon is the resource representing a Stripe coupon.
// For more details see https://stripe.com/docs/api#coupons.
type Coupon struct {
	AmountOff        int64             `json:"amount_off"`
	Created          int64             `json:"created"`
	Currency         Currency          `json:"currency"`
	Deleted          bool              `json:"deleted"`
	Duration         CouponDuration    `json:"duration"`
	DurationInMonths int64             `json:"duration_in_months"`
	ID               string            `json:"id"`
	Livemode         bool              `json:"livemode"`
	MaxRedemptions   int64             `json:"max_redemptions"`
	Metadata         map[string]string `json:"metadata"`
	Name             string            `json:"name"`
	PercentOff       int64             `json:"percent_off"`
	RedeemBy         int64             `json:"redeem_by"`
	TimesRedeemed    int64             `json:"times_redeemed"`
	Valid            bool              `json:"valid"`
}

// CouponList is a list of coupons as retrieved from a list endpoint.
type CouponList struct {
	ListMeta
	Data []*Coupon `json:"data"`
}

// UnmarshalJSON handles deserialization of a Coupon.
// This custom unmarshaling is needed because the resulting
// property may be an id or the full struct if it was expanded.
func (c *Coupon) UnmarshalJSON(data []byte) error {
	if id, ok := ParseID(data); ok {
		c.ID = id
		return nil
	}

	type coupon Coupon
	var v coupon
	if err := json.Unmarshal(data, &v); err != nil {
		return err
	}

	*c = Coupon(v)
	return nil
}
<|endoftext|>


<commit_before>package wstest_test

import (
	"context"
	"errors"
	"fmt"
	"net/http"
	"testing"
	"time"

	"github.com/gorilla/websocket"
	"github.com/posener/wstest"
	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/require"
)

// TestClient demonstrate the usage of wstest package
func TestClient(t *testing.T) {
	t.Parallel()
	var (
		s    = &handler{Upgraded: make(chan struct{})}
		d    = wstest.NewDialer(s)
		done = make(chan struct{})
	)

	c, resp, err := d.Dial("ws://example.org/ws", nil)
	require.Nil(t, err)

	<-s.Upgraded

	assert.Equal(t, http.StatusSwitchingProtocols, resp.StatusCode)

	for i := 0; i < 3; i++ {
		msg := fmt.Sprintf("hello, world! %d", i)

		go func() {
			err := c.WriteMessage(websocket.TextMessage, []byte(msg))
			require.Nil(t, err)
			done <- struct{}{}
		}()

		mT, m, err := s.ReadMessage()
		require.Nil(t, err)

		assert.Equal(t, msg, string(m))
		assert.Equal(t, websocket.TextMessage, mT)
		<-done

		go func() {
			err := s.WriteMessage(websocket.TextMessage, []byte(msg))
			require.Nil(t, err)
			done <- struct{}{}
		}()

		mT, m, err = c.ReadMessage()
		require.Nil(t, err)

		assert.Equal(t, msg, string(m))
		assert.Equal(t, websocket.TextMessage, mT)
		<-done
	}

	err = c.Close()
	require.Nil(t, err)

	err = s.Close()
	require.Nil(t, err)
}

// TestConcurrent tests concurrent reads and writes from a connection
func TestConcurrent(t *testing.T) {
	t.Parallel()
	var (
		s     = &handler{Upgraded: make(chan struct{})}
		d     = wstest.NewDialer(s)
		count = 20
	)

	c, _, err := d.Dial("ws://example.org/ws", nil)
	require.Nil(t, err)

	<-s.Upgraded

	for _, pair := range []struct{ src, dst *websocket.Conn }{{s.Conn, c}, {c, s.Conn}} {
		go func() {
			for i := 0; i < count; i++ {
				err := pair.src.WriteJSON(i)
				require.Nil(t, err)
			}
		}()

		received := make([]bool, count)

		for i := 0; i < count; i++ {
			var j int
			err := pair.dst.ReadJSON(&j)
			require.Nil(t, err)

			received[j] = true
		}

		var missing []int

		for i := range received {
			if !received[i] {
				missing = append(missing, i)
			}
		}
		assert.Equal(t, 0, len(missing), "%q -> %q: Did not received: %q", pair.src.LocalAddr(), pair.dst.LocalAddr(), missing)
	}

	err = c.Close()
	require.Nil(t, err)

	err = s.Close()
	require.Nil(t, err)
}

func TestBadAddress(t *testing.T) {
	t.Parallel()

	tests := []struct {
		url  string
		code int
	}{

		{
			url:  "ws://example.org/not-ws",
			code: http.StatusNotFound,
		},
		{
			url: "http://example.org/ws",
		},
	}

	for _, tt := range tests {
		t.Run(tt.url, func(t *testing.T) {
			s := &handler{Upgraded: make(chan struct{})}
			d := wstest.NewDialer(s)
			c, resp, err := d.Dial(tt.url, nil)
			assert.Nil(t, c)
			assert.NotNil(t, err)
			if tt.code != 0 {
				assert.Equal(t, tt.code, resp.StatusCode)
			}

			err = s.Close()
			require.Nil(t, err)
		})
	}
}

// TestConnectDeadline tests connection deadlines
func TestDeadlines(t *testing.T) {
	t.Parallel()
	h := &handler{Upgraded: make(chan struct{})}
	d := wstest.NewDialer(h)

	c, _, err := d.Dial("ws://example.org/ws", nil)
	require.Nil(t, err)

	<-h.Upgraded

	var i int

	for _, pair := range []struct{ src, dst *websocket.Conn }{{h.Conn, c}, {c, h.Conn}} {

		// set the deadline to now, and test for timeout
		pair.dst.SetReadDeadline(time.Now())
		err = pair.dst.ReadJSON(&i)
		assert.True(t, errors.As(err, &context.DeadlineExceeded))

		err = pair.dst.ReadJSON(&i)
		assert.True(t, errors.As(err, &context.DeadlineExceeded))

		go pair.src.WriteJSON(1)
		err = pair.dst.ReadJSON(&i)
		assert.True(t, errors.As(err, &context.DeadlineExceeded))

		// even after updating the deadline, should get an error
		pair.dst.SetReadDeadline(time.Now().Add(time.Second))
		err = pair.dst.ReadJSON(&i)
		assert.True(t, errors.As(err, &context.DeadlineExceeded))
	}
}

// TestConnectDeadline tests connection deadline
func TestConnectDeadline(t *testing.T) {
	t.Parallel()

	tests := []struct {
		path    string
		timeout time.Duration
		wantErr bool
	}{
		{
			path:    "/ws/delay",
			timeout: time.Millisecond,
			wantErr: true,
		},
		{
			path:    "/ws",
			timeout: time.Second,
		},
	}

	for _, tt := range tests {
		t.Run(fmt.Sprintf("%s/%s", tt.path, tt.timeout), func(t *testing.T) {
			s := &handler{Upgraded: make(chan struct{})}
			d := wstest.NewDialer(s)
			d.HandshakeTimeout = tt.timeout
			_, _, err := d.Dial("ws://example.org"+tt.path, nil)
			if tt.wantErr {
				assert.NotNil(t, err)
				return
			}

			assert.Nil(t, err)
			select {
			case <-s.Upgraded:
			case <-time.After(time.Second):
				t.Fatal("connection was not upgraded after 1s")
			}
		})
	}
}

// dialer for test purposes, can't handle multiple websocket connections concurrently
type handler struct {
	*websocket.Conn
	upgrader websocket.Upgrader
	Upgraded chan struct{}
}

func (s *handler) ServeHTTP(w http.ResponseWriter, r *http.Request) {

	switch r.URL.Path {
	case "/ws":
		s.connect(w, r)

	case "/ws/delay":
		<-time.After(500 * time.Millisecond)
		s.connect(w, r)

	default:
		w.WriteHeader(http.StatusNotFound)
	}

}

func (s *handler) connect(w http.ResponseWriter, r *http.Request) {
	defer close(s.Upgraded)
	var err error
	s.Conn, err = s.upgrader.Upgrade(w, r, nil)
	if err != nil {
		return
	}
}

func (s *handler) Close() error {
	if s.Conn == nil {
		return nil
	}
	return s.Conn.Close()
}
<commit_msg>Fix race condition in tests<commit_after>package wstest_test

import (
	"context"
	"errors"
	"fmt"
	"net/http"
	"testing"
	"time"

	"github.com/gorilla/websocket"
	"github.com/posener/wstest"
	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/require"
)

// TestClient demonstrate the usage of wstest package
func TestClient(t *testing.T) {
	t.Parallel()
	var (
		s    = &handler{Upgraded: make(chan struct{})}
		d    = wstest.NewDialer(s)
		done = make(chan struct{})
	)

	c, resp, err := d.Dial("ws://example.org/ws", nil)
	require.Nil(t, err)

	<-s.Upgraded

	assert.Equal(t, http.StatusSwitchingProtocols, resp.StatusCode)

	for i := 0; i < 3; i++ {
		msg := fmt.Sprintf("hello, world! %d", i)

		go func() {
			err := c.WriteMessage(websocket.TextMessage, []byte(msg))
			require.Nil(t, err)
			done <- struct{}{}
		}()

		mT, m, err := s.ReadMessage()
		require.Nil(t, err)

		assert.Equal(t, msg, string(m))
		assert.Equal(t, websocket.TextMessage, mT)
		<-done

		go func() {
			err := s.WriteMessage(websocket.TextMessage, []byte(msg))
			require.Nil(t, err)
			done <- struct{}{}
		}()

		mT, m, err = c.ReadMessage()
		require.Nil(t, err)

		assert.Equal(t, msg, string(m))
		assert.Equal(t, websocket.TextMessage, mT)
		<-done
	}

	err = c.Close()
	require.Nil(t, err)

	err = s.Close()
	require.Nil(t, err)
}

// TestConcurrent tests concurrent reads and writes from a connection
func TestConcurrent(t *testing.T) {
	t.Parallel()
	var (
		s     = &handler{Upgraded: make(chan struct{})}
		d     = wstest.NewDialer(s)
		count = 20
	)

	c, _, err := d.Dial("ws://example.org/ws", nil)
	require.Nil(t, err)

	<-s.Upgraded

	for _, pair := range []struct{ src, dst *websocket.Conn }{{s.Conn, c}, {c, s.Conn}} {
		go func() {
			for i := 0; i < count; i++ {
				err := pair.src.WriteJSON(i)
				require.Nil(t, err)
			}
		}()

		received := make([]bool, count)

		for i := 0; i < count; i++ {
			var j int
			err := pair.dst.ReadJSON(&j)
			require.Nil(t, err)

			received[j] = true
		}

		var missing []int

		for i := range received {
			if !received[i] {
				missing = append(missing, i)
			}
		}
		assert.Equal(t, 0, len(missing), "%q -> %q: Did not received: %q", pair.src.LocalAddr(), pair.dst.LocalAddr(), missing)
	}

	err = c.Close()
	require.Nil(t, err)

	err = s.Close()
	require.Nil(t, err)
}

func TestBadAddress(t *testing.T) {
	t.Parallel()

	tests := []struct {
		url  string
		code int
	}{

		{
			url:  "ws://example.org/not-ws",
			code: http.StatusNotFound,
		},
		{
			url: "http://example.org/ws",
		},
	}

	for _, tt := range tests {
		t.Run(tt.url, func(t *testing.T) {
			s := &handler{Upgraded: make(chan struct{})}
			d := wstest.NewDialer(s)
			c, resp, err := d.Dial(tt.url, nil)
			assert.Nil(t, c)
			assert.NotNil(t, err)
			if tt.code != 0 {
				assert.Equal(t, tt.code, resp.StatusCode)
			}

			err = s.Close()
			require.Nil(t, err)
		})
	}
}

// TestConnectDeadline tests connection deadlines
func TestDeadlines(t *testing.T) {
	t.Parallel()
	h := &handler{Upgraded: make(chan struct{})}
	d := wstest.NewDialer(h)

	c, _, err := d.Dial("ws://example.org/ws", nil)
	require.Nil(t, err)

	<-h.Upgraded

	assertDeadlineExceeded := func(t *testing.T, err error) {
		v := context.DeadlineExceeded
		assert.True(t, errors.As(err, &v))
	}

	var i int
	for _, pair := range []struct{ src, dst *websocket.Conn }{{h.Conn, c}, {c, h.Conn}} {

		// set the deadline to now, and test for timeout
		pair.dst.SetReadDeadline(time.Now())
		err = pair.dst.ReadJSON(&i)
		assertDeadlineExceeded(t, err)

		err = pair.dst.ReadJSON(&i)
		assertDeadlineExceeded(t, err)

		go pair.src.WriteJSON(1)
		err = pair.dst.ReadJSON(&i)
		assertDeadlineExceeded(t, err)

		// even after updating the deadline, should get an error
		pair.dst.SetReadDeadline(time.Now().Add(time.Second))
		err = pair.dst.ReadJSON(&i)
		assertDeadlineExceeded(t, err)
	}
}

// TestConnectDeadline tests connection deadline
func TestConnectDeadline(t *testing.T) {
	t.Parallel()

	tests := []struct {
		path    string
		timeout time.Duration
		wantErr bool
	}{
		{
			path:    "/ws/delay",
			timeout: time.Millisecond,
			wantErr: true,
		},
		{
			path:    "/ws",
			timeout: time.Second,
		},
	}

	for _, tt := range tests {
		t.Run(fmt.Sprintf("%s/%s", tt.path, tt.timeout), func(t *testing.T) {
			s := &handler{Upgraded: make(chan struct{})}
			d := wstest.NewDialer(s)
			d.HandshakeTimeout = tt.timeout
			_, _, err := d.Dial("ws://example.org"+tt.path, nil)
			if tt.wantErr {
				assert.NotNil(t, err)
				return
			}

			assert.Nil(t, err)
			select {
			case <-s.Upgraded:
			case <-time.After(time.Second):
				t.Fatal("connection was not upgraded after 1s")
			}
		})
	}
}

// dialer for test purposes, can't handle multiple websocket connections concurrently
type handler struct {
	*websocket.Conn
	upgrader websocket.Upgrader
	Upgraded chan struct{}
}

func (s *handler) ServeHTTP(w http.ResponseWriter, r *http.Request) {

	switch r.URL.Path {
	case "/ws":
		s.connect(w, r)

	case "/ws/delay":
		<-time.After(500 * time.Millisecond)
		s.connect(w, r)

	default:
		w.WriteHeader(http.StatusNotFound)
	}

}

func (s *handler) connect(w http.ResponseWriter, r *http.Request) {
	defer close(s.Upgraded)
	var err error
	s.Conn, err = s.upgrader.Upgrade(w, r, nil)
	if err != nil {
		return
	}
}

func (s *handler) Close() error {
	if s.Conn == nil {
		return nil
	}
	return s.Conn.Close()
}
<|endoftext|>


<commit_before>package bolt

import (
	"bytes"
	"fmt"
	"sort"
)

// Cursor represents an iterator that can traverse over all key/value pairs in a bucket in sorted order.
// Cursors see nested buckets with value == nil.
// Cursors can be obtained from a transaction and are valid as long as the transaction is open.
//
// Keys and values returned from the cursor are only valid for the life of the transaction.
//
// Changing data while traversing with a cursor may cause it to be invalidated
// and return unexpected keys and/or values. You must reposition your cursor
// after mutating data.
type Cursor struct {
	bucket *Bucket
	stack  []elemRef
}

// Bucket returns the bucket that this cursor was created from.
func (c *Cursor) Bucket() *Bucket {
	return c.bucket
}

// First moves the cursor to the first item in the bucket and returns its key and value.
// If the bucket is empty then a nil key and value are returned.
// The returned key and value are only valid for the life of the transaction.
func (c *Cursor) First() (key []byte, value []byte) {
	_assert(c.bucket.tx.db != nil, "tx closed")
	c.stack = c.stack[:0]
	p, n := c.bucket.pageNode(c.bucket.root)
	c.stack = append(c.stack, elemRef{page: p, node: n, index: 0})
	c.first()

	// If we land on an empty page then move to the next value.
	// https://github.com/boltdb/bolt/issues/450
	if c.stack[len(c.stack)-1].count() == 0 {
		c.next()
	}

	k, v, flags := c.keyValue()
	if (flags & uint32(bucketLeafFlag)) != 0 {
		return k, nil
	}
	return k, v

}

// Last moves the cursor to the last item in the bucket and returns its key and value.
// If the bucket is empty then a nil key and value are returned.
// The returned key and value are only valid for the life of the transaction.
func (c *Cursor) Last() (key []byte, value []byte) {
	_assert(c.bucket.tx.db != nil, "tx closed")
	c.stack = c.stack[:0]
	p, n := c.bucket.pageNode(c.bucket.root)
	ref := elemRef{page: p, node: n}
	ref.index = ref.count() - 1
	c.stack = append(c.stack, ref)
	c.last()
	k, v, flags := c.keyValue()
	if (flags & uint32(bucketLeafFlag)) != 0 {
		return k, nil
	}
	return k, v
}

// Next moves the cursor to the next item in the bucket and returns its key and value.
// If the cursor is at the end of the bucket then a nil key and value are returned.
// The returned key and value are only valid for the life of the transaction.
func (c *Cursor) Next() (key []byte, value []byte) {
	_assert(c.bucket.tx.db != nil, "tx closed")
	k, v, flags := c.next()
	if (flags & uint32(bucketLeafFlag)) != 0 {
		return k, nil
	}
	return k, v
}

// Prev moves the cursor to the previous item in the bucket and returns its key and value.
// If the cursor is at the beginning of the bucket then a nil key and value are returned.
// The returned key and value are only valid for the life of the transaction.
func (c *Cursor) Prev() (key []byte, value []byte) {
	_assert(c.bucket.tx.db != nil, "tx closed")

	// Attempt to move back one element until we're successful.
	// Move up the stack as we hit the beginning of each page in our stack.
	for i := len(c.stack) - 1; i >= 0; i-- {
		elem := &c.stack[i]
		if elem.index > 0 {
			elem.index--
			break
		}
		c.stack = c.stack[:i]
	}

	// If we've hit the end then return nil.
	if len(c.stack) == 0 {
		return nil, nil
	}

	// Move down the stack to find the last element of the last leaf under this branch.
	c.last()
	k, v, flags := c.keyValue()
	if (flags & uint32(bucketLeafFlag)) != 0 {
		return k, nil
	}
	return k, v
}

// Seek moves the cursor to a given key and returns it.
// If the key does not exist then the next key is used. If no keys
// follow, a nil key is returned.
// The returned key and value are only valid for the life of the transaction.
func (c *Cursor) Seek(seek []byte) (key []byte, value []byte) {
	k, v, flags := c.seek(seek)

	// If we ended up after the last element of a page then move to the next one.
	if ref := &c.stack[len(c.stack)-1]; ref.index >= ref.count() {
		k, v, flags = c.next()
	}

	if k == nil {
		return nil, nil
	} else if (flags & uint32(bucketLeafFlag)) != 0 {
		return k, nil
	}
	return k, v
}

// Delete removes the current key/value under the cursor from the bucket.
// Delete fails if current key/value is a bucket or if the transaction is not writable.
func (c *Cursor) Delete() error {
	if c.bucket.tx.db == nil {
		return ErrTxClosed
	} else if !c.bucket.Writable() {
		return ErrTxNotWritable
	}

	key, _, flags := c.keyValue()
	// Return an error if current value is a bucket.
	if (flags & bucketLeafFlag) != 0 {
		return ErrIncompatibleValue
	}
	c.node().del(key)

	return nil
}

// seek moves the cursor to a given key and returns it.
// If the key does not exist then the next key is used.
func (c *Cursor) seek(seek []byte) (key []byte, value []byte, flags uint32) {
	_assert(c.bucket.tx.db != nil, "tx closed")

	// Start from root page/node and traverse to correct page.
	c.stack = c.stack[:0]
	c.search(seek, c.bucket.root)
	ref := &c.stack[len(c.stack)-1]

	// If the cursor is pointing to the end of page/node then return nil.
	if ref.index >= ref.count() {
		return nil, nil, 0
	}

	// If this is a bucket then return a nil value.
	return c.keyValue()
}

// first moves the cursor to the first leaf element under the last page in the stack.
func (c *Cursor) first() {
	for {
		// Exit when we hit a leaf page.
		var ref = &c.stack[len(c.stack)-1]
		if ref.isLeaf() {
			break
		}

		// Keep adding pages pointing to the first element to the stack.
		var pgid pgid
		if ref.node != nil {
			pgid = ref.node.inodes[ref.index].pgid
		} else {
			pgid = ref.page.branchPageElement(uint16(ref.index)).pgid
		}
		p, n := c.bucket.pageNode(pgid)
		c.stack = append(c.stack, elemRef{page: p, node: n, index: 0})
	}
}

// last moves the cursor to the last leaf element under the last page in the stack.
func (c *Cursor) last() {
	for {
		// Exit when we hit a leaf page.
		ref := &c.stack[len(c.stack)-1]
		if ref.isLeaf() {
			break
		}

		// Keep adding pages pointing to the last element in the stack.
		var pgid pgid
		if ref.node != nil {
			pgid = ref.node.inodes[ref.index].pgid
		} else {
			pgid = ref.page.branchPageElement(uint16(ref.index)).pgid
		}
		p, n := c.bucket.pageNode(pgid)

		var nextRef = elemRef{page: p, node: n}
		nextRef.index = nextRef.count() - 1
		c.stack = append(c.stack, nextRef)
	}
}

// next moves to the next leaf element and returns the key and value.
// If the cursor is at the last leaf element then it stays there and returns nil.
func (c *Cursor) next() (key []byte, value []byte, flags uint32) {
	for {
		// Attempt to move over one element until we're successful.
		// Move up the stack as we hit the end of each page in our stack.
		var i int
		for i = len(c.stack) - 1; i >= 0; i-- {
			elem := &c.stack[i]
			if elem.index < elem.count()-1 {
				elem.index++
				break
			}
		}

		// If we've hit the root page then stop and return. This will leave the
		// cursor on the last element of the last page.
		if i == -1 {
			return nil, nil, 0
		}

		// Otherwise start from where we left off in the stack and find the
		// first element of the first leaf page.
		c.stack = c.stack[:i+1]
		c.first()

		// If this is an empty page then restart and move back up the stack.
		// https://github.com/boltdb/bolt/issues/450
		if c.stack[len(c.stack)-1].count() == 0 {
			continue
		}

		return c.keyValue()
	}
}

// search recursively performs a binary search against a given page/node until it finds a given key.
func (c *Cursor) search(key []byte, pgid pgid) {
	p, n := c.bucket.pageNode(pgid)
	if p != nil && (p.flags&(branchPageFlag|leafPageFlag)) == 0 {
		panic(fmt.Sprintf("invalid page type: %d: %x", p.id, p.flags))
	}
	e := elemRef{page: p, node: n}
	c.stack = append(c.stack, e)

	// If we're on a leaf page/node then find the specific node.
	if e.isLeaf() {
		c.nsearch(key)
		return
	}

	if n != nil {
		c.searchNode(key, n)
		return
	}
	c.searchPage(key, p)
}

func (c *Cursor) searchNode(key []byte, n *node) {
	var exact bool
	index := sort.Search(len(n.inodes), func(i int) bool {
		// TODO(benbjohnson): Optimize this range search. It's a bit hacky right now.
		// sort.Search() finds the lowest index where f() != -1 but we need the highest index.
		ret := bytes.Compare(n.inodes[i].key, key)
		if ret == 0 {
			exact = true
		}
		return ret != -1
	})
	if !exact && index > 0 {
		index--
	}
	c.stack[len(c.stack)-1].index = index

	// Recursively search to the next page.
	c.search(key, n.inodes[index].pgid)
}

func (c *Cursor) searchPage(key []byte, p *page) {
	// Binary search for the correct range.
	inodes := p.branchPageElements()

	var exact bool
	index := sort.Search(int(p.count), func(i int) bool {
		// TODO(benbjohnson): Optimize this range search. It's a bit hacky right now.
		// sort.Search() finds the lowest index where f() != -1 but we need the highest index.
		ret := bytes.Compare(inodes[i].key(), key)
		if ret == 0 {
			exact = true
		}
		return ret != -1
	})
	if !exact && index > 0 {
		index--
	}
	c.stack[len(c.stack)-1].index = index

	// Recursively search to the next page.
	c.search(key, inodes[index].pgid)
}

// nsearch searches the leaf node on the top of the stack for a key.
func (c *Cursor) nsearch(key []byte) {
	e := &c.stack[len(c.stack)-1]
	p, n := e.page, e.node

	// If we have a node then search its inodes.
	if n != nil {
		index := sort.Search(len(n.inodes), func(i int) bool {
			return bytes.Compare(n.inodes[i].key, key) != -1
		})
		e.index = index
		return
	}

	// If we have a page then search its leaf elements.
	inodes := p.leafPageElements()
	index := sort.Search(int(p.count), func(i int) bool {
		return bytes.Compare(inodes[i].key(), key) != -1
	})
	e.index = index
}

// keyValue returns the key and value of the current leaf element.
func (c *Cursor) keyValue() ([]byte, []byte, uint32) {
	ref := &c.stack[len(c.stack)-1]
	if ref.count() == 0 || ref.index >= ref.count() {
		return nil, nil, 0
	}

	// Retrieve value from node.
	if ref.node != nil {
		inode := &ref.node.inodes[ref.index]
		return inode.key, inode.value, inode.flags
	}

	// Or retrieve value from page.
	elem := ref.page.leafPageElement(uint16(ref.index))
	return elem.key(), elem.value(), elem.flags
}

// node returns the node that the cursor is currently positioned on.
func (c *Cursor) node() *node {
	_assert(len(c.stack) > 0, "accessing a node with a zero-length cursor stack")

	// If the top of the stack is a leaf node then just return it.
	if ref := &c.stack[len(c.stack)-1]; ref.node != nil && ref.isLeaf() {
		return ref.node
	}

	// Start from root and traverse down the hierarchy.
	var n = c.stack[0].node
	if n == nil {
		n = c.bucket.node(c.stack[0].page.id, nil)
	}
	for _, ref := range c.stack[:len(c.stack)-1] {
		_assert(!n.isLeaf, "expected branch node")
		n = n.childAt(int(ref.index))
	}
	_assert(n.isLeaf, "expected leaf node")
	return n
}

// elemRef represents a reference to an element on a given page/node.
type elemRef struct {
	page  *page
	node  *node
	index int
}

// isLeaf returns whether the ref is pointing at a leaf page/node.
func (r *elemRef) isLeaf() bool {
	if r.node != nil {
		return r.node.isLeaf
	}
	return (r.page.flags & leafPageFlag) != 0
}

// count returns the number of inodes or page elements.
func (r *elemRef) count() int {
	if r.node != nil {
		return len(r.node.inodes)
	}
	return int(r.page.count)
}
<commit_msg>cursor: remove redundant check when seek<commit_after>package bolt

import (
	"bytes"
	"fmt"
	"sort"
)

// Cursor represents an iterator that can traverse over all key/value pairs in a bucket in sorted order.
// Cursors see nested buckets with value == nil.
// Cursors can be obtained from a transaction and are valid as long as the transaction is open.
//
// Keys and values returned from the cursor are only valid for the life of the transaction.
//
// Changing data while traversing with a cursor may cause it to be invalidated
// and return unexpected keys and/or values. You must reposition your cursor
// after mutating data.
type Cursor struct {
	bucket *Bucket
	stack  []elemRef
}

// Bucket returns the bucket that this cursor was created from.
func (c *Cursor) Bucket() *Bucket {
	return c.bucket
}

// First moves the cursor to the first item in the bucket and returns its key and value.
// If the bucket is empty then a nil key and value are returned.
// The returned key and value are only valid for the life of the transaction.
func (c *Cursor) First() (key []byte, value []byte) {
	_assert(c.bucket.tx.db != nil, "tx closed")
	c.stack = c.stack[:0]
	p, n := c.bucket.pageNode(c.bucket.root)
	c.stack = append(c.stack, elemRef{page: p, node: n, index: 0})
	c.first()

	// If we land on an empty page then move to the next value.
	// https://github.com/boltdb/bolt/issues/450
	if c.stack[len(c.stack)-1].count() == 0 {
		c.next()
	}

	k, v, flags := c.keyValue()
	if (flags & uint32(bucketLeafFlag)) != 0 {
		return k, nil
	}
	return k, v

}

// Last moves the cursor to the last item in the bucket and returns its key and value.
// If the bucket is empty then a nil key and value are returned.
// The returned key and value are only valid for the life of the transaction.
func (c *Cursor) Last() (key []byte, value []byte) {
	_assert(c.bucket.tx.db != nil, "tx closed")
	c.stack = c.stack[:0]
	p, n := c.bucket.pageNode(c.bucket.root)
	ref := elemRef{page: p, node: n}
	ref.index = ref.count() - 1
	c.stack = append(c.stack, ref)
	c.last()
	k, v, flags := c.keyValue()
	if (flags & uint32(bucketLeafFlag)) != 0 {
		return k, nil
	}
	return k, v
}

// Next moves the cursor to the next item in the bucket and returns its key and value.
// If the cursor is at the end of the bucket then a nil key and value are returned.
// The returned key and value are only valid for the life of the transaction.
func (c *Cursor) Next() (key []byte, value []byte) {
	_assert(c.bucket.tx.db != nil, "tx closed")
	k, v, flags := c.next()
	if (flags & uint32(bucketLeafFlag)) != 0 {
		return k, nil
	}
	return k, v
}

// Prev moves the cursor to the previous item in the bucket and returns its key and value.
// If the cursor is at the beginning of the bucket then a nil key and value are returned.
// The returned key and value are only valid for the life of the transaction.
func (c *Cursor) Prev() (key []byte, value []byte) {
	_assert(c.bucket.tx.db != nil, "tx closed")

	// Attempt to move back one element until we're successful.
	// Move up the stack as we hit the beginning of each page in our stack.
	for i := len(c.stack) - 1; i >= 0; i-- {
		elem := &c.stack[i]
		if elem.index > 0 {
			elem.index--
			break
		}
		c.stack = c.stack[:i]
	}

	// If we've hit the end then return nil.
	if len(c.stack) == 0 {
		return nil, nil
	}

	// Move down the stack to find the last element of the last leaf under this branch.
	c.last()
	k, v, flags := c.keyValue()
	if (flags & uint32(bucketLeafFlag)) != 0 {
		return k, nil
	}
	return k, v
}

// Seek moves the cursor to a given key and returns it.
// If the key does not exist then the next key is used. If no keys
// follow, a nil key is returned.
// The returned key and value are only valid for the life of the transaction.
func (c *Cursor) Seek(seek []byte) (key []byte, value []byte) {
	k, v, flags := c.seek(seek)

	// If we ended up after the last element of a page then move to the next one.
	if ref := &c.stack[len(c.stack)-1]; ref.index >= ref.count() {
		k, v, flags = c.next()
	}

	if k == nil {
		return nil, nil
	} else if (flags & uint32(bucketLeafFlag)) != 0 {
		return k, nil
	}
	return k, v
}

// Delete removes the current key/value under the cursor from the bucket.
// Delete fails if current key/value is a bucket or if the transaction is not writable.
func (c *Cursor) Delete() error {
	if c.bucket.tx.db == nil {
		return ErrTxClosed
	} else if !c.bucket.Writable() {
		return ErrTxNotWritable
	}

	key, _, flags := c.keyValue()
	// Return an error if current value is a bucket.
	if (flags & bucketLeafFlag) != 0 {
		return ErrIncompatibleValue
	}
	c.node().del(key)

	return nil
}

// seek moves the cursor to a given key and returns it.
// If the key does not exist then the next key is used.
func (c *Cursor) seek(seek []byte) (key []byte, value []byte, flags uint32) {
	_assert(c.bucket.tx.db != nil, "tx closed")

	// Start from root page/node and traverse to correct page.
	c.stack = c.stack[:0]
	c.search(seek, c.bucket.root)

	// If this is a bucket then return a nil value.
	return c.keyValue()
}

// first moves the cursor to the first leaf element under the last page in the stack.
func (c *Cursor) first() {
	for {
		// Exit when we hit a leaf page.
		var ref = &c.stack[len(c.stack)-1]
		if ref.isLeaf() {
			break
		}

		// Keep adding pages pointing to the first element to the stack.
		var pgid pgid
		if ref.node != nil {
			pgid = ref.node.inodes[ref.index].pgid
		} else {
			pgid = ref.page.branchPageElement(uint16(ref.index)).pgid
		}
		p, n := c.bucket.pageNode(pgid)
		c.stack = append(c.stack, elemRef{page: p, node: n, index: 0})
	}
}

// last moves the cursor to the last leaf element under the last page in the stack.
func (c *Cursor) last() {
	for {
		// Exit when we hit a leaf page.
		ref := &c.stack[len(c.stack)-1]
		if ref.isLeaf() {
			break
		}

		// Keep adding pages pointing to the last element in the stack.
		var pgid pgid
		if ref.node != nil {
			pgid = ref.node.inodes[ref.index].pgid
		} else {
			pgid = ref.page.branchPageElement(uint16(ref.index)).pgid
		}
		p, n := c.bucket.pageNode(pgid)

		var nextRef = elemRef{page: p, node: n}
		nextRef.index = nextRef.count() - 1
		c.stack = append(c.stack, nextRef)
	}
}

// next moves to the next leaf element and returns the key and value.
// If the cursor is at the last leaf element then it stays there and returns nil.
func (c *Cursor) next() (key []byte, value []byte, flags uint32) {
	for {
		// Attempt to move over one element until we're successful.
		// Move up the stack as we hit the end of each page in our stack.
		var i int
		for i = len(c.stack) - 1; i >= 0; i-- {
			elem := &c.stack[i]
			if elem.index < elem.count()-1 {
				elem.index++
				break
			}
		}

		// If we've hit the root page then stop and return. This will leave the
		// cursor on the last element of the last page.
		if i == -1 {
			return nil, nil, 0
		}

		// Otherwise start from where we left off in the stack and find the
		// first element of the first leaf page.
		c.stack = c.stack[:i+1]
		c.first()

		// If this is an empty page then restart and move back up the stack.
		// https://github.com/boltdb/bolt/issues/450
		if c.stack[len(c.stack)-1].count() == 0 {
			continue
		}

		return c.keyValue()
	}
}

// search recursively performs a binary search against a given page/node until it finds a given key.
func (c *Cursor) search(key []byte, pgid pgid) {
	p, n := c.bucket.pageNode(pgid)
	if p != nil && (p.flags&(branchPageFlag|leafPageFlag)) == 0 {
		panic(fmt.Sprintf("invalid page type: %d: %x", p.id, p.flags))
	}
	e := elemRef{page: p, node: n}
	c.stack = append(c.stack, e)

	// If we're on a leaf page/node then find the specific node.
	if e.isLeaf() {
		c.nsearch(key)
		return
	}

	if n != nil {
		c.searchNode(key, n)
		return
	}
	c.searchPage(key, p)
}

func (c *Cursor) searchNode(key []byte, n *node) {
	var exact bool
	index := sort.Search(len(n.inodes), func(i int) bool {
		// TODO(benbjohnson): Optimize this range search. It's a bit hacky right now.
		// sort.Search() finds the lowest index where f() != -1 but we need the highest index.
		ret := bytes.Compare(n.inodes[i].key, key)
		if ret == 0 {
			exact = true
		}
		return ret != -1
	})
	if !exact && index > 0 {
		index--
	}
	c.stack[len(c.stack)-1].index = index

	// Recursively search to the next page.
	c.search(key, n.inodes[index].pgid)
}

func (c *Cursor) searchPage(key []byte, p *page) {
	// Binary search for the correct range.
	inodes := p.branchPageElements()

	var exact bool
	index := sort.Search(int(p.count), func(i int) bool {
		// TODO(benbjohnson): Optimize this range search. It's a bit hacky right now.
		// sort.Search() finds the lowest index where f() != -1 but we need the highest index.
		ret := bytes.Compare(inodes[i].key(), key)
		if ret == 0 {
			exact = true
		}
		return ret != -1
	})
	if !exact && index > 0 {
		index--
	}
	c.stack[len(c.stack)-1].index = index

	// Recursively search to the next page.
	c.search(key, inodes[index].pgid)
}

// nsearch searches the leaf node on the top of the stack for a key.
func (c *Cursor) nsearch(key []byte) {
	e := &c.stack[len(c.stack)-1]
	p, n := e.page, e.node

	// If we have a node then search its inodes.
	if n != nil {
		index := sort.Search(len(n.inodes), func(i int) bool {
			return bytes.Compare(n.inodes[i].key, key) != -1
		})
		e.index = index
		return
	}

	// If we have a page then search its leaf elements.
	inodes := p.leafPageElements()
	index := sort.Search(int(p.count), func(i int) bool {
		return bytes.Compare(inodes[i].key(), key) != -1
	})
	e.index = index
}

// keyValue returns the key and value of the current leaf element.
func (c *Cursor) keyValue() ([]byte, []byte, uint32) {
	ref := &c.stack[len(c.stack)-1]

	// If the cursor is pointing to the end of page/node then return nil.
	if ref.count() == 0 || ref.index >= ref.count() {
		return nil, nil, 0
	}

	// Retrieve value from node.
	if ref.node != nil {
		inode := &ref.node.inodes[ref.index]
		return inode.key, inode.value, inode.flags
	}

	// Or retrieve value from page.
	elem := ref.page.leafPageElement(uint16(ref.index))
	return elem.key(), elem.value(), elem.flags
}

// node returns the node that the cursor is currently positioned on.
func (c *Cursor) node() *node {
	_assert(len(c.stack) > 0, "accessing a node with a zero-length cursor stack")

	// If the top of the stack is a leaf node then just return it.
	if ref := &c.stack[len(c.stack)-1]; ref.node != nil && ref.isLeaf() {
		return ref.node
	}

	// Start from root and traverse down the hierarchy.
	var n = c.stack[0].node
	if n == nil {
		n = c.bucket.node(c.stack[0].page.id, nil)
	}
	for _, ref := range c.stack[:len(c.stack)-1] {
		_assert(!n.isLeaf, "expected branch node")
		n = n.childAt(int(ref.index))
	}
	_assert(n.isLeaf, "expected leaf node")
	return n
}

// elemRef represents a reference to an element on a given page/node.
type elemRef struct {
	page  *page
	node  *node
	index int
}

// isLeaf returns whether the ref is pointing at a leaf page/node.
func (r *elemRef) isLeaf() bool {
	if r.node != nil {
		return r.node.isLeaf
	}
	return (r.page.flags & leafPageFlag) != 0
}

// count returns the number of inodes or page elements.
func (r *elemRef) count() int {
	if r.node != nil {
		return len(r.node.inodes)
	}
	return int(r.page.count)
}
<|endoftext|>


<commit_before>// Licensed to the Apache Software Foundation (ASF) under one or more
// contributor license agreements.  See the NOTICE file distributed with
// this work for additional information regarding copyright ownership.
// The ASF licenses this file to You under the Apache License, Version 2.0
// (the "License"); you may not use this file except in compliance with
// the License.  You may obtain a copy of the License at
//
//    http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

// boot is the boot code for the Java SDK harness container. It is responsible
// for retrieving staged files and invoking the JVM correctly.
package main

import (
	"context"
	"flag"
	"fmt"
	"log"
	"os"
	"path/filepath"
	"strconv"
	"strings"

	"github.com/apache/beam/sdks/v2/go/pkg/beam/artifact"
	fnpb "github.com/apache/beam/sdks/v2/go/pkg/beam/model/fnexecution_v1"
	pipepb "github.com/apache/beam/sdks/v2/go/pkg/beam/model/pipeline_v1"
	"github.com/apache/beam/sdks/v2/go/pkg/beam/provision"
	"github.com/apache/beam/sdks/v2/go/pkg/beam/util/execx"
	"github.com/apache/beam/sdks/v2/go/pkg/beam/util/grpcx"
	"github.com/apache/beam/sdks/v2/go/pkg/beam/util/syscallx"
	"github.com/golang/protobuf/proto"
)

var (
	// Contract: https://s.apache.org/beam-fn-api-container-contract.

	id                = flag.String("id", "", "Local identifier (required).")
	loggingEndpoint   = flag.String("logging_endpoint", "", "Logging endpoint (required).")
	artifactEndpoint  = flag.String("artifact_endpoint", "", "Artifact endpoint (required).")
	provisionEndpoint = flag.String("provision_endpoint", "", "Provision endpoint (required).")
	controlEndpoint   = flag.String("control_endpoint", "", "Control endpoint (required).")
	semiPersistDir    = flag.String("semi_persist_dir", "/tmp", "Local semi-persistent directory (optional).")
)

const (
	enableGoogleCloudProfilerOption     = "enable_google_cloud_profiler"
	enableGoogleCloudHeapSamplingOption = "enable_google_cloud_heap_sampling"
	googleCloudProfilerAgentBaseArgs    = "-agentpath:/opt/google_cloud_profiler/profiler_java_agent.so=-logtostderr,-cprof_service=%s,-cprof_service_version=%s"
	googleCloudProfilerAgentHeapArgs    = googleCloudProfilerAgentBaseArgs + ",-cprof_enable_heap_sampling,-cprof_heap_sampling_interval=2097152"
)

func main() {
	flag.Parse()
	if *id == "" {
		log.Fatal("No id provided.")
	}
	if *provisionEndpoint == "" {
		log.Fatal("No provision endpoint provided.")
	}

	ctx := grpcx.WriteWorkerID(context.Background(), *id)

	info, err := provision.Info(ctx, *provisionEndpoint)
	if err != nil {
		log.Fatalf("Failed to obtain provisioning information: %v", err)
	}
	log.Printf("Provision info:\n%v", info)

	// TODO(BEAM-8201): Simplify once flags are no longer used.
	if info.GetLoggingEndpoint().GetUrl() != "" {
		*loggingEndpoint = info.GetLoggingEndpoint().GetUrl()
	}
	if info.GetArtifactEndpoint().GetUrl() != "" {
		*artifactEndpoint = info.GetArtifactEndpoint().GetUrl()
	}
	if info.GetControlEndpoint().GetUrl() != "" {
		*controlEndpoint = info.GetControlEndpoint().GetUrl()
	}

	if *loggingEndpoint == "" {
		log.Fatal("No logging endpoint provided.")
	}
	if *artifactEndpoint == "" {
		log.Fatal("No artifact endpoint provided.")
	}
	if *controlEndpoint == "" {
		log.Fatal("No control endpoint provided.")
	}

	log.Printf("Initializing java harness: %v", strings.Join(os.Args, " "))

	// (1) Obtain the pipeline options

	options, err := provision.ProtoToJSON(info.GetPipelineOptions())
	if err != nil {
		log.Fatalf("Failed to convert pipeline options: %v", err)
	}

	// (2) Retrieve the staged user jars. We ignore any disk limit,
	// because the staged jars are mandatory.

	// Using the SDK Harness ID in the artifact destination path to make sure that dependencies used by multiple
	// SDK Harnesses in the same VM do not conflict. This is needed since some runners (for example, Dataflow)
	// may share the artifact staging directory across multiple SDK Harnesses
	// TODO(BEAM-9455): consider removing the SDK Harness ID from the staging path after Dataflow can properly
	// seperate out dependencies per environment.
	dir := filepath.Join(*semiPersistDir, *id, "staged")

	artifacts, err := artifact.Materialize(ctx, *artifactEndpoint, info.GetDependencies(), info.GetRetrievalToken(), dir)
	if err != nil {
		log.Fatalf("Failed to retrieve staged files: %v", err)
	}

	// (3) Invoke the Java harness, preserving artifact ordering in classpath.

	os.Setenv("HARNESS_ID", *id)
	os.Setenv("PIPELINE_OPTIONS", options)
	os.Setenv("LOGGING_API_SERVICE_DESCRIPTOR", proto.MarshalTextString(&pipepb.ApiServiceDescriptor{Url: *loggingEndpoint}))
	os.Setenv("CONTROL_API_SERVICE_DESCRIPTOR", proto.MarshalTextString(&pipepb.ApiServiceDescriptor{Url: *controlEndpoint}))
	os.Setenv("RUNNER_CAPABILITIES", strings.Join(info.GetRunnerCapabilities(), " "))

	if info.GetStatusEndpoint() != nil {
		os.Setenv("STATUS_API_SERVICE_DESCRIPTOR", proto.MarshalTextString(info.GetStatusEndpoint()))
	}

	const jarsDir = "/opt/apache/beam/jars"
	cp := []string{
		filepath.Join(jarsDir, "slf4j-api.jar"),
		filepath.Join(jarsDir, "slf4j-jdk14.jar"),
		filepath.Join(jarsDir, "beam-sdks-java-harness.jar"),
		filepath.Join(jarsDir, "beam-sdks-java-io-kafka.jar"),
		filepath.Join(jarsDir, "kafka-clients.jar"),
	}

	var hasWorkerExperiment = strings.Contains(options, "use_staged_dataflow_worker_jar")
	for _, a := range artifacts {
		name, _ := artifact.MustExtractFilePayload(a)
		if hasWorkerExperiment {
			if strings.HasPrefix(name, "beam-runners-google-cloud-dataflow-java-fn-api-worker") {
				continue
			}
			if name == "dataflow-worker.jar" {
				continue
			}
		}
		cp = append(cp, filepath.Join(dir, filepath.FromSlash(name)))
	}

	args := []string{
		"-Xmx" + strconv.FormatUint(heapSizeLimit(info), 10),
		"-XX:-OmitStackTraceInFastThrow",
		"-cp", strings.Join(cp, ":"),
	}

	enableGoogleCloudProfiler := strings.Contains(options, enableGoogleCloudProfilerOption)
	enableGoogleCloudHeapSampling := strings.Contains(options, enableGoogleCloudHeapSamplingOption)
	if enableGoogleCloudProfiler {
		if metadata := info.GetMetadata(); metadata != nil {
			if jobName, nameExists := metadata["job_name"]; nameExists {
				if jobId, idExists := metadata["job_id"]; idExists {
					if enableGoogleCloudHeapSampling {
						args = append(args, fmt.Sprintf(googleCloudProfilerAgentHeapArgs, jobName, jobId))
					} else {
						args = append(args, fmt.Sprintf(googleCloudProfilerAgentBaseArgs, jobName, jobId))
					}
					log.Printf("Turning on Cloud Profiling. Profile heap: %t", enableGoogleCloudHeapSampling)
				} else {
					log.Println("Required job_id missing from metadata, profiling will not be enabled without it.")
				}
			} else {
				log.Println("Required job_name missing from metadata, profiling will not be enabled without it.")
			}
		} else {
			log.Println("enable_google_cloud_profiler is set to true, but no metadata is received from provision server, profiling will not be enabled.")
		}
	}

	args = append(args, "org.apache.beam.fn.harness.FnHarness")
	log.Printf("Executing: java %v", strings.Join(args, " "))

	log.Fatalf("Java exited: %v", execx.Execute("java", args...))
}

// heapSizeLimit returns 80% of the runner limit, if provided. If not provided,
// it returns 70% of the physical memory on the machine. If it cannot determine
// that value, it returns 1GB. This is an imperfect heuristic. It aims to
// ensure there is memory for non-heap use and other overhead, while also not
// underutilizing the machine.
func heapSizeLimit(info *fnpb.ProvisionInfo) uint64 {
	if size, err := syscallx.PhysicalMemorySize(); err == nil {
		return (size * 70) / 100
	}
	return 1 << 30
}
<commit_msg>Add -XX:+AlwaysActAsServerClassMachine to Java SDK container<commit_after>// Licensed to the Apache Software Foundation (ASF) under one or more
// contributor license agreements.  See the NOTICE file distributed with
// this work for additional information regarding copyright ownership.
// The ASF licenses this file to You under the Apache License, Version 2.0
// (the "License"); you may not use this file except in compliance with
// the License.  You may obtain a copy of the License at
//
//    http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

// boot is the boot code for the Java SDK harness container. It is responsible
// for retrieving staged files and invoking the JVM correctly.
package main

import (
	"context"
	"flag"
	"fmt"
	"log"
	"os"
	"path/filepath"
	"strconv"
	"strings"

	"github.com/apache/beam/sdks/v2/go/pkg/beam/artifact"
	fnpb "github.com/apache/beam/sdks/v2/go/pkg/beam/model/fnexecution_v1"
	pipepb "github.com/apache/beam/sdks/v2/go/pkg/beam/model/pipeline_v1"
	"github.com/apache/beam/sdks/v2/go/pkg/beam/provision"
	"github.com/apache/beam/sdks/v2/go/pkg/beam/util/execx"
	"github.com/apache/beam/sdks/v2/go/pkg/beam/util/grpcx"
	"github.com/apache/beam/sdks/v2/go/pkg/beam/util/syscallx"
	"github.com/golang/protobuf/proto"
)

var (
	// Contract: https://s.apache.org/beam-fn-api-container-contract.

	id                = flag.String("id", "", "Local identifier (required).")
	loggingEndpoint   = flag.String("logging_endpoint", "", "Logging endpoint (required).")
	artifactEndpoint  = flag.String("artifact_endpoint", "", "Artifact endpoint (required).")
	provisionEndpoint = flag.String("provision_endpoint", "", "Provision endpoint (required).")
	controlEndpoint   = flag.String("control_endpoint", "", "Control endpoint (required).")
	semiPersistDir    = flag.String("semi_persist_dir", "/tmp", "Local semi-persistent directory (optional).")
)

const (
	enableGoogleCloudProfilerOption     = "enable_google_cloud_profiler"
	enableGoogleCloudHeapSamplingOption = "enable_google_cloud_heap_sampling"
	googleCloudProfilerAgentBaseArgs    = "-agentpath:/opt/google_cloud_profiler/profiler_java_agent.so=-logtostderr,-cprof_service=%s,-cprof_service_version=%s"
	googleCloudProfilerAgentHeapArgs    = googleCloudProfilerAgentBaseArgs + ",-cprof_enable_heap_sampling,-cprof_heap_sampling_interval=2097152"
)

func main() {
	flag.Parse()
	if *id == "" {
		log.Fatal("No id provided.")
	}
	if *provisionEndpoint == "" {
		log.Fatal("No provision endpoint provided.")
	}

	ctx := grpcx.WriteWorkerID(context.Background(), *id)

	info, err := provision.Info(ctx, *provisionEndpoint)
	if err != nil {
		log.Fatalf("Failed to obtain provisioning information: %v", err)
	}
	log.Printf("Provision info:\n%v", info)

	// TODO(BEAM-8201): Simplify once flags are no longer used.
	if info.GetLoggingEndpoint().GetUrl() != "" {
		*loggingEndpoint = info.GetLoggingEndpoint().GetUrl()
	}
	if info.GetArtifactEndpoint().GetUrl() != "" {
		*artifactEndpoint = info.GetArtifactEndpoint().GetUrl()
	}
	if info.GetControlEndpoint().GetUrl() != "" {
		*controlEndpoint = info.GetControlEndpoint().GetUrl()
	}

	if *loggingEndpoint == "" {
		log.Fatal("No logging endpoint provided.")
	}
	if *artifactEndpoint == "" {
		log.Fatal("No artifact endpoint provided.")
	}
	if *controlEndpoint == "" {
		log.Fatal("No control endpoint provided.")
	}

	log.Printf("Initializing java harness: %v", strings.Join(os.Args, " "))

	// (1) Obtain the pipeline options

	options, err := provision.ProtoToJSON(info.GetPipelineOptions())
	if err != nil {
		log.Fatalf("Failed to convert pipeline options: %v", err)
	}

	// (2) Retrieve the staged user jars. We ignore any disk limit,
	// because the staged jars are mandatory.

	// Using the SDK Harness ID in the artifact destination path to make sure that dependencies used by multiple
	// SDK Harnesses in the same VM do not conflict. This is needed since some runners (for example, Dataflow)
	// may share the artifact staging directory across multiple SDK Harnesses
	// TODO(BEAM-9455): consider removing the SDK Harness ID from the staging path after Dataflow can properly
	// seperate out dependencies per environment.
	dir := filepath.Join(*semiPersistDir, *id, "staged")

	artifacts, err := artifact.Materialize(ctx, *artifactEndpoint, info.GetDependencies(), info.GetRetrievalToken(), dir)
	if err != nil {
		log.Fatalf("Failed to retrieve staged files: %v", err)
	}

	// (3) Invoke the Java harness, preserving artifact ordering in classpath.

	os.Setenv("HARNESS_ID", *id)
	os.Setenv("PIPELINE_OPTIONS", options)
	os.Setenv("LOGGING_API_SERVICE_DESCRIPTOR", proto.MarshalTextString(&pipepb.ApiServiceDescriptor{Url: *loggingEndpoint}))
	os.Setenv("CONTROL_API_SERVICE_DESCRIPTOR", proto.MarshalTextString(&pipepb.ApiServiceDescriptor{Url: *controlEndpoint}))
	os.Setenv("RUNNER_CAPABILITIES", strings.Join(info.GetRunnerCapabilities(), " "))

	if info.GetStatusEndpoint() != nil {
		os.Setenv("STATUS_API_SERVICE_DESCRIPTOR", proto.MarshalTextString(info.GetStatusEndpoint()))
	}

	const jarsDir = "/opt/apache/beam/jars"
	cp := []string{
		filepath.Join(jarsDir, "slf4j-api.jar"),
		filepath.Join(jarsDir, "slf4j-jdk14.jar"),
		filepath.Join(jarsDir, "beam-sdks-java-harness.jar"),
		filepath.Join(jarsDir, "beam-sdks-java-io-kafka.jar"),
		filepath.Join(jarsDir, "kafka-clients.jar"),
	}

	var hasWorkerExperiment = strings.Contains(options, "use_staged_dataflow_worker_jar")
	for _, a := range artifacts {
		name, _ := artifact.MustExtractFilePayload(a)
		if hasWorkerExperiment {
			if strings.HasPrefix(name, "beam-runners-google-cloud-dataflow-java-fn-api-worker") {
				continue
			}
			if name == "dataflow-worker.jar" {
				continue
			}
		}
		cp = append(cp, filepath.Join(dir, filepath.FromSlash(name)))
	}

	args := []string{
		"-Xmx" + strconv.FormatUint(heapSizeLimit(info), 10),
		"-XX:+AlwaysActAsServerClassMachine",
		"-XX:-OmitStackTraceInFastThrow",
		"-cp", strings.Join(cp, ":"),
	}

	enableGoogleCloudProfiler := strings.Contains(options, enableGoogleCloudProfilerOption)
	enableGoogleCloudHeapSampling := strings.Contains(options, enableGoogleCloudHeapSamplingOption)
	if enableGoogleCloudProfiler {
		if metadata := info.GetMetadata(); metadata != nil {
			if jobName, nameExists := metadata["job_name"]; nameExists {
				if jobId, idExists := metadata["job_id"]; idExists {
					if enableGoogleCloudHeapSampling {
						args = append(args, fmt.Sprintf(googleCloudProfilerAgentHeapArgs, jobName, jobId))
					} else {
						args = append(args, fmt.Sprintf(googleCloudProfilerAgentBaseArgs, jobName, jobId))
					}
					log.Printf("Turning on Cloud Profiling. Profile heap: %t", enableGoogleCloudHeapSampling)
				} else {
					log.Println("Required job_id missing from metadata, profiling will not be enabled without it.")
				}
			} else {
				log.Println("Required job_name missing from metadata, profiling will not be enabled without it.")
			}
		} else {
			log.Println("enable_google_cloud_profiler is set to true, but no metadata is received from provision server, profiling will not be enabled.")
		}
	}

	args = append(args, "org.apache.beam.fn.harness.FnHarness")
	log.Printf("Executing: java %v", strings.Join(args, " "))

	log.Fatalf("Java exited: %v", execx.Execute("java", args...))
}

// heapSizeLimit returns 80% of the runner limit, if provided. If not provided,
// it returns 70% of the physical memory on the machine. If it cannot determine
// that value, it returns 1GB. This is an imperfect heuristic. It aims to
// ensure there is memory for non-heap use and other overhead, while also not
// underutilizing the machine.
func heapSizeLimit(info *fnpb.ProvisionInfo) uint64 {
	if size, err := syscallx.PhysicalMemorySize(); err == nil {
		return (size * 70) / 100
	}
	return 1 << 30
}
<|endoftext|>


<commit_before>// Copyright (c) 2016, Daniel Martí <mvdan@mvdan.cc>
// See LICENSE for licensing information

//go:build go1.18
// +build go1.18

package syntax

import (
	"io"
	"os/exec"
	"strings"
	"testing"
)

func FuzzQuote(f *testing.F) {
	if _, err := exec.LookPath("bash"); err != nil {
		f.Skipf("requires bash to verify quoted strings")
	}

	// Keep in sync with ExampleQuote.
	f.Add("foo")
	f.Add("bar $baz")
	f.Add(`"won't"`)
	f.Add(`~/home`)
	f.Add("#1304")
	f.Add("name=value")
	f.Add(`glob-*`)
	f.Add("invalid-\xe2'")
	f.Add("nonprint-\x0b\x1b")
	f.Fuzz(func(t *testing.T, s string) {
		quoted, ok := Quote(s)
		if !ok {
			// Contains a null byte; not interesting.
			t.Skip()
		}
		// Beware that this might run arbitrary code
		// if Quote is too naive and allows ';' or '$'.
		//
		// Also note that this fuzzing would not catch '=',
		// as we don't use the quoted string as a first argument
		// to avoid running random commands.
		//
		// We could consider ways to fully sandbox the bash process,
		// but for now that feels overkill.
		out, err := exec.Command("bash", "-c", "printf %s "+quoted).CombinedOutput()
		if err != nil {
			t.Fatalf("bash error on %q quoted as %s: %v: %s", s, quoted, err, out)
		}
		want, got := s, string(out)
		if want != got {
			t.Fatalf("output mismatch on %q quoted as %s: got %q (len=%d)", want, quoted, got, len(got))
		}
	})
}

func FuzzParsePrint(f *testing.F) {
	// TODO: probably use f.Add on table-driven test cases too?
	// In the past, any crashers found by go-fuzz got put there.

	add := func(src string) {
		// For now, default to just LangBash with KeepComments.
		f.Add(src, uint8(LangBash), true, false,
			uint8(0), false, false, false, false, false, false, false)
	}

	add("echo foo")
	add("'foo' \"bar\" $'baz'")
	add("if foo; then bar; baz; fi")
	add("{ (foo; bar); baz; }")
	add("$foo ${bar} ${baz[@]} $")
	add("foo >bar <<EOF\nbaz\nEOF")
	add("foo=bar baz=(x y z)")
	add("foo && bar || baz")
	add("foo | bar # baz")
	add("foo \\\n bar \\\\ba\\z")

	f.Fuzz(func(t *testing.T,
		src string,

		// parser options
		// TODO: also fuzz StopAt
		langVariant uint8, // 0-3
		keepComments bool,

		simplify bool,

		// printer options
		indent uint8, // 0-255
		binaryNextLine bool,
		switchCaseIndent bool,
		spaceRedirects bool,
		keepPadding bool,
		minify bool,
		singleLine bool,
		functionNextLine bool,
	) {
		if langVariant > 3 {
			t.Skip() // lang variants are 0-3
		}
		if indent > 16 {
			t.Skip() // more indentation won't really be interesting
		}

		parser := NewParser()
		Variant(LangVariant(langVariant))(parser)
		KeepComments(keepComments)(parser)

		prog, err := parser.Parse(strings.NewReader(src), "")
		if err != nil {
			t.Skip() // not valid shell syntax
		}

		if simplify {
			Simplify(prog)
		}

		printer := NewPrinter()
		Indent(uint(indent))(printer)
		BinaryNextLine(binaryNextLine)(printer)
		SwitchCaseIndent(switchCaseIndent)(printer)
		SpaceRedirects(spaceRedirects)(printer)
		KeepPadding(keepPadding)(printer)
		Minify(minify)(printer)
		SingleLine(singleLine)(printer)
		FunctionNextLine(functionNextLine)(printer)

		if err := printer.Print(io.Discard, prog); err != nil {
			t.Skip() // e.g. invalid option
		}
	})
}
<commit_msg>syntax: feed the fuzzer with lots of initial corpus<commit_after>// Copyright (c) 2016, Daniel Martí <mvdan@mvdan.cc>
// See LICENSE for licensing information

//go:build go1.18
// +build go1.18

package syntax

import (
	"io"
	"os/exec"
	"strings"
	"testing"
)

func FuzzQuote(f *testing.F) {
	if _, err := exec.LookPath("bash"); err != nil {
		f.Skipf("requires bash to verify quoted strings")
	}

	// Keep in sync with ExampleQuote.
	f.Add("foo")
	f.Add("bar $baz")
	f.Add(`"won't"`)
	f.Add(`~/home`)
	f.Add("#1304")
	f.Add("name=value")
	f.Add(`glob-*`)
	f.Add("invalid-\xe2'")
	f.Add("nonprint-\x0b\x1b")
	f.Fuzz(func(t *testing.T, s string) {
		quoted, ok := Quote(s)
		if !ok {
			// Contains a null byte; not interesting.
			t.Skip()
		}
		// Beware that this might run arbitrary code
		// if Quote is too naive and allows ';' or '$'.
		//
		// Also note that this fuzzing would not catch '=',
		// as we don't use the quoted string as a first argument
		// to avoid running random commands.
		//
		// We could consider ways to fully sandbox the bash process,
		// but for now that feels overkill.
		out, err := exec.Command("bash", "-c", "printf %s "+quoted).CombinedOutput()
		if err != nil {
			t.Fatalf("bash error on %q quoted as %s: %v: %s", s, quoted, err, out)
		}
		want, got := s, string(out)
		if want != got {
			t.Fatalf("output mismatch on %q quoted as %s: got %q (len=%d)", want, quoted, got, len(got))
		}
	})
}

func FuzzParsePrint(f *testing.F) {
	// TODO: probably use f.Add on table-driven test cases too?
	// In the past, any crashers found by go-fuzz got put there.

	add := func(src string, variant LangVariant) {
		// For now, default to just KeepComments.
		f.Add(src, uint8(variant), true, false,
			uint8(0), false, false, false, false, false, false, false)
	}

	for _, test := range shellTests {
		add(test.in, LangBash)
	}
	for _, test := range printTests {
		add(test.in, LangBash)
	}
	for _, test := range fileTests {
		for _, in := range test.Strs {
			if test.Bash != nil {
				add(in, LangBash)
			}
			if test.Posix != nil {
				add(in, LangPOSIX)
			}
			if test.MirBSDKorn != nil {
				add(in, LangMirBSDKorn)
			}
			if test.Bats != nil {
				add(in, LangBats)
			}
		}
	}

	f.Fuzz(func(t *testing.T,
		src string,

		// parser options
		// TODO: also fuzz StopAt
		langVariant uint8, // 0-3
		keepComments bool,

		simplify bool,

		// printer options
		indent uint8, // 0-255
		binaryNextLine bool,
		switchCaseIndent bool,
		spaceRedirects bool,
		keepPadding bool,
		minify bool,
		singleLine bool,
		functionNextLine bool,
	) {
		if langVariant > 3 {
			t.Skip() // lang variants are 0-3
		}
		if indent > 16 {
			t.Skip() // more indentation won't really be interesting
		}

		parser := NewParser()
		Variant(LangVariant(langVariant))(parser)
		KeepComments(keepComments)(parser)

		prog, err := parser.Parse(strings.NewReader(src), "")
		if err != nil {
			t.Skip() // not valid shell syntax
		}

		if simplify {
			Simplify(prog)
		}

		printer := NewPrinter()
		Indent(uint(indent))(printer)
		BinaryNextLine(binaryNextLine)(printer)
		SwitchCaseIndent(switchCaseIndent)(printer)
		SpaceRedirects(spaceRedirects)(printer)
		KeepPadding(keepPadding)(printer)
		Minify(minify)(printer)
		SingleLine(singleLine)(printer)
		FunctionNextLine(functionNextLine)(printer)

		if err := printer.Print(io.Discard, prog); err != nil {
			t.Skip() // e.g. invalid option
		}
	})
}
<|endoftext|>