<commit_before>package dlidparser

import (
	"errors"
	"strconv"
	"strings"
	"time"
)

func parseV3(data string, issuer string) (license *DLIDLicense, err error) {

	start, end, err := dataRangeV2(data)

	if end >= len(data) {
		err = errors.New("Payload location does not exist in data")
	}

	payload := data[start:end]

	if err != nil {
		return
	}

	license, err = parseDataV3(payload, issuer)

	if err != nil {
		return
	}

	return
}

func parseDataV3(licenceData string, issuer string) (license *DLIDLicense, err error) {

	// Version 3 of the DLID card spec was published in 2005.  It is currently
	// (as of 2012) used in Wisconsin.

	if !strings.HasPrefix(licenceData, "DL") {
		err = errors.New("Missing header in licence data chunk")
		return
	}

	licenceData = licenceData[2:]

	components := strings.Split(licenceData, "\n")

	license = new(DLIDLicense)

	license.SetIssuerId(issuer)
	license.SetIssuerName(issuers[issuer])

	var dateOfBirth string
	var expiryDate string
	var issueDate string

	for component := range components {

		if len(components[component]) < 3 {
			continue
		}

		identifier := components[component][0:3]
		data := components[component][3:]

		data = strings.Trim(data, " ")

		switch identifier {
		case "DCA":
			license.SetVehicleClass(data)

		case "DCB":
			license.SetRestrictionCodes(data)

		case "DCD":
			license.SetEndorsementCodes(data)

		case "DCS":
			license.SetLastName(data)

		case "DCG":
			license.SetCountry(data)

		case "DCT":

			// This field contains all of the licencee's names except last
			// name.  The V3 spec doc doesn't specify how the names are
			// separated and doesn't provide an example (unlike the 2000
			// doc).  Wisconsin use a space and Virginia use a comma.  This
			// is why standards have to be adequately documented.

			separator := " "

			if strings.Index(data, separator) == -1 {
				separator = ","
			}

			names := strings.Split(data, separator)

			license.SetFirstName(names[0])

			if len(names) > 1 {
				license.SetMiddleNames(names[1:])
			}

		case "DAG":
			license.SetStreet(data)

		case "DAI":
			license.SetCity(data)

		case "DAJ":
			license.SetState(data)

		case "DAK":
			license.SetPostal(data)

		case "DAQ":
			license.SetCustomerId(data)

		case "DBA":
			expiryDate = data

		case "DBB":
			dateOfBirth = data

		case "DBC":
			switch data {
			case "1":
				license.SetSex(DriverSexMale)
			case "2":
				license.SetSex(DriverSexFemale)
			default:
				license.SetSex(DriverSexNone)
			}

		case "DBD":
			issueDate = data
		}
	}

	// At this point we should know the country and the postal code (both are
	// mandatory fields) so we can undo the desperate mess the standards body
	// made of the postal code field.

	if license.Country() == "USA" && len(license.Postal()) > 0 {

		// For some reason known only to themselves, the standards guys took
		// the V1 and 2 postal code standards (code padded to 11 characters with
		// spaces) and replaced the spaces with zeros if a) the country is "USA"
		// and b) if the trailing "+4" portion of the postal code is unknown.
		// Quite what happens to pad Canadian postal codes (they are always 6
		// alphanumeric characters, like British postal codes) is undocumented.
		//
		// We will extract the 5-digit zip and the +4 section.  If the +4 is all
		// zeros we can discard it.  The last two digits are always useless (the
		// next version of the spec shortens the field to 9 characters) so we
		// can ignore them.

		// Naturally, some Texas licences ignore the spec and just use 5
		// characters if they don't have a +4 section.

		if len(licence.Postal()) > 5 {
			zip := license.Postal()[:5]
			plus4 := license.Postal()[5:9]

			if plus4 == "0000" || plus4 == "    " {
				license.SetPostal(zip)
			} else {
				license.SetPostal(zip + "+" + plus4)
			}
		}
	}

	// Now we can parse the birth date, too.
	if len(license.Country()) > 0 {
		license.SetDateOfBirth(parseDateV3(dateOfBirth, license.Country()))
		license.SetExpiryDate(parseDateV3(expiryDate, license.Country()))
		license.SetIssueDate(parseDateV3(issueDate, license.Country()))
	}

	return
}

func parseDateV3(data string, country string) time.Time {

	// And now we get the payoff for the previous awful decision to switch to
	// Lumpy Date Format: we're now supporting the international big-endian
	// date format used in Canada and V1 of the spec (yyyyMMdd) and the US
	// lumpy date format *in the same field*.  I can understand that different
	// versions of a standard don't agree with each other, but now we've got two
	// implementations of a standard within a single field in a single version
	// of the standard.  Breathtakingly stupid.

	var day int
	var month int
	var year int
	var err error
	var location *time.Location

	if country == "USA" {
		month, err = strconv.Atoi(data[:2])

		if err != nil {
			return time.Unix(0, 0)
		}

		day, err = strconv.Atoi(data[2:4])

		if err != nil {
			return time.Unix(0, 0)
		}

		year, err = strconv.Atoi(data[4:8])

		if err != nil {
			return time.Unix(0, 0)
		}
	} else {
		year, err = strconv.Atoi(data[:4])

		if err != nil {
			return time.Unix(0, 0)
		}

		month, err = strconv.Atoi(data[4:6])

		if err != nil {
			return time.Unix(0, 0)
		}

		day, err = strconv.Atoi(data[6:8])

		if err != nil {
			return time.Unix(0, 0)
		}
	}

	location, err = time.LoadLocation("UTC")

	return time.Date(year, time.Month(month), day, 0, 0, 0, 0, location)
}
<commit_msg>Fixed typo.<commit_after>package dlidparser

import (
	"errors"
	"strconv"
	"strings"
	"time"
)

func parseV3(data string, issuer string) (license *DLIDLicense, err error) {

	start, end, err := dataRangeV2(data)

	if end >= len(data) {
		err = errors.New("Payload location does not exist in data")
	}

	payload := data[start:end]

	if err != nil {
		return
	}

	license, err = parseDataV3(payload, issuer)

	if err != nil {
		return
	}

	return
}

func parseDataV3(licenceData string, issuer string) (license *DLIDLicense, err error) {

	// Version 3 of the DLID card spec was published in 2005.  It is currently
	// (as of 2012) used in Wisconsin.

	if !strings.HasPrefix(licenceData, "DL") {
		err = errors.New("Missing header in licence data chunk")
		return
	}

	licenceData = licenceData[2:]

	components := strings.Split(licenceData, "\n")

	license = new(DLIDLicense)

	license.SetIssuerId(issuer)
	license.SetIssuerName(issuers[issuer])

	var dateOfBirth string
	var expiryDate string
	var issueDate string

	for component := range components {

		if len(components[component]) < 3 {
			continue
		}

		identifier := components[component][0:3]
		data := components[component][3:]

		data = strings.Trim(data, " ")

		switch identifier {
		case "DCA":
			license.SetVehicleClass(data)

		case "DCB":
			license.SetRestrictionCodes(data)

		case "DCD":
			license.SetEndorsementCodes(data)

		case "DCS":
			license.SetLastName(data)

		case "DCG":
			license.SetCountry(data)

		case "DCT":

			// This field contains all of the licencee's names except last
			// name.  The V3 spec doc doesn't specify how the names are
			// separated and doesn't provide an example (unlike the 2000
			// doc).  Wisconsin use a space and Virginia use a comma.  This
			// is why standards have to be adequately documented.

			separator := " "

			if strings.Index(data, separator) == -1 {
				separator = ","
			}

			names := strings.Split(data, separator)

			license.SetFirstName(names[0])

			if len(names) > 1 {
				license.SetMiddleNames(names[1:])
			}

		case "DAG":
			license.SetStreet(data)

		case "DAI":
			license.SetCity(data)

		case "DAJ":
			license.SetState(data)

		case "DAK":
			license.SetPostal(data)

		case "DAQ":
			license.SetCustomerId(data)

		case "DBA":
			expiryDate = data

		case "DBB":
			dateOfBirth = data

		case "DBC":
			switch data {
			case "1":
				license.SetSex(DriverSexMale)
			case "2":
				license.SetSex(DriverSexFemale)
			default:
				license.SetSex(DriverSexNone)
			}

		case "DBD":
			issueDate = data
		}
	}

	// At this point we should know the country and the postal code (both are
	// mandatory fields) so we can undo the desperate mess the standards body
	// made of the postal code field.

	if license.Country() == "USA" && len(license.Postal()) > 0 {

		// For some reason known only to themselves, the standards guys took
		// the V1 and 2 postal code standards (code padded to 11 characters with
		// spaces) and replaced the spaces with zeros if a) the country is "USA"
		// and b) if the trailing "+4" portion of the postal code is unknown.
		// Quite what happens to pad Canadian postal codes (they are always 6
		// alphanumeric characters, like British postal codes) is undocumented.
		//
		// We will extract the 5-digit zip and the +4 section.  If the +4 is all
		// zeros we can discard it.  The last two digits are always useless (the
		// next version of the spec shortens the field to 9 characters) so we
		// can ignore them.

		// Naturally, some Texas licences ignore the spec and just use 5
		// characters if they don't have a +4 section.

		if len(license.Postal()) > 5 {
			zip := license.Postal()[:5]
			plus4 := license.Postal()[5:9]

			if plus4 == "0000" || plus4 == "    " {
				license.SetPostal(zip)
			} else {
				license.SetPostal(zip + "+" + plus4)
			}
		}
	}

	// Now we can parse the birth date, too.
	if len(license.Country()) > 0 {
		license.SetDateOfBirth(parseDateV3(dateOfBirth, license.Country()))
		license.SetExpiryDate(parseDateV3(expiryDate, license.Country()))
		license.SetIssueDate(parseDateV3(issueDate, license.Country()))
	}

	return
}

func parseDateV3(data string, country string) time.Time {

	// And now we get the payoff for the previous awful decision to switch to
	// Lumpy Date Format: we're now supporting the international big-endian
	// date format used in Canada and V1 of the spec (yyyyMMdd) and the US
	// lumpy date format *in the same field*.  I can understand that different
	// versions of a standard don't agree with each other, but now we've got two
	// implementations of a standard within a single field in a single version
	// of the standard.  Breathtakingly stupid.

	var day int
	var month int
	var year int
	var err error
	var location *time.Location

	if country == "USA" {
		month, err = strconv.Atoi(data[:2])

		if err != nil {
			return time.Unix(0, 0)
		}

		day, err = strconv.Atoi(data[2:4])

		if err != nil {
			return time.Unix(0, 0)
		}

		year, err = strconv.Atoi(data[4:8])

		if err != nil {
			return time.Unix(0, 0)
		}
	} else {
		year, err = strconv.Atoi(data[:4])

		if err != nil {
			return time.Unix(0, 0)
		}

		month, err = strconv.Atoi(data[4:6])

		if err != nil {
			return time.Unix(0, 0)
		}

		day, err = strconv.Atoi(data[6:8])

		if err != nil {
			return time.Unix(0, 0)
		}
	}

	location, err = time.LoadLocation("UTC")

	return time.Date(year, time.Month(month), day, 0, 0, 0, 0, location)
}
<|endoftext|>


<commit_before>package ssh

import (
	"bufio"
	"fmt"
	"io/ioutil"
	"os"
	"strings"

	"golang.org/x/crypto/ssh"
	"golang.org/x/crypto/ssh/testdata"

	. "gopkg.in/check.v1"
)

type (
	SuiteCommon struct{}

	mockKnownHosts struct{}
)

func (mockKnownHosts) host() string { return "github.com" }
func (mockKnownHosts) knownHosts() []byte {
	return []byte(`github.com ssh-rsa AAAAB3NzaC1yc2EAAAABIwAAAQEAq2A7hRGmdnm9tUDbO9IDSwBK6TbQa+PXYPCPy6rbTrTtw7PHkccKrpp0yVhp5HdEIcKr6pLlVDBfOLX9QUsyCOV0wzfjIJNlGEYsdlLJizHhbn2mUjvSAHQqZETYP81eFzLQNnPHt4EVVUh7VfDESU84KezmD5QlWpXLmvU31/yMf+Se8xhHTvKSCZIFImWwoG6mbUoWf9nzpIoaSjB+weqqUUmpaaasXVal72J+UX2B+2RPW3RcT0eOzQgqlJL3RKrTJvdsjE3JEAvGq3lGHSZXy28G3skua2SmVi/w4yCE6gbODqnTWlg7+wC604ydGXA8VJiS5ap43JXiUFFAaQ==`)
}
func (mockKnownHosts) Network() string { return "tcp" }
func (mockKnownHosts) String() string  { return "github.com:22" }

var _ = Suite(&SuiteCommon{})

func (s *SuiteCommon) TestKeyboardInteractiveName(c *C) {
	a := &KeyboardInteractive{
		User:      "test",
		Challenge: nil,
	}
	c.Assert(a.Name(), Equals, KeyboardInteractiveName)
}

func (s *SuiteCommon) TestKeyboardInteractiveString(c *C) {
	a := &KeyboardInteractive{
		User:      "test",
		Challenge: nil,
	}
	c.Assert(a.String(), Equals, fmt.Sprintf("user: test, name: %s", KeyboardInteractiveName))
}

func (s *SuiteCommon) TestPasswordName(c *C) {
	a := &Password{
		User:     "test",
		Password: "",
	}
	c.Assert(a.Name(), Equals, PasswordName)
}

func (s *SuiteCommon) TestPasswordString(c *C) {
	a := &Password{
		User:     "test",
		Password: "",
	}
	c.Assert(a.String(), Equals, fmt.Sprintf("user: test, name: %s", PasswordName))
}

func (s *SuiteCommon) TestPasswordCallbackName(c *C) {
	a := &PasswordCallback{
		User:     "test",
		Callback: nil,
	}
	c.Assert(a.Name(), Equals, PasswordCallbackName)
}

func (s *SuiteCommon) TestPasswordCallbackString(c *C) {
	a := &PasswordCallback{
		User:     "test",
		Callback: nil,
	}
	c.Assert(a.String(), Equals, fmt.Sprintf("user: test, name: %s", PasswordCallbackName))
}

func (s *SuiteCommon) TestPublicKeysName(c *C) {
	a := &PublicKeys{
		User:   "test",
		Signer: nil,
	}
	c.Assert(a.Name(), Equals, PublicKeysName)
}

func (s *SuiteCommon) TestPublicKeysString(c *C) {
	a := &PublicKeys{
		User:   "test",
		Signer: nil,
	}
	c.Assert(a.String(), Equals, fmt.Sprintf("user: test, name: %s", PublicKeysName))
}

func (s *SuiteCommon) TestPublicKeysCallbackName(c *C) {
	a := &PublicKeysCallback{
		User:     "test",
		Callback: nil,
	}
	c.Assert(a.Name(), Equals, PublicKeysCallbackName)
}

func (s *SuiteCommon) TestPublicKeysCallbackString(c *C) {
	a := &PublicKeysCallback{
		User:     "test",
		Callback: nil,
	}
	c.Assert(a.String(), Equals, fmt.Sprintf("user: test, name: %s", PublicKeysCallbackName))
}
func (s *SuiteCommon) TestNewSSHAgentAuth(c *C) {
	if os.Getenv("SSH_AUTH_SOCK") == "" {
		c.Skip("SSH_AUTH_SOCK or SSH_TEST_PRIVATE_KEY are required")
	}

	auth, err := NewSSHAgentAuth("foo")
	c.Assert(err, IsNil)
	c.Assert(auth, NotNil)
}

func (s *SuiteCommon) TestNewSSHAgentAuthNoAgent(c *C) {
	addr := os.Getenv("SSH_AUTH_SOCK")
	err := os.Unsetenv("SSH_AUTH_SOCK")
	c.Assert(err, IsNil)

	defer func() {
		err := os.Setenv("SSH_AUTH_SOCK", addr)
		c.Assert(err, IsNil)
	}()

	k, err := NewSSHAgentAuth("foo")
	c.Assert(k, IsNil)
	c.Assert(err, ErrorMatches, ".*SSH_AUTH_SOCK.*|.*SSH agent .* not running.*")
}

func (*SuiteCommon) TestNewPublicKeys(c *C) {
	auth, err := NewPublicKeys("foo", testdata.PEMBytes["rsa"], "")
	c.Assert(err, IsNil)
	c.Assert(auth, NotNil)
}

func (*SuiteCommon) TestNewPublicKeysWithEncryptedPEM(c *C) {
	f := testdata.PEMEncryptedKeys[0]
	auth, err := NewPublicKeys("foo", f.PEMBytes, f.EncryptionKey)
	c.Assert(err, IsNil)
	c.Assert(auth, NotNil)
}

func (*SuiteCommon) TestNewPublicKeysFromFile(c *C) {
	f, err := ioutil.TempFile("", "ssh-test")
	c.Assert(err, IsNil)
	_, err = f.Write(testdata.PEMBytes["rsa"])
	c.Assert(err, IsNil)
	c.Assert(f.Close(), IsNil)
	defer os.RemoveAll(f.Name())

	auth, err := NewPublicKeysFromFile("foo", f.Name(), "")
	c.Assert(err, IsNil)
	c.Assert(auth, NotNil)
}

func (*SuiteCommon) TestNewPublicKeysWithInvalidPEM(c *C) {
	auth, err := NewPublicKeys("foo", []byte("bar"), "")
	c.Assert(err, NotNil)
	c.Assert(auth, IsNil)
}

func (*SuiteCommon) TestNewKnownHostsCallback(c *C) {
	var mock = mockKnownHosts{}

	f, err := ioutil.TempFile("", "known-hosts")
	c.Assert(err, IsNil)

	_, err = f.Write(mock.knownHosts())
	c.Assert(err, IsNil)

	err = f.Close()
	c.Assert(err, IsNil)

	defer os.RemoveAll(f.Name())

	f, err = os.Open(f.Name())
	c.Assert(err, IsNil)

	defer f.Close()

	var hostKey ssh.PublicKey
	scanner := bufio.NewScanner(f)
	for scanner.Scan() {
		fields := strings.Split(scanner.Text(), " ")
		if len(fields) != 3 {
			continue
		}
		if strings.Contains(fields[0], mock.host()) {
			var err error
			hostKey, _, _, _, err = ssh.ParseAuthorizedKey(scanner.Bytes())
			if err != nil {
				c.Fatalf("error parsing %q: %v", fields[2], err)
			}
			break
		}
	}
	if hostKey == nil {
		c.Fatalf("no hostkey for %s", mock.host())
	}

	clb, err := NewKnownHostsCallback(f.Name())
	c.Assert(err, IsNil)

	err = clb(mock.String(), mock, hostKey)
	c.Assert(err, IsNil)
}
<commit_msg>plumbing: transport/ssh, fix no agent test on windows<commit_after>package ssh

import (
	"bufio"
	"fmt"
	"io/ioutil"
	"os"
	"strings"

	"golang.org/x/crypto/ssh"
	"golang.org/x/crypto/ssh/testdata"

	. "gopkg.in/check.v1"
)

type (
	SuiteCommon struct{}

	mockKnownHosts struct{}
)

func (mockKnownHosts) host() string { return "github.com" }
func (mockKnownHosts) knownHosts() []byte {
	return []byte(`github.com ssh-rsa AAAAB3NzaC1yc2EAAAABIwAAAQEAq2A7hRGmdnm9tUDbO9IDSwBK6TbQa+PXYPCPy6rbTrTtw7PHkccKrpp0yVhp5HdEIcKr6pLlVDBfOLX9QUsyCOV0wzfjIJNlGEYsdlLJizHhbn2mUjvSAHQqZETYP81eFzLQNnPHt4EVVUh7VfDESU84KezmD5QlWpXLmvU31/yMf+Se8xhHTvKSCZIFImWwoG6mbUoWf9nzpIoaSjB+weqqUUmpaaasXVal72J+UX2B+2RPW3RcT0eOzQgqlJL3RKrTJvdsjE3JEAvGq3lGHSZXy28G3skua2SmVi/w4yCE6gbODqnTWlg7+wC604ydGXA8VJiS5ap43JXiUFFAaQ==`)
}
func (mockKnownHosts) Network() string { return "tcp" }
func (mockKnownHosts) String() string  { return "github.com:22" }

var _ = Suite(&SuiteCommon{})

func (s *SuiteCommon) TestKeyboardInteractiveName(c *C) {
	a := &KeyboardInteractive{
		User:      "test",
		Challenge: nil,
	}
	c.Assert(a.Name(), Equals, KeyboardInteractiveName)
}

func (s *SuiteCommon) TestKeyboardInteractiveString(c *C) {
	a := &KeyboardInteractive{
		User:      "test",
		Challenge: nil,
	}
	c.Assert(a.String(), Equals, fmt.Sprintf("user: test, name: %s", KeyboardInteractiveName))
}

func (s *SuiteCommon) TestPasswordName(c *C) {
	a := &Password{
		User:     "test",
		Password: "",
	}
	c.Assert(a.Name(), Equals, PasswordName)
}

func (s *SuiteCommon) TestPasswordString(c *C) {
	a := &Password{
		User:     "test",
		Password: "",
	}
	c.Assert(a.String(), Equals, fmt.Sprintf("user: test, name: %s", PasswordName))
}

func (s *SuiteCommon) TestPasswordCallbackName(c *C) {
	a := &PasswordCallback{
		User:     "test",
		Callback: nil,
	}
	c.Assert(a.Name(), Equals, PasswordCallbackName)
}

func (s *SuiteCommon) TestPasswordCallbackString(c *C) {
	a := &PasswordCallback{
		User:     "test",
		Callback: nil,
	}
	c.Assert(a.String(), Equals, fmt.Sprintf("user: test, name: %s", PasswordCallbackName))
}

func (s *SuiteCommon) TestPublicKeysName(c *C) {
	a := &PublicKeys{
		User:   "test",
		Signer: nil,
	}
	c.Assert(a.Name(), Equals, PublicKeysName)
}

func (s *SuiteCommon) TestPublicKeysString(c *C) {
	a := &PublicKeys{
		User:   "test",
		Signer: nil,
	}
	c.Assert(a.String(), Equals, fmt.Sprintf("user: test, name: %s", PublicKeysName))
}

func (s *SuiteCommon) TestPublicKeysCallbackName(c *C) {
	a := &PublicKeysCallback{
		User:     "test",
		Callback: nil,
	}
	c.Assert(a.Name(), Equals, PublicKeysCallbackName)
}

func (s *SuiteCommon) TestPublicKeysCallbackString(c *C) {
	a := &PublicKeysCallback{
		User:     "test",
		Callback: nil,
	}
	c.Assert(a.String(), Equals, fmt.Sprintf("user: test, name: %s", PublicKeysCallbackName))
}
func (s *SuiteCommon) TestNewSSHAgentAuth(c *C) {
	if os.Getenv("SSH_AUTH_SOCK") == "" {
		c.Skip("SSH_AUTH_SOCK or SSH_TEST_PRIVATE_KEY are required")
	}

	auth, err := NewSSHAgentAuth("foo")
	c.Assert(err, IsNil)
	c.Assert(auth, NotNil)
}

func (s *SuiteCommon) TestNewSSHAgentAuthNoAgent(c *C) {
	addr := os.Getenv("SSH_AUTH_SOCK")
	err := os.Unsetenv("SSH_AUTH_SOCK")
	c.Assert(err, IsNil)

	defer func() {
		err := os.Setenv("SSH_AUTH_SOCK", addr)
		c.Assert(err, IsNil)
	}()

	k, err := NewSSHAgentAuth("foo")
	c.Assert(k, IsNil)
	c.Assert(err, ErrorMatches, ".*SSH_AUTH_SOCK.*|.*SSH agent .* not detect.*")
}

func (*SuiteCommon) TestNewPublicKeys(c *C) {
	auth, err := NewPublicKeys("foo", testdata.PEMBytes["rsa"], "")
	c.Assert(err, IsNil)
	c.Assert(auth, NotNil)
}

func (*SuiteCommon) TestNewPublicKeysWithEncryptedPEM(c *C) {
	f := testdata.PEMEncryptedKeys[0]
	auth, err := NewPublicKeys("foo", f.PEMBytes, f.EncryptionKey)
	c.Assert(err, IsNil)
	c.Assert(auth, NotNil)
}

func (*SuiteCommon) TestNewPublicKeysFromFile(c *C) {
	f, err := ioutil.TempFile("", "ssh-test")
	c.Assert(err, IsNil)
	_, err = f.Write(testdata.PEMBytes["rsa"])
	c.Assert(err, IsNil)
	c.Assert(f.Close(), IsNil)
	defer os.RemoveAll(f.Name())

	auth, err := NewPublicKeysFromFile("foo", f.Name(), "")
	c.Assert(err, IsNil)
	c.Assert(auth, NotNil)
}

func (*SuiteCommon) TestNewPublicKeysWithInvalidPEM(c *C) {
	auth, err := NewPublicKeys("foo", []byte("bar"), "")
	c.Assert(err, NotNil)
	c.Assert(auth, IsNil)
}

func (*SuiteCommon) TestNewKnownHostsCallback(c *C) {
	var mock = mockKnownHosts{}

	f, err := ioutil.TempFile("", "known-hosts")
	c.Assert(err, IsNil)

	_, err = f.Write(mock.knownHosts())
	c.Assert(err, IsNil)

	err = f.Close()
	c.Assert(err, IsNil)

	defer os.RemoveAll(f.Name())

	f, err = os.Open(f.Name())
	c.Assert(err, IsNil)

	defer f.Close()

	var hostKey ssh.PublicKey
	scanner := bufio.NewScanner(f)
	for scanner.Scan() {
		fields := strings.Split(scanner.Text(), " ")
		if len(fields) != 3 {
			continue
		}
		if strings.Contains(fields[0], mock.host()) {
			var err error
			hostKey, _, _, _, err = ssh.ParseAuthorizedKey(scanner.Bytes())
			if err != nil {
				c.Fatalf("error parsing %q: %v", fields[2], err)
			}
			break
		}
	}
	if hostKey == nil {
		c.Fatalf("no hostkey for %s", mock.host())
	}

	clb, err := NewKnownHostsCallback(f.Name())
	c.Assert(err, IsNil)

	err = clb(mock.String(), mock, hostKey)
	c.Assert(err, IsNil)
}
<|endoftext|>


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

package quickfix

import (
	"bytes"
	"fmt"
	"go/build"
	"go/token"
	"path/filepath"
	"regexp"
	"strconv"
	"strings"

	"nvim-go/config"
	"nvim-go/context"
	"nvim-go/pathutil"

	"github.com/neovim-go/vim"
)

// SetLoclist set the error results data to current buffer's locationlist.
func SetLoclist(v *vim.Vim, loclist []*vim.QuickfixError) error {
	// setloclist({nr}, {list} [, {action}])
	// v.Call(fname string, result interface{}, args ...interface{})
	if len(loclist) > 0 {
		v.Call("setloclist", nil, 0, loclist)
	} else {
		v.Command("lexpr ''")
	}

	return nil
}

var (
	listtype                  ErrorListType
	openlistCmd, closelistCmd func() error
	setlistCmd, clearlistCmd  func() error
)

type ErrorListType string

const (
	Quickfix     ErrorListType = "quickfix"
	LocationList ErrorListType = "locationlist"
)

// SetErrorlist set the error results data to current buffer's locationlist.
func SetErrorlist(v *vim.Vim, errlist []*vim.QuickfixError) error {
	if setlistCmd == nil || clearlistCmd == nil {
		switch listtype {
		case Quickfix:
			setlistCmd = func() error { return v.Call("setqflist", nil, errlist) }
			clearlistCmd = func() error { return v.Command("cgetexpr ''") }
		case LocationList:
			setlistCmd = func() error { return v.Call("setloclist", nil, 0, errlist) }
			clearlistCmd = func() error { return v.Command("lgetexpr ''") }
		}
	}

	if len(errlist) == 0 {
		return clearlistCmd()
	}

	return setlistCmd()
}

// ErrorList merges the errlist map items and open the locationlist window.
// TODO(zchee): This function will reports the errors with open the quickfix window, but will close
// the quickfix window if no errors.
// Do ErrorList function name is appropriate?
func ErrorList(v *vim.Vim, w vim.Window, errors map[string][]*vim.QuickfixError, keep bool) error {
	if listtype == "" {
		listtype = ErrorListType(config.ErrorListType)
		switch listtype {
		case Quickfix:
			openlistCmd = func() error { return v.Command("copen") }
			closelistCmd = func() error { return v.Command("cclose") }
		case LocationList:
			openlistCmd = func() error { return v.Command("lopen") }
			closelistCmd = func() error { return v.Command("lclose") }
		}
	}

	if errors == nil || len(errors) == 0 {
		return closelistCmd()
	}

	var errlist []*vim.QuickfixError
	for _, err := range errors {
		errlist = append(errlist, err...)
	}

	if err := SetErrorlist(v, errlist); err != nil {
		return err
	}

	if keep {
		defer v.SetCurrentWindow(w)
	}
	return openlistCmd()
}

// OpenLoclist open or close the current buffer's locationlist window.
func OpenLoclist(v *vim.Vim, w vim.Window, loclist []*vim.QuickfixError, keep bool) error {
	if len(loclist) == 0 {
		return v.Command("lclose")
	}

	v.Command("lopen")
	if keep {
		return v.SetCurrentWindow(w)
	}
	return nil
}

// CloseLoclist close the current buffer's locationlist window.
func CloseLoclist(v *vim.Vim) error {
	return v.Command("lclose")
}

// SetQuickfix set the error results data to quickfix list.
func SetQuickfix(p *vim.Pipeline, qflist []*vim.QuickfixError) error {
	p.Call("setqflist", nil, qflist)

	return nil
}

// OpenOuickfix open the quickfix list window.
func OpenOuickfix(p *vim.Pipeline, w vim.Window, keep bool) error {
	p.Command("copen")
	if keep {
		p.SetCurrentWindow(w)
	}
	if err := p.Wait(); err != nil {
		return err
	}

	return nil
}

// CloseQuickfix close the quickfix list window.
func CloseQuickfix(v *vim.Vim) error {
	return v.Command("cclose")
}

func GotoPos(v *vim.Vim, w vim.Window, pos token.Position, cwd string) error {
	fname, line, col := SplitPos(pos.String(), cwd)

	v.Command(fmt.Sprintf("edit %s", pathutil.Expand(fname)))
	v.SetWindowCursor(w, [2]int{line, col - 1})
	defer v.Command(`lclose`)

	return v.Command(`normal zz`)
}

// SplitPos parses a string of form 'token.Pos', and return the relative
// filepath from the current working directory path.
func SplitPos(pos string, cwd string) (string, int, int) {
	position := strings.Split(pos, ":")

	fname := position[0]
	line, err := strconv.ParseInt(position[1], 10, 64)
	if err != nil {
		line = 0
	}
	col, err := strconv.ParseInt(position[2], 10, 64)
	if err != nil {
		col = 0
	}

	if strings.HasPrefix(fname, cwd) {
		frel := strings.TrimPrefix(fname, cwd+string(filepath.Separator))
		if fname != frel {
			return frel, int(line), int(col)
		}
	}

	return fname, int(line), int(col)
}

var errRe = regexp.MustCompile(`(?m)^([^:]+):(\d+)(?::(\d+))?:\s(.*)`)

// ParseError parses a typical error message of Go compile tools.
// TODO(zchee): More elegant way
func ParseError(errors []byte, cwd string, ctxt *context.BuildContext) ([]*vim.QuickfixError, error) {
	var (
		parentDir string
		fpath     string
		errlist   []*vim.QuickfixError
	)

	// m[1]: relative file path of error file
	// m[2]: line number of error point
	// m[3]: column number of error point
	// m[4]: error description text
	for _, m := range errRe.FindAllSubmatch(errors, -1) {
		fpath = filepath.Base(string(bytes.Replace(m[1], []byte{'\t'}, nil, -1)))
		if bytes.Contains(m[1], []byte{'#'}) {
			// m[1][2:] is trim '# ' from errors directory
			// '# nvim-go/nvim' -> 'nvim-go/nvim'
			path := bytes.Split(m[1][2:], []byte{'\n'})

			// save the parent directory path for the second and subsequent error description
			parentDir = string(path[0])
			fpath = filepath.Base(string(bytes.Replace(path[1], []byte{'\t'}, nil, -1)))
		}

		filename := filepath.Join(parentDir, fpath)

		switch ctxt.Tool {
		case "go":
			goSrcPath := filepath.Join(build.Default.GOPATH, "src")
			currentDir := strings.TrimPrefix(cwd, goSrcPath+string(filepath.Separator))
			filename = strings.TrimPrefix(filepath.Clean(filename), currentDir+string(filepath.Separator))

		case "gb":
			if !filepath.IsAbs(filename) {
				filename = filepath.Join(ctxt.GbProjectDir, "src", filename)
			}
			if frel, err := filepath.Rel(cwd, filename); err == nil {
				filename = frel
			}
		}

		line, err := strconv.Atoi(string(m[2]))
		if err != nil {
			return nil, err
		}

		col, err := strconv.Atoi(string(m[3]))
		// fallback if cannot convert col to type int
		// Basically, col == ""
		if err != nil {
			col = 0
		}

		errlist = append(errlist, &vim.QuickfixError{
			FileName: filename,
			LNum:     line,
			Col:      col,
			Text:     string(bytes.TrimSpace(m[4])),
		})
	}

	return errlist, nil
}
<commit_msg>nvim/quickfix: Fix display the old error when Fmt after Bulid<commit_after>// Copyright 2016 Koichi Shiraishi. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package quickfix

import (
	"bytes"
	"fmt"
	"go/build"
	"go/token"
	"path/filepath"
	"regexp"
	"strconv"
	"strings"

	"nvim-go/config"
	"nvim-go/context"
	"nvim-go/pathutil"

	"github.com/neovim-go/vim"
)

// SetLoclist set the error results data to current buffer's locationlist.
func SetLoclist(v *vim.Vim, loclist []*vim.QuickfixError) error {
	// setloclist({nr}, {list} [, {action}])
	// v.Call(fname string, result interface{}, args ...interface{})
	if len(loclist) > 0 {
		v.Call("setloclist", nil, 0, loclist)
	} else {
		v.Command("lexpr ''")
	}

	return nil
}

var (
	listtype     ErrorListType
	openlistCmd  func() error
	closelistCmd func() error
	clearlistCmd func() error
	setlistCmd   func(errlist []*vim.QuickfixError) error
)

type ErrorListType string

const (
	Quickfix     ErrorListType = "quickfix"
	LocationList ErrorListType = "locationlist"
)

// ErrorList merges the errlist map items and open the locationlist window.
// TODO(zchee): This function will reports the errors with open the quickfix window, but will close
// the quickfix window if no errors.
// Do ErrorList function name is appropriate?
func ErrorList(v *vim.Vim, errors map[string][]*vim.QuickfixError, keep bool) error {
	if listtype == "" {
		listtype = ErrorListType(config.ErrorListType)
		switch listtype {
		case Quickfix:
			openlistCmd = func() error { return v.Command("copen") }
			closelistCmd = func() error { return v.Command("cclose") }
			clearlistCmd = func() error { return v.Command("cgetexpr ''") }
		case LocationList:
			openlistCmd = func() error { return v.Command("lopen") }
			closelistCmd = func() error { return v.Command("lclose") }
			clearlistCmd = func() error { return v.Command("lgetexpr ''") }
		}
	}

	if errors == nil || len(errors) == 0 {
		defer clearlistCmd()
		return closelistCmd()
	}

	var errlist []*vim.QuickfixError
	for _, err := range errors {
		errlist = append(errlist, err...)
	}
	if err := SetErrorlist(v, errlist); err != nil {
		return err
	}

	if keep {
		w, err := v.CurrentWindow()
		if err != nil {
			return err
		}
		defer v.SetCurrentWindow(w)
	}
	return openlistCmd()
}

// SetErrorlist set the error results data to Neovim error list.
func SetErrorlist(v *vim.Vim, errlist []*vim.QuickfixError) error {
	if setlistCmd == nil {
		switch listtype {
		case Quickfix:
			setlistCmd = func(errlist []*vim.QuickfixError) error { return v.Call("setqflist", nil, errlist, "r") }
		case LocationList:
			setlistCmd = func(errlist []*vim.QuickfixError) error { return v.Call("setloclist", nil, 0, errlist, "r") }
		}
	}

	return setlistCmd(errlist)
}

// ClearErrorlist clear the Neovim error list.
func ClearErrorlist(v *vim.Vim) error {
	if clearlistCmd == nil {
		switch listtype {
		case Quickfix:
			clearlistCmd = func() error { return v.Command("cgetexpr ''") }
		case LocationList:
			clearlistCmd = func() error { return v.Command("lgetexpr ''") }
		}
	}

	return clearlistCmd()
}

// OpenLoclist open or close the current buffer's locationlist window.
func OpenLoclist(v *vim.Vim, w vim.Window, loclist []*vim.QuickfixError, keep bool) error {
	if len(loclist) == 0 {
		return v.Command("lclose")
	}

	v.Command("lopen")
	if keep {
		return v.SetCurrentWindow(w)
	}
	return nil
}

// CloseLoclist close the current buffer's locationlist window.
func CloseLoclist(v *vim.Vim) error {
	return v.Command("lclose")
}

// SetQuickfix set the error results data to quickfix list.
func SetQuickfix(p *vim.Pipeline, qflist []*vim.QuickfixError) error {
	p.Call("setqflist", nil, qflist)

	return nil
}

// OpenOuickfix open the quickfix list window.
func OpenOuickfix(p *vim.Pipeline, w vim.Window, keep bool) error {
	p.Command("copen")
	if keep {
		p.SetCurrentWindow(w)
	}
	if err := p.Wait(); err != nil {
		return err
	}

	return nil
}

// CloseQuickfix close the quickfix list window.
func CloseQuickfix(v *vim.Vim) error {
	return v.Command("cclose")
}

func GotoPos(v *vim.Vim, w vim.Window, pos token.Position, cwd string) error {
	fname, line, col := SplitPos(pos.String(), cwd)

	v.Command(fmt.Sprintf("edit %s", pathutil.Expand(fname)))
	v.SetWindowCursor(w, [2]int{line, col - 1})
	defer v.Command(`lclose`)

	return v.Command(`normal zz`)
}

// SplitPos parses a string of form 'token.Pos', and return the relative
// filepath from the current working directory path.
func SplitPos(pos string, cwd string) (string, int, int) {
	position := strings.Split(pos, ":")

	fname := position[0]
	line, err := strconv.ParseInt(position[1], 10, 64)
	if err != nil {
		line = 0
	}
	col, err := strconv.ParseInt(position[2], 10, 64)
	if err != nil {
		col = 0
	}

	if strings.HasPrefix(fname, cwd) {
		frel := strings.TrimPrefix(fname, cwd+string(filepath.Separator))
		if fname != frel {
			return frel, int(line), int(col)
		}
	}

	return fname, int(line), int(col)
}

var errRe = regexp.MustCompile(`(?m)^([^:]+):(\d+)(?::(\d+))?:\s(.*)`)

// ParseError parses a typical error message of Go compile tools.
// TODO(zchee): More elegant way
func ParseError(errors []byte, cwd string, ctxt *context.BuildContext) ([]*vim.QuickfixError, error) {
	var (
		parentDir string
		fpath     string
		errlist   []*vim.QuickfixError
	)

	// m[1]: relative file path of error file
	// m[2]: line number of error point
	// m[3]: column number of error point
	// m[4]: error description text
	for _, m := range errRe.FindAllSubmatch(errors, -1) {
		fpath = filepath.Base(string(bytes.Replace(m[1], []byte{'\t'}, nil, -1)))
		if bytes.Contains(m[1], []byte{'#'}) {
			// m[1][2:] is trim '# ' from errors directory
			// '# nvim-go/nvim' -> 'nvim-go/nvim'
			path := bytes.Split(m[1][2:], []byte{'\n'})

			// save the parent directory path for the second and subsequent error description
			parentDir = string(path[0])
			fpath = filepath.Base(string(bytes.Replace(path[1], []byte{'\t'}, nil, -1)))
		}

		filename := filepath.Join(parentDir, fpath)

		switch ctxt.Tool {
		case "go":
			goSrcPath := filepath.Join(build.Default.GOPATH, "src")
			currentDir := strings.TrimPrefix(cwd, goSrcPath+string(filepath.Separator))
			filename = strings.TrimPrefix(filepath.Clean(filename), currentDir+string(filepath.Separator))

		case "gb":
			if !filepath.IsAbs(filename) {
				filename = filepath.Join(ctxt.GbProjectDir, "src", filename)
			}
			if frel, err := filepath.Rel(cwd, filename); err == nil {
				filename = frel
			}
		}

		line, err := strconv.Atoi(string(m[2]))
		if err != nil {
			return nil, err
		}

		col, err := strconv.Atoi(string(m[3]))
		// fallback if cannot convert col to type int
		// Basically, col == ""
		if err != nil {
			col = 0
		}

		errlist = append(errlist, &vim.QuickfixError{
			FileName: filename,
			LNum:     line,
			Col:      col,
			Text:     string(bytes.TrimSpace(m[4])),
		})
	}

	return errlist, nil
}
<|endoftext|>


<commit_before>// (C) Copyright 2017 Hewlett Packard Enterprise Development LP
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package main

import (
	"flag"
	"fmt"
	"net/http"
	"os"
	"sync"
	"time"

	"github.com/HewlettPackard/lustre_exporter/sources"
	"github.com/prometheus/client_golang/prometheus"
	"github.com/prometheus/client_golang/prometheus/promhttp"
	"github.com/prometheus/common/log"
	"github.com/prometheus/common/version"
)

var (
	scrapeDurations = prometheus.NewSummaryVec(
		prometheus.SummaryOpts{
			Namespace: sources.Namespace,
			Subsystem: "exporter",
			Name:      "scrape_duration_seconds",
			Help:      "lustre_exporter: Duration of a scrape job.",
		},
		[]string{"source", "result"},
	)
)

//LustreSource is a list of all sources that the user would like to collect.
type LustreSource struct {
	sourceList map[string]sources.LustreSource
}

//Describe implements the prometheus.Describe interface
func (l LustreSource) Describe(ch chan<- *prometheus.Desc) {
	scrapeDurations.Describe(ch)
}

//Collect implements the prometheus.Collect interface
func (l LustreSource) Collect(ch chan<- prometheus.Metric) {
	wg := sync.WaitGroup{}
	wg.Add(len(l.sourceList))
	for name, c := range l.sourceList {
		go func(name string, c sources.LustreSource) {
			collectFromSource(name, c, ch)
			wg.Done()
		}(name, c)
	}
	wg.Wait()
	scrapeDurations.Collect(ch)
}

func collectFromSource(name string, s sources.LustreSource, ch chan<- prometheus.Metric) {
	result := "success"
	begin := time.Now()
	err := s.Update(ch)
	duration := time.Since(begin)
	if err != nil {
		log.Errorf("ERROR: %q source failed after %f seconds: %s", name, duration.Seconds(), err)
		result = "error"
	} else {
		log.Debugf("OK: %q source succeeded after %f seconds: %s", name, duration.Seconds(), err)
	}
	scrapeDurations.WithLabelValues(name, result).Observe(duration.Seconds())
}

func loadSources(list []string) (map[string]sources.LustreSource, error) {
	sourceList := map[string]sources.LustreSource{}
	for _, name := range list {
		fn, ok := sources.Factories[name]
		if !ok {
			return nil, fmt.Errorf("source %q not available", name)
		}
		c, err := fn()
		if err != nil {
			return nil, err
		}
		sourceList[name] = c
	}
	return sourceList, nil
}

func init() {
	prometheus.MustRegister(version.NewCollector("lustre_exporter"))
}

func main() {
	var (
		showVersion   = flag.Bool("version", false, "Print version information.")
		listenAddress = flag.String("web.listen-address", ":9169", "Address to use to expose Lustre metrics.")
		metricsPath   = flag.String("web.telemetry-path", "/metrics", "Path to use to expose Lustre metrics.")
	)
	flag.Parse()

	if *showVersion {
		fmt.Fprintln(os.Stdout, version.Print("lustre_exporter"))
		os.Exit(0)
	}

	log.Infoln("Starting lustre_exporter", version.Info())
	log.Infoln("Build context", version.BuildContext())

	//expand to include more sources eventually (CLI, other?)
	enabledSources := []string{"procfs"}

	sourceList, err := loadSources(enabledSources)
	if err != nil {
		log.Fatalf("Couldn't load sources: %q", err)
	}

	log.Infof("Enabled sources:")
	for s := range sourceList {
		log.Infof(" - %s", s)
	}

	prometheus.MustRegister(LustreSource{sourceList: sourceList})
	handler := promhttp.HandlerFor(prometheus.DefaultGatherer, promhttp.HandlerOpts{ErrorLog: log.NewErrorLogger()})

	http.Handle(*metricsPath, prometheus.InstrumentHandler("prometheus", handler))
	http.HandleFunc("/", func(w http.ResponseWriter, r *http.Request) {
		w.Write([]byte(`
			<title>Lustre Exporter</title>

Lustre Exporter

<a href="` + *metricsPath + `">Metrics

			`))
	})

	log.Infoln("Listening on", *listenAddress)
	err = http.ListenAndServe(*listenAddress, nil)
	if err != nil {
		log.Fatal(err)
	}
}
<commit_msg>Enable procsys as a source<commit_after>// (C) Copyright 2017 Hewlett Packard Enterprise Development LP
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package main

import (
	"flag"
	"fmt"
	"net/http"
	"os"
	"sync"
	"time"

	"github.com/HewlettPackard/lustre_exporter/sources"
	"github.com/prometheus/client_golang/prometheus"
	"github.com/prometheus/client_golang/prometheus/promhttp"
	"github.com/prometheus/common/log"
	"github.com/prometheus/common/version"
)

var (
	scrapeDurations = prometheus.NewSummaryVec(
		prometheus.SummaryOpts{
			Namespace: sources.Namespace,
			Subsystem: "exporter",
			Name:      "scrape_duration_seconds",
			Help:      "lustre_exporter: Duration of a scrape job.",
		},
		[]string{"source", "result"},
	)
)

//LustreSource is a list of all sources that the user would like to collect.
type LustreSource struct {
	sourceList map[string]sources.LustreSource
}

//Describe implements the prometheus.Describe interface
func (l LustreSource) Describe(ch chan<- *prometheus.Desc) {
	scrapeDurations.Describe(ch)
}

//Collect implements the prometheus.Collect interface
func (l LustreSource) Collect(ch chan<- prometheus.Metric) {
	wg := sync.WaitGroup{}
	wg.Add(len(l.sourceList))
	for name, c := range l.sourceList {
		go func(name string, c sources.LustreSource) {
			collectFromSource(name, c, ch)
			wg.Done()
		}(name, c)
	}
	wg.Wait()
	scrapeDurations.Collect(ch)
}

func collectFromSource(name string, s sources.LustreSource, ch chan<- prometheus.Metric) {
	result := "success"
	begin := time.Now()
	err := s.Update(ch)
	duration := time.Since(begin)
	if err != nil {
		log.Errorf("ERROR: %q source failed after %f seconds: %s", name, duration.Seconds(), err)
		result = "error"
	} else {
		log.Debugf("OK: %q source succeeded after %f seconds: %s", name, duration.Seconds(), err)
	}
	scrapeDurations.WithLabelValues(name, result).Observe(duration.Seconds())
}

func loadSources(list []string) (map[string]sources.LustreSource, error) {
	sourceList := map[string]sources.LustreSource{}
	for _, name := range list {
		fn, ok := sources.Factories[name]
		if !ok {
			return nil, fmt.Errorf("source %q not available", name)
		}
		c, err := fn()
		if err != nil {
			return nil, err
		}
		sourceList[name] = c
	}
	return sourceList, nil
}

func init() {
	prometheus.MustRegister(version.NewCollector("lustre_exporter"))
}

func main() {
	var (
		showVersion   = flag.Bool("version", false, "Print version information.")
		listenAddress = flag.String("web.listen-address", ":9169", "Address to use to expose Lustre metrics.")
		metricsPath   = flag.String("web.telemetry-path", "/metrics", "Path to use to expose Lustre metrics.")
	)
	flag.Parse()

	if *showVersion {
		fmt.Fprintln(os.Stdout, version.Print("lustre_exporter"))
		os.Exit(0)
	}

	log.Infoln("Starting lustre_exporter", version.Info())
	log.Infoln("Build context", version.BuildContext())

	//expand to include more sources eventually (CLI, other?)
	enabledSources := []string{"procfs", "procsys"}

	sourceList, err := loadSources(enabledSources)
	if err != nil {
		log.Fatalf("Couldn't load sources: %q", err)
	}

	log.Infof("Enabled sources:")
	for s := range sourceList {
		log.Infof(" - %s", s)
	}

	prometheus.MustRegister(LustreSource{sourceList: sourceList})
	handler := promhttp.HandlerFor(prometheus.DefaultGatherer, promhttp.HandlerOpts{ErrorLog: log.NewErrorLogger()})

	http.Handle(*metricsPath, prometheus.InstrumentHandler("prometheus", handler))
	http.HandleFunc("/", func(w http.ResponseWriter, r *http.Request) {
		w.Write([]byte(`
			<title>Lustre Exporter</title>

Lustre Exporter

<a href="` + *metricsPath + `">Metrics

			`))
	})

	log.Infoln("Listening on", *listenAddress)
	err = http.ListenAndServe(*listenAddress, nil)
	if err != nil {
		log.Fatal(err)
	}
}
<|endoftext|>


<commit_before>package main

import (
	"os"
	"os/exec"
	"sync"
	"syscall"
	"github.com/go-clog/clog"
	"github.com/mitchellh/go-ps"
)

var (
	// 11 ttys from 0 to 12 but tty7 is excluded as it is frequently used for a X display
	ttys = [11]string{"tty1", "tty2", "tty3", "tty4", "tty5", "tty6", "tty8", "tty9", "tty10", "tty11", "tty12"}

	// GettysList of managed or unmanaged
	GettysList 		= make(map[int]*FollowGetty)
	GettysListMu	= sync.RWMutex{}
)

// FollowGetty struct with tracking infos
type FollowGetty struct {
	TTY			string
	PID			int
	Managed		bool
	Autologin	string
}

// ManageGettys to deal with reexec
func ManageGettys() {
	if MainConfig.StartedReexec {
		// Defaulting gettys as unmanaged
		GettysListMu.Lock()
		for i := range GettysList {
			GettysList[i].Managed = false
		}
		GettysListMu.Unlock()

		waitAndSpawnGettys()
	} else {
		manageAndSpawnGettys()
	}
}

// We are watching unmanaged gettys and spawn one if persistency
func aliveOrSpawn(getty *FollowGetty) {
	for {
		_, err := ps.FindProcess(getty.PID)
		if err != nil {
			clog.Error(2, "can't get running process for pid %d of unmanaged getty %s", getty.PID, getty.TTY)
			// since getty exited, we are now managing it
			getty.Managed = true
			return
		}
		// getty still running
	}
}

func waitAndSpawnGettys() {
	wg := sync.WaitGroup{}

	wg.Add(len(GettysList))
	for _, getty := range GettysList {
		go func(followGetty *FollowGetty) {
			defer wg.Done()
			for {
				if followGetty.Managed {
					if err := spawnGetty(followGetty.Autologin, followGetty.TTY, followGetty); err != nil {
						followGetty.PID = 0
						clog.Error(2, err.Error())
					}
				} else {
					aliveOrSpawn(followGetty)
				}

				// if no persistency or going shutdown, exit loop
				if !MainConfig.Persist || ShuttingDown {
					return
				}

			}
		}(getty)
	}
	wg.Wait()
}

// Gettys spawn the number of ttys required for len(autologins) to login.
// If persist is true, they'll be respawned if they die.
func manageAndSpawnGettys() {
	if MainConfig.StartedReexec {}
	autologins := MainConfig.Autologins

	switch len(autologins) {
	case 0, 1:
		// If there's no autologins, we still want to spawn a tty, and if
		// there's one, there's no need to coordinate goroutines
		var user string
		if len(autologins) == 1 {
			user = autologins[0]
		}

		for {
			GettysList[0] = &FollowGetty{TTY: ttys[0], Managed: true}
			if user != "" {
				GettysList[0].Autologin = user
			}
			if err := spawnGetty(user, ttys[0], GettysList[0]); err != nil {
				GettysList[0].PID = 0
				clog.Error(2, err.Error())
			}

			// if no persistency or going shutdown, exit loop
			if !MainConfig.Persist || ShuttingDown {
				return
			}
		}
	default:
		// getty(user, tty) blocks, so spawn a goroutine for each one and wait
		// for them to finish with a waitgroup, respawning as necessary in the
		// goroutine if it happens to quit. (NB: if persist is true they will
		// never finish.)
		wg := sync.WaitGroup{}

		// Sanity check length of autologins
		if len(autologins) > 11 {
			autologins = autologins[:11]
		}

		wg.Add(len(autologins))
		for i, user := range autologins {
			go func(user, tty string, idx int) {
				defer wg.Done()
				for {
					GettysList[idx] = &FollowGetty{TTY: tty, Managed: true, Autologin: user}
					clog.Trace("getty idx %d tty %s autologin %s", idx, tty, user)
					if err := spawnGetty(user, tty, GettysList[idx]); err != nil {
						GettysList[idx].PID = 0
						clog.Error(2, err.Error())
					}

					// if no persistency or going shutdown, exit loop
					if !MainConfig.Persist || ShuttingDown {
						return
					}
				}
			}(user, ttys[i], i)
		}
		// Block until all the ttys we spawned in goroutines are finished instead of
		// returning right away (and shutting down the system.)
		wg.Wait()
	}
}

// Spawn a single tty on tty, logging in with user autologin.
func spawnGetty(autologin, tty string, gettyFollow *FollowGetty) error {
	clog.Info("Spawning getty on %s with user %s", tty, autologin)

	var cmd *exec.Cmd
	if autologin != "" {
		cmd = exec.Command("getty", "--noclear", tty, "--autologin", autologin)
	} else {
		cmd = exec.Command("getty", "--noclear", tty)
	}

	// If we don't Setsid, we'll get an "inappropriate ioctl for device"
	// error upon starting the login shell.
	cmd.SysProcAttr = &syscall.SysProcAttr{
		Setsid: true,
	}
	cmd.Stdout = os.Stdout
	cmd.Stdin = os.Stdin
	cmd.Stderr = os.Stderr

	if err := cmd.Start(); err != nil {
		clog.Error(2, "[lutra] Getty %s exited with error: %s", tty, err.Error())
		return err
	}

	gettyFollow.PID = cmd.Process.Pid

	return cmd.Wait()
}
<commit_msg>This should help #4 but no<commit_after>package main

import (
	"os"
	"os/exec"
	"sync"
	"syscall"
	"github.com/go-clog/clog"
	"github.com/mitchellh/go-ps"
)

var (
	// 11 ttys from 0 to 12 but tty7 is excluded as it is frequently used for a X display
	ttys = [11]string{"tty1", "tty2", "tty3", "tty4", "tty5", "tty6", "tty8", "tty9", "tty10", "tty11", "tty12"}

	// GettysList of managed or unmanaged
	GettysList 		= make(map[int]*FollowGetty)
	GettysListMu	= sync.RWMutex{}
)

// FollowGetty struct with tracking infos
type FollowGetty struct {
	TTY			string
	PID			int
	Managed		bool
	Autologin	string
}

// ManageGettys to deal with reexec
func ManageGettys() {
	if MainConfig.StartedReexec {
		// Defaulting gettys as unmanaged
		GettysListMu.Lock()
		for i := range GettysList {
			GettysList[i].Managed = false
		}
		GettysListMu.Unlock()

		waitAndSpawnGettys()
	} else {
		manageAndSpawnGettys()
	}
}

// We are watching unmanaged gettys and spawn one if persistency
func aliveOrSpawn(getty *FollowGetty) {
	for {
		_, err := ps.FindProcess(getty.PID)
		if err != nil {
			clog.Error(2, "can't get running process for pid %d of unmanaged getty %s", getty.PID, getty.TTY)
			// since getty exited, we are now managing it
			getty.Managed = true
			return
		}
		// getty still running
	}
}

func waitAndSpawnGettys() {
	wg := sync.WaitGroup{}

	wg.Add(len(GettysList))
	for idx, getty := range GettysList {
		go func(followGetty *FollowGetty) {
			defer wg.Done()
			for {
				if followGetty.Managed {
					if err := spawnGetty(followGetty.Autologin, followGetty.TTY, idx); err != nil {
						followGetty.PID = 0
						clog.Error(2, err.Error())
					}
				} else {
					aliveOrSpawn(followGetty)
				}

				// if no persistency or going shutdown, exit loop
				if !MainConfig.Persist || ShuttingDown {
					return
				}

			}
		}(getty)
	}
	wg.Wait()
}

// Gettys spawn the number of ttys required for len(autologins) to login.
// If persist is true, they'll be respawned if they die.
func manageAndSpawnGettys() {
	if MainConfig.StartedReexec {}
	autologins := MainConfig.Autologins

	switch len(autologins) {
	case 0, 1:
		// If there's no autologins, we still want to spawn a tty, and if
		// there's one, there's no need to coordinate goroutines
		var user string
		if len(autologins) == 1 {
			user = autologins[0]
		}

		for {
			GettysList[0] = &FollowGetty{TTY: ttys[0], Managed: true}
			if user != "" {
				GettysList[0].Autologin = user
			}
			if err := spawnGetty(user, ttys[0], 0); err != nil {
				GettysList[0].PID = 0
				clog.Error(2, err.Error())
			}

			// if no persistency or going shutdown, exit loop
			if !MainConfig.Persist || ShuttingDown {
				return
			}
		}
	default:
		// getty(user, tty) blocks, so spawn a goroutine for each one and wait
		// for them to finish with a waitgroup, respawning as necessary in the
		// goroutine if it happens to quit. (NB: if persist is true they will
		// never finish.)
		wg := sync.WaitGroup{}

		// Sanity check length of autologins
		if len(autologins) > 11 {
			autologins = autologins[:11]
		}

		wg.Add(len(autologins))
		for i, user := range autologins {
			go func(user, tty string, idx int) {
				defer wg.Done()
				for {
					GettysListMu.Lock()
					GettysList[idx] = &FollowGetty{TTY: tty, Managed: true, Autologin: user}
					GettysListMu.Unlock()

					clog.Trace("getty idx %d tty %s autologin %s", idx, tty, user)
					if err := spawnGetty(user, tty, idx); err != nil {
						GettysListMu.Lock()
						GettysList[idx].PID = 0
						GettysListMu.Unlock()

						clog.Error(2, err.Error())
					}

					// if no persistency or going shutdown, exit loop
					if !MainConfig.Persist || ShuttingDown {
						return
					}
				}
			}(user, ttys[i], i)
		}
		// Block until all the ttys we spawned in goroutines are finished instead of
		// returning right away (and shutting down the system.)
		wg.Wait()
	}
}

// Spawn a single tty on tty, logging in with user autologin.
func spawnGetty(autologin, tty string, idx int) error {
	clog.Info("Spawning getty on %s with user %s", tty, autologin)

	var cmd *exec.Cmd
	if autologin != "" {
		cmd = exec.Command("getty", "--noclear", tty, "--autologin", autologin)
	} else {
		cmd = exec.Command("getty", "--noclear", tty)
	}

	// If we don't Setsid, we'll get an "inappropriate ioctl for device"
	// error upon starting the login shell.
	cmd.SysProcAttr = &syscall.SysProcAttr{
		Setsid: true,
	}
	cmd.Stdout = os.Stdout
	cmd.Stdin = os.Stdin
	cmd.Stderr = os.Stderr

	if err := cmd.Start(); err != nil {
		clog.Error(2, "[lutra] Getty %s exited with error: %s", tty, err.Error())
		return err
	}

	GettysList[idx].PID = cmd.Process.Pid

	return cmd.Wait()
}
<|endoftext|>


<commit_before>/*
 * portal - marshal
 *
 * a library that implements an algorithm for doing consumer coordination within Kafka, rather
 * than using Zookeeper or another external system.
 *
 */

package marshal

import (
	"crypto/md5"
	"encoding/binary"
	"errors"
	"fmt"
	"math/rand"
	"sync"
	"sync/atomic"
	"time"

	"github.com/dropbox/kafka"
)

// KafkaCluster is a user-agnostic view of the world. It connects to a Kafka cluster
// and runs rationalizers to observe the complete world state.
type KafkaCluster struct {
	// These members are not protected by the lock and can be read at any
	// time as they're write-once or only ever atomically updated. They must
	// never be overwritten once a KafkaCluster is created.
	quit       *int32
	name       string
	broker     *kafka.Broker
	producer   kafka.Producer
	partitions int
	jitters    chan time.Duration
	options    MarshalOptions

	// Lock protects the following members; you must have this lock in order to
	// read from or write to these.
	lock       sync.RWMutex
	marshalers []*Marshaler
	topics     map[string]int
	groups     map[string]map[string]*topicState
	// pausedGroups stores the expiry time for groups that are paused.
	pausedGroups map[string]time.Time

	// This WaitGroup is used for signalling when all of the rationalizers have
	// finished processing.
	rationalizers sync.WaitGroup

	// rsteps is updated whenever a rationalizer processes a log entry, this is
	// used mainly by the test suite.
	rsteps *int32

	// This is for testing only. When this is non-zero, the rationalizer will answer
	// queries based on THIS time instead of the current, actual time.
	ts int64
}

// MarshalOptions contains various tunables that can be used to adjust the configuration
// of the underlying system.
type MarshalOptions struct {
	// BrokerConnectionLimit is used to set the maximum simultaneous number of connections
	// that can be made to each broker.
	// Default: 30.
	BrokerConnectionLimit int

	// ConsumeRequestTimeout sets the time that we ask Kafka to wait before returning any
	// data to us. Setting this high uses more connections and can lead to some latency
	// but keeps the load on Kafka minimal. Use this to balance QPS against latency.
	// Default: 3 seconds.
	ConsumeRequestTimeout time.Duration

	// MarshalRequestTimeout is used for our coordination requests. This should be reasonable
	// at default, but is left as a tunable in case you have clients that are claiming an
	// extremely large number of partitions and are too slow.
	// Default: 1 second.
	MarshalRequestTimeout time.Duration
}

// NewMarshalOptions returns a set of MarshalOptions populated with defaults.
func NewMarshalOptions() MarshalOptions {
	return MarshalOptions{
		BrokerConnectionLimit: 30,
		ConsumeRequestTimeout: 3 * time.Second,
		MarshalRequestTimeout: 1 * time.Second,
	}
}

// Dial returns a new cluster object which can be used to instantiate a number of Marshalers
// that all use the same cluster. You may pass brokerConf or may set it to nil.
func Dial(name string, brokers []string, options MarshalOptions) (*KafkaCluster, error) {
	// Connect to Kafka
	brokerConf := kafka.NewBrokerConf("PortalMarshal")
	brokerConf.ConnectionLimit = options.BrokerConnectionLimit
	broker, err := kafka.Dial(brokers, brokerConf)
	if err != nil {
		return nil, err
	}

	c := &KafkaCluster{
		quit:         new(int32),
		rsteps:       new(int32),
		name:         name,
		options:      options,
		broker:       broker,
		producer:     broker.Producer(kafka.NewProducerConf()),
		topics:       make(map[string]int),
		groups:       make(map[string]map[string]*topicState),
		pausedGroups: make(map[string]time.Time),
		jitters:      make(chan time.Duration, 100),
	}

	// Do an initial metadata fetch, this will block a bit
	err = c.refreshMetadata()
	if err != nil {
		return nil, fmt.Errorf("Failed to get metadata: %s", err)
	}

	// If there is no marshal topic, then we can't run. The admins must go create the topic
	// before they can use this library. Please see the README.
	c.partitions = c.getTopicPartitions(MarshalTopic)
	if c.partitions == 0 {
		return nil, errors.New("Marshalling topic not found. Please see the documentation.")
	}

	// Now we start a goroutine to start consuming each of the partitions in the marshal
	// topic. Note that this doesn't handle increasing the partition count on that topic
	// without stopping all consumers.
	c.rationalizers.Add(c.partitions)
	for id := 0; id < c.partitions; id++ {
		go c.rationalize(id, c.kafkaConsumerChannel(id))
	}

	// A jitter calculator, just fills a channel with random numbers so that other
	// people don't have to build their own random generator...
	go func() {
		rnd := rand.New(rand.NewSource(time.Now().UnixNano()))
		for {
			jitter := rnd.Intn(HeartbeatInterval/2) + (HeartbeatInterval / 2)
			c.jitters <- time.Duration(jitter) * time.Second
		}
	}()

	// Now start the metadata refreshing goroutine
	go func() {
		for !c.Terminated() {
			time.Sleep(<-c.jitters)
			log.Infof("[%s] Refreshing topic metadata.", c.name)
			c.refreshMetadata()

			// See if the number of partitions in the marshal topic changed. This is bad if
			// it happens, since it means we can no longer coordinate correctly.
			if c.getTopicPartitions(MarshalTopic) != c.partitions {
				log.Errorf("[%s] Marshal topic partition count changed. Terminating!", c.name)
				c.Terminate()
			}
		}
	}()

	// Wait for all rationalizers to come alive
	log.Infof("[%s] Waiting for all rationalizers to come alive.", c.name)
	c.rationalizers.Wait()
	log.Infof("[%s] All rationalizers alive, KafkaCluster now alive.", c.name)

	return c, nil
}

// NewMarshaler creates a Marshaler off of an existing cluster. This is more efficient
// if you're creating multiple instances, since they can share the same underlying cluster.
func (c *KafkaCluster) NewMarshaler(clientID, groupID string) (*Marshaler, error) {
	if c.Terminated() {
		return nil, errors.New("Cluster is terminated.")
	}

	// Get offset coordinator so we can look up (and save) committed offsets later.
	coordinator, err := c.getOffsetCoordinator(groupID)
	if err != nil {
		return nil, err
	}

	m := &Marshaler{
		quit:       new(int32),
		cluster:    c,
		instanceID: newInstanceID(),
		clientID:   clientID,
		groupID:    groupID,
		offsets:    coordinator,
	}

	c.lock.Lock()
	defer c.lock.Unlock()

	// Remove any dead marshalers from our slice and add the new one
	filtered := make([]*Marshaler, 0)
	for _, marshaler := range c.marshalers {
		if !marshaler.Terminated() {
			filtered = append(filtered, marshaler)
		}
	}
	filtered = append(filtered, m)
	c.marshalers = filtered

	return m, nil
}

// refreshMetadata is periodically used to update our internal state with topic information
// about the world.
func (c *KafkaCluster) refreshMetadata() error {
	md, err := c.broker.Metadata()
	if err != nil {
		return err
	}

	newTopics := make(map[string]int)
	for _, topic := range md.Topics {
		newTopics[topic.Name] = len(topic.Partitions)
	}

	c.lock.Lock()
	defer c.lock.Unlock()
	c.topics = newTopics
	return nil
}

// getOffsetCoordinator returns a kafka.OffsetCoordinator for a specific group.
func (c *KafkaCluster) getOffsetCoordinator(groupID string) (kafka.OffsetCoordinator, error) {
	return c.broker.OffsetCoordinator(
		kafka.NewOffsetCoordinatorConf(groupID))
}

// getClaimPartition calculates which partition a topic should use for coordination. This uses
// a hashing function (non-cryptographic) to predictably partition the topic space.
func (c *KafkaCluster) getClaimPartition(topicName string) int {
	// We use MD5 because it's a fast and good hashing algorithm and we don't need cryptographic
	// properties. We then take the first 8 bytes and treat them as a uint64 and modulo that
	// across how many partitions we have.
	hash := md5.Sum([]byte(topicName))
	uval := binary.LittleEndian.Uint64(hash[0:8])
	return int(uval % uint64(c.partitions))
}

// getPartitionState returns a topicState and possibly creates it and the partition state within
// the State.
func (c *KafkaCluster) getPartitionState(groupID, topicName string, partID int) *topicState {
	c.lock.Lock()
	defer c.lock.Unlock()

	group, ok := c.groups[groupID]
	if !ok {
		group = make(map[string]*topicState)
		c.groups[groupID] = group
	}

	topic, ok := group[topicName]
	if !ok {
		topic = &topicState{
			claimPartition: c.getClaimPartition(topicName),
			partitions:     make([]PartitionClaim, partID+1),
		}
		group[topicName] = topic
	}

	// Take the topic lock if we can
	topic.lock.Lock()
	defer topic.lock.Unlock()

	// They might be referring to a partition we don't know about, maybe extend it
	if len(topic.partitions) < partID+1 {
		for i := len(topic.partitions); i <= partID; i++ {
			topic.partitions = append(topic.partitions, PartitionClaim{})
		}
	}

	return topic
}

// getTopics returns the list of known topics.
func (c *KafkaCluster) getTopics() []string {
	c.lock.RLock()
	defer c.lock.RUnlock()

	topics := make([]string, 0, len(c.topics))
	for topic := range c.topics {
		topics = append(topics, topic)
	}
	return topics
}

// getTopicPartitions returns the count of how many partitions are in a given topic. Returns 0 if a
// topic is unknown.
func (c *KafkaCluster) getTopicPartitions(topicName string) int {
	c.lock.RLock()
	defer c.lock.RUnlock()

	count, _ := c.topics[topicName]
	return count
}

// removeMarshal removes a terminated Marshal from a cluster's list.
func (c *KafkaCluster) removeMarshal(m *Marshaler) {
	c.lock.Lock()
	defer c.lock.Unlock()

	for i, ml := range c.marshalers {
		if ml == m {
			c.marshalers = append(c.marshalers[:i], c.marshalers[i+1:]...)
			break
		}
	}
}

// waitForRsteps is used by the test suite to ask the rationalizer to wait until some number
// of events have been processed. This also returns the current rsteps when it returns.
func (c *KafkaCluster) waitForRsteps(steps int) int {
	for {
		cval := atomic.LoadInt32(c.rsteps)
		if cval >= int32(steps) {
			return int(cval)
		}
		time.Sleep(5 * time.Millisecond)
	}
}

// pauseConsumerGroup stores an expiry time for consumer groups that we'd like to pause.
func (c *KafkaCluster) pauseConsumerGroup(groupID string, adminID string, expiry time.Time) {
	c.lock.Lock()
	defer c.lock.Unlock()

	log.Warningf("Cluster marking group %s paused with expiry: %s", groupID, expiry.Format(time.UnixDate))
	c.pausedGroups[groupID] = expiry
}

// IsGroupPaused returns true if the given consumer group is paused.
// TODO(pihu) This just checks the expiry time, and not the admin ID.
func (c *KafkaCluster) IsGroupPaused(groupID string) bool {
	c.lock.RLock()
	defer c.lock.RUnlock()

	if res, ok := c.pausedGroups[groupID]; !ok {
		return false
	} else {
		return time.Now().Before(res)
	}
}

// Terminate is called when we're done with the marshaler and want to shut down.
func (c *KafkaCluster) Terminate() {
	if !atomic.CompareAndSwapInt32(c.quit, 0, 1) {
		return
	}

	c.lock.Lock()
	defer c.lock.Unlock()

	// Terminate all Marshalers which will in turn terminate all Consumers and
	// let everybody know we're all done.
	for _, marshaler := range c.marshalers {
		marshaler.terminateAndCleanup(false)
	}
	c.marshalers = nil

	// Close the broker asynchronously to prevent blocking on potential network I/O
	go c.broker.Close()
}

// Terminated returns whether or not we have been terminated.
func (c *KafkaCluster) Terminated() bool {
	return atomic.LoadInt32(c.quit) == 1
}
<commit_msg>Refresh consumer metadata hourly.<commit_after>/*
 * portal - marshal
 *
 * a library that implements an algorithm for doing consumer coordination within Kafka, rather
 * than using Zookeeper or another external system.
 *
 */

package marshal

import (
	"crypto/md5"
	"encoding/binary"
	"errors"
	"fmt"
	"math/rand"
	"sync"
	"sync/atomic"
	"time"

	"github.com/dropbox/kafka"
)

// KafkaCluster is a user-agnostic view of the world. It connects to a Kafka cluster
// and runs rationalizers to observe the complete world state.
type KafkaCluster struct {
	// These members are not protected by the lock and can be read at any
	// time as they're write-once or only ever atomically updated. They must
	// never be overwritten once a KafkaCluster is created.
	quit       *int32
	name       string
	broker     *kafka.Broker
	producer   kafka.Producer
	partitions int
	jitters    chan time.Duration
	options    MarshalOptions

	// Lock protects the following members; you must have this lock in order to
	// read from or write to these.
	lock       sync.RWMutex
	marshalers []*Marshaler
	topics     map[string]int
	groups     map[string]map[string]*topicState
	// pausedGroups stores the expiry time for groups that are paused.
	pausedGroups map[string]time.Time

	// This WaitGroup is used for signalling when all of the rationalizers have
	// finished processing.
	rationalizers sync.WaitGroup

	// rsteps is updated whenever a rationalizer processes a log entry, this is
	// used mainly by the test suite.
	rsteps *int32

	// This is for testing only. When this is non-zero, the rationalizer will answer
	// queries based on THIS time instead of the current, actual time.
	ts int64
}

// MarshalOptions contains various tunables that can be used to adjust the configuration
// of the underlying system.
type MarshalOptions struct {
	// BrokerConnectionLimit is used to set the maximum simultaneous number of connections
	// that can be made to each broker.
	// Default: 30.
	BrokerConnectionLimit int

	// ConsumeRequestTimeout sets the time that we ask Kafka to wait before returning any
	// data to us. Setting this high uses more connections and can lead to some latency
	// but keeps the load on Kafka minimal. Use this to balance QPS against latency.
	// Default: 3 seconds.
	ConsumeRequestTimeout time.Duration

	// MarshalRequestTimeout is used for our coordination requests. This should be reasonable
	// at default, but is left as a tunable in case you have clients that are claiming an
	// extremely large number of partitions and are too slow.
	// Default: 1 second.
	MarshalRequestTimeout time.Duration
}

// NewMarshalOptions returns a set of MarshalOptions populated with defaults.
func NewMarshalOptions() MarshalOptions {
	return MarshalOptions{
		BrokerConnectionLimit: 30,
		ConsumeRequestTimeout: 3 * time.Second,
		MarshalRequestTimeout: 1 * time.Second,
	}
}

// Dial returns a new cluster object which can be used to instantiate a number of Marshalers
// that all use the same cluster. You may pass brokerConf or may set it to nil.
func Dial(name string, brokers []string, options MarshalOptions) (*KafkaCluster, error) {
	// Connect to Kafka
	brokerConf := kafka.NewBrokerConf("PortalMarshal")
	brokerConf.MetadataRefreshFrequency = time.Hour
	brokerConf.ConnectionLimit = options.BrokerConnectionLimit
	broker, err := kafka.Dial(brokers, brokerConf)
	if err != nil {
		return nil, err
	}

	c := &KafkaCluster{
		quit:         new(int32),
		rsteps:       new(int32),
		name:         name,
		options:      options,
		broker:       broker,
		producer:     broker.Producer(kafka.NewProducerConf()),
		topics:       make(map[string]int),
		groups:       make(map[string]map[string]*topicState),
		pausedGroups: make(map[string]time.Time),
		jitters:      make(chan time.Duration, 100),
	}

	// Do an initial metadata fetch, this will block a bit
	err = c.refreshMetadata()
	if err != nil {
		return nil, fmt.Errorf("Failed to get metadata: %s", err)
	}

	// If there is no marshal topic, then we can't run. The admins must go create the topic
	// before they can use this library. Please see the README.
	c.partitions = c.getTopicPartitions(MarshalTopic)
	if c.partitions == 0 {
		return nil, errors.New("Marshalling topic not found. Please see the documentation.")
	}

	// Now we start a goroutine to start consuming each of the partitions in the marshal
	// topic. Note that this doesn't handle increasing the partition count on that topic
	// without stopping all consumers.
	c.rationalizers.Add(c.partitions)
	for id := 0; id < c.partitions; id++ {
		go c.rationalize(id, c.kafkaConsumerChannel(id))
	}

	// A jitter calculator, just fills a channel with random numbers so that other
	// people don't have to build their own random generator...
	go func() {
		rnd := rand.New(rand.NewSource(time.Now().UnixNano()))
		for {
			jitter := rnd.Intn(HeartbeatInterval/2) + (HeartbeatInterval / 2)
			c.jitters <- time.Duration(jitter) * time.Second
		}
	}()

	// Now start the metadata refreshing goroutine
	go func() {
		for !c.Terminated() {
			time.Sleep(<-c.jitters)
			log.Infof("[%s] Refreshing topic metadata.", c.name)
			c.refreshMetadata()

			// See if the number of partitions in the marshal topic changed. This is bad if
			// it happens, since it means we can no longer coordinate correctly.
			if c.getTopicPartitions(MarshalTopic) != c.partitions {
				log.Errorf("[%s] Marshal topic partition count changed. Terminating!", c.name)
				c.Terminate()
			}
		}
	}()

	// Wait for all rationalizers to come alive
	log.Infof("[%s] Waiting for all rationalizers to come alive.", c.name)
	c.rationalizers.Wait()
	log.Infof("[%s] All rationalizers alive, KafkaCluster now alive.", c.name)

	return c, nil
}

// NewMarshaler creates a Marshaler off of an existing cluster. This is more efficient
// if you're creating multiple instances, since they can share the same underlying cluster.
func (c *KafkaCluster) NewMarshaler(clientID, groupID string) (*Marshaler, error) {
	if c.Terminated() {
		return nil, errors.New("Cluster is terminated.")
	}

	// Get offset coordinator so we can look up (and save) committed offsets later.
	coordinator, err := c.getOffsetCoordinator(groupID)
	if err != nil {
		return nil, err
	}

	m := &Marshaler{
		quit:       new(int32),
		cluster:    c,
		instanceID: newInstanceID(),
		clientID:   clientID,
		groupID:    groupID,
		offsets:    coordinator,
	}

	c.lock.Lock()
	defer c.lock.Unlock()

	// Remove any dead marshalers from our slice and add the new one
	filtered := make([]*Marshaler, 0)
	for _, marshaler := range c.marshalers {
		if !marshaler.Terminated() {
			filtered = append(filtered, marshaler)
		}
	}
	filtered = append(filtered, m)
	c.marshalers = filtered

	return m, nil
}

// refreshMetadata is periodically used to update our internal state with topic information
// about the world.
func (c *KafkaCluster) refreshMetadata() error {
	md, err := c.broker.Metadata()
	if err != nil {
		return err
	}

	newTopics := make(map[string]int)
	for _, topic := range md.Topics {
		newTopics[topic.Name] = len(topic.Partitions)
	}

	c.lock.Lock()
	defer c.lock.Unlock()
	c.topics = newTopics
	return nil
}

// getOffsetCoordinator returns a kafka.OffsetCoordinator for a specific group.
func (c *KafkaCluster) getOffsetCoordinator(groupID string) (kafka.OffsetCoordinator, error) {
	return c.broker.OffsetCoordinator(
		kafka.NewOffsetCoordinatorConf(groupID))
}

// getClaimPartition calculates which partition a topic should use for coordination. This uses
// a hashing function (non-cryptographic) to predictably partition the topic space.
func (c *KafkaCluster) getClaimPartition(topicName string) int {
	// We use MD5 because it's a fast and good hashing algorithm and we don't need cryptographic
	// properties. We then take the first 8 bytes and treat them as a uint64 and modulo that
	// across how many partitions we have.
	hash := md5.Sum([]byte(topicName))
	uval := binary.LittleEndian.Uint64(hash[0:8])
	return int(uval % uint64(c.partitions))
}

// getPartitionState returns a topicState and possibly creates it and the partition state within
// the State.
func (c *KafkaCluster) getPartitionState(groupID, topicName string, partID int) *topicState {
	c.lock.Lock()
	defer c.lock.Unlock()

	group, ok := c.groups[groupID]
	if !ok {
		group = make(map[string]*topicState)
		c.groups[groupID] = group
	}

	topic, ok := group[topicName]
	if !ok {
		topic = &topicState{
			claimPartition: c.getClaimPartition(topicName),
			partitions:     make([]PartitionClaim, partID+1),
		}
		group[topicName] = topic
	}

	// Take the topic lock if we can
	topic.lock.Lock()
	defer topic.lock.Unlock()

	// They might be referring to a partition we don't know about, maybe extend it
	if len(topic.partitions) < partID+1 {
		for i := len(topic.partitions); i <= partID; i++ {
			topic.partitions = append(topic.partitions, PartitionClaim{})
		}
	}

	return topic
}

// getTopics returns the list of known topics.
func (c *KafkaCluster) getTopics() []string {
	c.lock.RLock()
	defer c.lock.RUnlock()

	topics := make([]string, 0, len(c.topics))
	for topic := range c.topics {
		topics = append(topics, topic)
	}
	return topics
}

// getTopicPartitions returns the count of how many partitions are in a given topic. Returns 0 if a
// topic is unknown.
func (c *KafkaCluster) getTopicPartitions(topicName string) int {
	c.lock.RLock()
	defer c.lock.RUnlock()

	count, _ := c.topics[topicName]
	return count
}

// removeMarshal removes a terminated Marshal from a cluster's list.
func (c *KafkaCluster) removeMarshal(m *Marshaler) {
	c.lock.Lock()
	defer c.lock.Unlock()

	for i, ml := range c.marshalers {
		if ml == m {
			c.marshalers = append(c.marshalers[:i], c.marshalers[i+1:]...)
			break
		}
	}
}

// waitForRsteps is used by the test suite to ask the rationalizer to wait until some number
// of events have been processed. This also returns the current rsteps when it returns.
func (c *KafkaCluster) waitForRsteps(steps int) int {
	for {
		cval := atomic.LoadInt32(c.rsteps)
		if cval >= int32(steps) {
			return int(cval)
		}
		time.Sleep(5 * time.Millisecond)
	}
}

// pauseConsumerGroup stores an expiry time for consumer groups that we'd like to pause.
func (c *KafkaCluster) pauseConsumerGroup(groupID string, adminID string, expiry time.Time) {
	c.lock.Lock()
	defer c.lock.Unlock()

	log.Warningf("Cluster marking group %s paused with expiry: %s", groupID, expiry.Format(time.UnixDate))
	c.pausedGroups[groupID] = expiry
}

// IsGroupPaused returns true if the given consumer group is paused.
// TODO(pihu) This just checks the expiry time, and not the admin ID.
func (c *KafkaCluster) IsGroupPaused(groupID string) bool {
	c.lock.RLock()
	defer c.lock.RUnlock()

	if res, ok := c.pausedGroups[groupID]; !ok {
		return false
	} else {
		return time.Now().Before(res)
	}
}

// Terminate is called when we're done with the marshaler and want to shut down.
func (c *KafkaCluster) Terminate() {
	if !atomic.CompareAndSwapInt32(c.quit, 0, 1) {
		return
	}

	c.lock.Lock()
	defer c.lock.Unlock()

	// Terminate all Marshalers which will in turn terminate all Consumers and
	// let everybody know we're all done.
	for _, marshaler := range c.marshalers {
		marshaler.terminateAndCleanup(false)
	}
	c.marshalers = nil

	// Close the broker asynchronously to prevent blocking on potential network I/O
	go c.broker.Close()
}

// Terminated returns whether or not we have been terminated.
func (c *KafkaCluster) Terminated() bool {
	return atomic.LoadInt32(c.quit) == 1
}
<|endoftext|>


<commit_before>/*
Copyright 2017 The Kubernetes Authors.

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
*/

package snapshotter

import (
	"fmt"

	"github.com/golang/glog"

	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	"k8s.io/apimachinery/pkg/runtime"
	"k8s.io/client-go/kubernetes"
	v1 "k8s.io/client-go/pkg/api/v1"
	"k8s.io/client-go/rest"
	"k8s.io/kubernetes/pkg/util/goroutinemap"
	"k8s.io/kubernetes/pkg/util/goroutinemap/exponentialbackoff"

	tprv1 "github.com/rootfs/snapshot/pkg/apis/tpr/v1"
	"github.com/rootfs/snapshot/pkg/controller/cache"

	"github.com/rootfs/snapshot/pkg/volume/hostpath"
)

const (
	defaultExponentialBackOffOnError = true
)

// VolumeSnapshotter does the "heavy lifting": it spawns gouroutines that talk to the
// backend to actually perform the operations on the storage devices.
// It creates and deletes the snapshots and promotes snapshots to volumes (PV). The create
// and delete operations need to be idempotent and count with the fact the API object writes
type VolumeSnapshotter interface {
	CreateVolumeSnapshot(snapshotName string, snapshotSpec *tprv1.VolumeSnapshotSpec)
	DeleteVolumeSnapshot(snapshotName string, snapshotSpec *tprv1.VolumeSnapshotSpec)
	PromoteVolumeSnapshotToPV(snapshotName string, snapshotSpec *tprv1.VolumeSnapshotSpec)
}

type volumeSnapshotter struct {
	restClient         *rest.RESTClient
	coreClient         kubernetes.Interface
	scheme             *runtime.Scheme
	actualStateOfWorld cache.ActualStateOfWorld
	runningOperation   goroutinemap.GoRoutineMap
}

const (
	snapshotOpCreatePrefix  string = "create"
	snapshotOpDeletePrefix  string = "delete"
	snapshotOpPromotePrefix string = "promote"
)

func NewVolumeSnapshotter(
	restClient *rest.RESTClient,
	scheme *runtime.Scheme,
	clientset kubernetes.Interface,
	asw cache.ActualStateOfWorld) VolumeSnapshotter {
	return &volumeSnapshotter{
		restClient:         restClient,
		coreClient:         clientset,
		scheme:             scheme,
		actualStateOfWorld: asw,
		runningOperation:   goroutinemap.NewGoRoutineMap(defaultExponentialBackOffOnError),
	}
}

// This is the function responsible for determining the correct volume plugin to use,
// asking it to make a snapshot and assignig it some name that it returns to the caller.
func (vs *volumeSnapshotter) takeSnapshot(spec *v1.PersistentVolumeSpec) (*tprv1.VolumeSnapshotDataSource, error) {
	// TODO: Find a volume snapshot plugin to use for taking the snapshot and do so
	if spec.HostPath != nil {
		snap, err := hostpath.Snapshot(spec.HostPath.Path)
		if err != nil {
			glog.Warningf("failed to snapshot %s, err: %v", spec.HostPath.Path, err)
		} else {
			glog.Infof("snapshot %#v to snap %#v", spec.HostPath, snap.HostPath)
			return snap, nil
		}
	}

	return nil, nil
}

// Below are the closures meant to build the functions for the GoRoutineMap operations.

func (vs *volumeSnapshotter) getSnapshotCreateFunc(snapshotName string, snapshotSpec *tprv1.VolumeSnapshotSpec) func() error {
	// Create a snapshot:
	// 1. If Snapshot referencs SnapshotData object, try to find it
	//   1a. If doesn't exist, log error and finish, if it exists already, check its SnapshotRef
	//   1b. If it's empty, check its Spec UID (or fins out what PV/PVC does and copyt the mechanism)
	//   1c. If it matches the user (TODO: how to find out?), bind the two objects and finish
	//   1d. If it doesn't match, log error and finish.
	// 2. Create the SnapshotData object
	// 3. Ask the backend to create the snapshot (device)
	// 4. If OK, update the SnapshotData and Snapshot objects
	// 5. Add the Snapshot to the ActualStateOfWorld
	// 6. Finish (we have created snapshot for an user)
	return func() error {
		// TODO: return if VolumeSnapshotData is already created
		if snapshotSpec.SnapshotDataName != "" {
			// This spec has the SnapshotDataName already set: this means importing admin-created snapshots
			// TODO: Not implemented yet
			return fmt.Errorf("Importing snapshots is not implemented yet")
		}

		pvcName := snapshotSpec.PersistentVolumeClaimName
		if pvcName == "" {
			return fmt.Errorf("The PVC name is not specified in snapshot %s", snapshotName)
		}
		snapNameSpace, snapName, err := cache.GetNameAndNameSpaceFromSnapshotName(snapshotName)
		if err != nil {
			return fmt.Errorf("Snapshot %s is malformed", snapshotName)
		}
		pvc, err := vs.coreClient.CoreV1().PersistentVolumeClaims(snapNameSpace).Get(pvcName, metav1.GetOptions{})
		if err != nil {
			return fmt.Errorf("Failed to retrieve PVC %s from the API server: %q", pvcName, err)
		}
		if pvc.Status.Phase != v1.ClaimBound {
			return fmt.Errorf("The PVC %s not yet bound to a PV, will not attempt to take a snapshot yet.")
		}

		pvName := pvc.Spec.VolumeName
		pv, err := vs.coreClient.CoreV1().PersistentVolumes().Get(pvName, metav1.GetOptions{})
		if err != nil {
			return fmt.Errorf("Failed to retrieve PV %s from the API server: %q", pvName, err)
		}

		snapshotDataSource, err := vs.takeSnapshot(&pv.Spec)
		if err != nil || snapshotDataSource == nil {
			return fmt.Errorf("Failed to take snapshot of the volume %s: %q", pvName, err)
		}
		// Snapshot has been created, made an object for it
		snapshotData := &tprv1.VolumeSnapshotData{
			Metadata: metav1.ObjectMeta{
				// FIXME: make a unique ID
				Name: snapName,
			},
			Spec: tprv1.VolumeSnapshotDataSpec{
				VolumeSnapshotRef: &v1.ObjectReference{
					Kind: "VolumeSnapshot",
					Name: snapshotName,
				},
				PersistentVolumeRef: &v1.ObjectReference{
					Kind: "PersistentVolume",
					Name: pvName,
				},
				VolumeSnapshotDataSource: *snapshotDataSource,
			},
		}
		var result tprv1.VolumeSnapshotData
		err = vs.restClient.Post().
			Resource(tprv1.VolumeSnapshotDataResourcePlural).
			Namespace(v1.NamespaceDefault).
			Body(snapshotData).
			Do().Into(&result)

		if err != nil {
			//FIXME: Errors writing to the API server are common: this needs to be re-tried
			glog.Warningf("Error creating the VolumeSnapshotData %s: %v", snapshotName, err)
		}
		vs.actualStateOfWorld.AddSnapshot(snapshotName, snapshotSpec)
		// TODO: Update the VolumeSnapshot object too

		return nil
	}
}

func (vs *volumeSnapshotter) getSnapshotDeleteFunc(snapshotName string, snapshotSpec *tprv1.VolumeSnapshotSpec) func() error {
	// Delete a snapshot
	// 1. Find the SnapshotData corresponding to Snapshot
	//   1a: Not found => finish (it's been deleted already)
	// 2. Ask the backend to remove the snapshot device
	// 3. Delete the SnapshotData object
	// 4. Remove the Snapshot from ActualStateOfWorld
	// 5. Finish
	return func() error {
		vs.actualStateOfWorld.DeleteSnapshot(snapshotName)

		return nil
	}
}

func (vs *volumeSnapshotter) getSnapshotPromoteFunc(snapshotName string, snapshotSpec *tprv1.VolumeSnapshotSpec) func() error {
	// Promote snapshot to a PVC
	// 1. We have a PVC referencing a Snapshot object
	// 2. Find the SnapshotData corresponding to tha Snapshot
	// 3. Ask the backend to give us a device (PV) made from the snapshot device
	// 4. Bind it to the PVC
	// 5. Finish
	return func() error { return nil }
}

func (vs *volumeSnapshotter) CreateVolumeSnapshot(snapshotName string, snapshotSpec *tprv1.VolumeSnapshotSpec) {
	operationName := snapshotOpCreatePrefix + snapshotName + snapshotSpec.PersistentVolumeClaimName
	glog.Infof("Snapshotter is about to create volume snapshot operation named %s, spec %#v", operationName, snapshotSpec)

	err := vs.runningOperation.Run(operationName, vs.getSnapshotCreateFunc(snapshotName, snapshotSpec))

	if err != nil {
		switch {
		case goroutinemap.IsAlreadyExists(err):
			glog.V(4).Infof("operation %q is already running, skipping", operationName)
		case exponentialbackoff.IsExponentialBackoff(err):
			glog.V(4).Infof("operation %q postponed due to exponential backoff", operationName)
		default:
			glog.Errorf("Failed to schedule the operation %q: %v", operationName, err)
		}
	}
}

func (vs *volumeSnapshotter) DeleteVolumeSnapshot(snapshotName string, snapshotSpec *tprv1.VolumeSnapshotSpec) {
	operationName := snapshotOpDeletePrefix + snapshotName + snapshotSpec.PersistentVolumeClaimName
	glog.Infof("Snapshotter is about to create volume snapshot operation named %s", operationName)

	err := vs.runningOperation.Run(operationName, vs.getSnapshotDeleteFunc(snapshotName, snapshotSpec))

	if err != nil {
		switch {
		case goroutinemap.IsAlreadyExists(err):
			glog.V(4).Infof("operation %q is already running, skipping", operationName)
		case exponentialbackoff.IsExponentialBackoff(err):
			glog.V(4).Infof("operation %q postponed due to exponential backoff", operationName)
		default:
			glog.Errorf("Failed to schedule the operation %q: %v", operationName, err)
		}
	}
}

func (vs *volumeSnapshotter) PromoteVolumeSnapshotToPV(snapshotName string, snapshotSpec *tprv1.VolumeSnapshotSpec) {
	operationName := snapshotOpPromotePrefix + snapshotName + snapshotSpec.PersistentVolumeClaimName
	glog.Infof("Snapshotter is about to create volume snapshot operation named %s", operationName)

	err := vs.runningOperation.Run(operationName, vs.getSnapshotPromoteFunc(snapshotName, snapshotSpec))

	if err != nil {
		switch {
		case goroutinemap.IsAlreadyExists(err):
			glog.V(4).Infof("operation %q is already running, skipping", operationName)
		case exponentialbackoff.IsExponentialBackoff(err):
			glog.V(4).Infof("operation %q postponed due to exponential backoff", operationName)
		default:
			glog.Errorf("Failed to schedule the operation %q: %v", operationName, err)
		}
	}
}
<commit_msg>WIP: update the volume snapshot object on successful creation<commit_after>/*
Copyright 2017 The Kubernetes Authors.

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
*/

package snapshotter

import (
	"fmt"

	"github.com/golang/glog"

	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	"k8s.io/apimachinery/pkg/runtime"
	"k8s.io/client-go/kubernetes"
	v1 "k8s.io/client-go/pkg/api/v1"
	"k8s.io/client-go/rest"
	"k8s.io/kubernetes/pkg/util/goroutinemap"
	"k8s.io/kubernetes/pkg/util/goroutinemap/exponentialbackoff"

	tprv1 "github.com/rootfs/snapshot/pkg/apis/tpr/v1"
	"github.com/rootfs/snapshot/pkg/controller/cache"

	"github.com/rootfs/snapshot/pkg/volume/hostpath"
)

const (
	defaultExponentialBackOffOnError = true
)

// VolumeSnapshotter does the "heavy lifting": it spawns gouroutines that talk to the
// backend to actually perform the operations on the storage devices.
// It creates and deletes the snapshots and promotes snapshots to volumes (PV). The create
// and delete operations need to be idempotent and count with the fact the API object writes
type VolumeSnapshotter interface {
	CreateVolumeSnapshot(snapshotName string, snapshotSpec *tprv1.VolumeSnapshotSpec)
	DeleteVolumeSnapshot(snapshotName string, snapshotSpec *tprv1.VolumeSnapshotSpec)
	PromoteVolumeSnapshotToPV(snapshotName string, snapshotSpec *tprv1.VolumeSnapshotSpec)
}

type volumeSnapshotter struct {
	restClient         *rest.RESTClient
	coreClient         kubernetes.Interface
	scheme             *runtime.Scheme
	actualStateOfWorld cache.ActualStateOfWorld
	runningOperation   goroutinemap.GoRoutineMap
}

const (
	snapshotOpCreatePrefix  string = "create"
	snapshotOpDeletePrefix  string = "delete"
	snapshotOpPromotePrefix string = "promote"
)

func NewVolumeSnapshotter(
	restClient *rest.RESTClient,
	scheme *runtime.Scheme,
	clientset kubernetes.Interface,
	asw cache.ActualStateOfWorld) VolumeSnapshotter {
	return &volumeSnapshotter{
		restClient:         restClient,
		coreClient:         clientset,
		scheme:             scheme,
		actualStateOfWorld: asw,
		runningOperation:   goroutinemap.NewGoRoutineMap(defaultExponentialBackOffOnError),
	}
}

// This is the function responsible for determining the correct volume plugin to use,
// asking it to make a snapshot and assignig it some name that it returns to the caller.
func (vs *volumeSnapshotter) takeSnapshot(spec *v1.PersistentVolumeSpec) (*tprv1.VolumeSnapshotDataSource, error) {
	// TODO: Find a volume snapshot plugin to use for taking the snapshot and do so
	if spec.HostPath != nil {
		snap, err := hostpath.Snapshot(spec.HostPath.Path)
		if err != nil {
			glog.Warningf("failed to snapshot %s, err: %v", spec.HostPath.Path, err)
		} else {
			glog.Infof("snapshot %#v to snap %#v", spec.HostPath, snap.HostPath)
			return snap, nil
		}
	}

	return nil, nil
}

// Below are the closures meant to build the functions for the GoRoutineMap operations.

func (vs *volumeSnapshotter) getSnapshotCreateFunc(snapshotName string, snapshotSpec *tprv1.VolumeSnapshotSpec) func() error {
	// Create a snapshot:
	// 1. If Snapshot referencs SnapshotData object, try to find it
	//   1a. If doesn't exist, log error and finish, if it exists already, check its SnapshotRef
	//   1b. If it's empty, check its Spec UID (or fins out what PV/PVC does and copyt the mechanism)
	//   1c. If it matches the user (TODO: how to find out?), bind the two objects and finish
	//   1d. If it doesn't match, log error and finish.
	// 2. Create the SnapshotData object
	// 3. Ask the backend to create the snapshot (device)
	// 4. If OK, update the SnapshotData and Snapshot objects
	// 5. Add the Snapshot to the ActualStateOfWorld
	// 6. Finish (we have created snapshot for an user)
	return func() error {
		// TODO: return if VolumeSnapshotData is already created
		if snapshotSpec.SnapshotDataName != "" {
			// This spec has the SnapshotDataName already set: this means importing admin-created snapshots
			// TODO: Not implemented yet
			return fmt.Errorf("Importing snapshots is not implemented yet")
		}

		pvcName := snapshotSpec.PersistentVolumeClaimName
		if pvcName == "" {
			return fmt.Errorf("The PVC name is not specified in snapshot %s", snapshotName)
		}
		snapNameSpace, snapName, err := cache.GetNameAndNameSpaceFromSnapshotName(snapshotName)
		if err != nil {
			return fmt.Errorf("Snapshot %s is malformed", snapshotName)
		}
		pvc, err := vs.coreClient.CoreV1().PersistentVolumeClaims(snapNameSpace).Get(pvcName, metav1.GetOptions{})
		if err != nil {
			return fmt.Errorf("Failed to retrieve PVC %s from the API server: %q", pvcName, err)
		}
		if pvc.Status.Phase != v1.ClaimBound {
			return fmt.Errorf("The PVC %s not yet bound to a PV, will not attempt to take a snapshot yet.")
		}

		pvName := pvc.Spec.VolumeName
		pv, err := vs.coreClient.CoreV1().PersistentVolumes().Get(pvName, metav1.GetOptions{})
		if err != nil {
			return fmt.Errorf("Failed to retrieve PV %s from the API server: %q", pvName, err)
		}

		snapshotDataSource, err := vs.takeSnapshot(&pv.Spec)
		if err != nil || snapshotDataSource == nil {
			return fmt.Errorf("Failed to take snapshot of the volume %s: %q", pvName, err)
		}
		// Snapshot has been created, made an object for it
		readyCondition := tprv1.VolumeSnapshotDataCondition{
			Type:    tprv1.VolumeSnapshotDataConditionReady,
			Status:  v1.ConditionTrue,
			Message: "Snapsot created succsessfully",
		}
		snapshotData := &tprv1.VolumeSnapshotData{
			Metadata: metav1.ObjectMeta{
				// FIXME: make a unique ID
				Name: snapName,
			},
			Spec: tprv1.VolumeSnapshotDataSpec{
				VolumeSnapshotRef: &v1.ObjectReference{
					Kind: "VolumeSnapshot",
					Name: snapshotName,
				},
				PersistentVolumeRef: &v1.ObjectReference{
					Kind: "PersistentVolume",
					Name: pvName,
				},
				VolumeSnapshotDataSource: *snapshotDataSource,
			},
			Status: tprv1.VolumeSnapshotDataStatus{
				Conditions: []tprv1.VolumeSnapshotDataCondition{
					readyCondition,
				},
			},
		}
		var result tprv1.VolumeSnapshotData
		err = vs.restClient.Post().
			Resource(tprv1.VolumeSnapshotDataResourcePlural).
			Namespace(v1.NamespaceDefault).
			Body(snapshotData).
			Do().Into(&result)

		if err != nil {
			//FIXME: Errors writing to the API server are common: this needs to be re-tried
			glog.Errorf("Error creating the VolumeSnapshotData %s: %v", snapshotName, err)
		}
		vs.actualStateOfWorld.AddSnapshot(snapshotName, snapshotSpec)
		// TODO: Update the VolumeSnapshot object too

		var snapshotObj tprv1.VolumeSnapshot
		err = vs.restClient.Get().
			Name(snapName).
			Resource(tprv1.VolumeSnapshotResourcePlural).
			Namespace(v1.NamespaceDefault).
			Do().Into(&snapshotObj)

		objCopy, err := vs.scheme.DeepCopy(&snapshotObj)
		if err != nil {
			glog.Errorf("Error copying snapshot object %s object from API server", snapName)
		}
		snapshotCopy, ok := objCopy.(*tprv1.VolumeSnapshot)
		if !ok {
			glog.Warning("expecting type VolumeSnapshot but received type %T", objCopy)
		}
		snapshotCopy.Spec.SnapshotDataName = snapName
		snapshotCopy.Status.Conditions = []tprv1.VolumeSnapshotCondition{
			{
				Type:    tprv1.VolumeSnapshotConditionReady,
				Status:  v1.ConditionTrue,
				Message: "Snapsot created succsessfully",
			},
		}
		// TODO: Make diff of the two objects and then use restClient.Patch to update it
		return nil
	}
}

func (vs *volumeSnapshotter) getSnapshotDeleteFunc(snapshotName string, snapshotSpec *tprv1.VolumeSnapshotSpec) func() error {
	// Delete a snapshot
	// 1. Find the SnapshotData corresponding to Snapshot
	//   1a: Not found => finish (it's been deleted already)
	// 2. Ask the backend to remove the snapshot device
	// 3. Delete the SnapshotData object
	// 4. Remove the Snapshot from ActualStateOfWorld
	// 5. Finish
	return func() error {
		vs.actualStateOfWorld.DeleteSnapshot(snapshotName)

		return nil
	}
}

func (vs *volumeSnapshotter) getSnapshotPromoteFunc(snapshotName string, snapshotSpec *tprv1.VolumeSnapshotSpec) func() error {
	// Promote snapshot to a PVC
	// 1. We have a PVC referencing a Snapshot object
	// 2. Find the SnapshotData corresponding to tha Snapshot
	// 3. Ask the backend to give us a device (PV) made from the snapshot device
	// 4. Bind it to the PVC
	// 5. Finish
	return func() error { return nil }
}

func (vs *volumeSnapshotter) CreateVolumeSnapshot(snapshotName string, snapshotSpec *tprv1.VolumeSnapshotSpec) {
	operationName := snapshotOpCreatePrefix + snapshotName + snapshotSpec.PersistentVolumeClaimName
	glog.Infof("Snapshotter is about to create volume snapshot operation named %s, spec %#v", operationName, snapshotSpec)

	err := vs.runningOperation.Run(operationName, vs.getSnapshotCreateFunc(snapshotName, snapshotSpec))

	if err != nil {
		switch {
		case goroutinemap.IsAlreadyExists(err):
			glog.V(4).Infof("operation %q is already running, skipping", operationName)
		case exponentialbackoff.IsExponentialBackoff(err):
			glog.V(4).Infof("operation %q postponed due to exponential backoff", operationName)
		default:
			glog.Errorf("Failed to schedule the operation %q: %v", operationName, err)
		}
	}
}

func (vs *volumeSnapshotter) DeleteVolumeSnapshot(snapshotName string, snapshotSpec *tprv1.VolumeSnapshotSpec) {
	operationName := snapshotOpDeletePrefix + snapshotName + snapshotSpec.PersistentVolumeClaimName
	glog.Infof("Snapshotter is about to create volume snapshot operation named %s", operationName)

	err := vs.runningOperation.Run(operationName, vs.getSnapshotDeleteFunc(snapshotName, snapshotSpec))

	if err != nil {
		switch {
		case goroutinemap.IsAlreadyExists(err):
			glog.V(4).Infof("operation %q is already running, skipping", operationName)
		case exponentialbackoff.IsExponentialBackoff(err):
			glog.V(4).Infof("operation %q postponed due to exponential backoff", operationName)
		default:
			glog.Errorf("Failed to schedule the operation %q: %v", operationName, err)
		}
	}
}

func (vs *volumeSnapshotter) PromoteVolumeSnapshotToPV(snapshotName string, snapshotSpec *tprv1.VolumeSnapshotSpec) {
	operationName := snapshotOpPromotePrefix + snapshotName + snapshotSpec.PersistentVolumeClaimName
	glog.Infof("Snapshotter is about to create volume snapshot operation named %s", operationName)

	err := vs.runningOperation.Run(operationName, vs.getSnapshotPromoteFunc(snapshotName, snapshotSpec))

	if err != nil {
		switch {
		case goroutinemap.IsAlreadyExists(err):
			glog.V(4).Infof("operation %q is already running, skipping", operationName)
		case exponentialbackoff.IsExponentialBackoff(err):
			glog.V(4).Infof("operation %q postponed due to exponential backoff", operationName)
		default:
			glog.Errorf("Failed to schedule the operation %q: %v", operationName, err)
		}
	}
}
<|endoftext|>