<commit_before>/*
 * Copyright (C) 2014 Michael Wendland
 *
 * This program is free software: you can redistribute it and/or modify
 * it under the terms of the GNU Affero General Public License as published
 * by the Free Software Foundation, either version 3 of the License, or
 * (at your option) any later version.
 *
 * This program is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
 * GNU Affero General Public License for more details.
 *
 * You should have received a copy of the GNU Affero General Public License
 * along with this program.  If not, see <http://www.gnu.org/licenses/>.
 *
 * Authors:
 *   Michael Wendland <michael@michiwend.com>
 */

package goefa

import (
	"net/url"
	"strconv"
	"time"
)

// EFADepartureArrival represents either an arrival or a departure and is not
// directly used but embedded by EFAArrival and EFADeparture.
type EFADepartureArrival struct {
	Area         int    `xml:"area,attr"`
	Countdown    int    `xml:"countdown,attr"`
	MapName      string `xml:"mapName,attr"`
	Platform     string `xml:"platform,attr"`
	PlatformName string `xml:"platformName,attr"`
	StopID       int    `xml:"displayName,attr"`
	StopName     string `xml:"stopName,attr"`
	Lat          int64  `xml:"x,attr"`
	Lng          int64  `xml:"y,attr"`

	DateTime    EFATime        `xml:"itdDateTime"`
	ServingLine EFAServingLine `xml:"itdServingLine"`
}

// EFAArrival represents a single arrival for a specific stop.
type EFAArrival struct {
	EFADepartureArrival `xml:"itdArrival"`
}

// EFADeparture represents a single departure for a specific stop.
type EFADeparture struct {
	EFADepartureArrival `xml:"itdDeparture"`
}

type departureMonitorResult struct {
	efaResult
	Odv struct {
		OdvPlace struct {
		}
		OdvName struct {
			State string `xml:"state,attr"`
		} `xml:"itdOdvName"`
	} `xml:"itdDepartureMonitorRequest>itdOdv"`
	Departures []*EFADeparture `xml:"itdDepartureMonitorRequest>itdDepartureList>itdDeparture"`
}

func (d *departureMonitorResult) endpoint() string {
	return "XML_DM_REQUEST"
}

// Departures performs a stateless dm_request for the corresponding stopID and
// returns an array of EFADepartures. Use time.Now() as the second argument in
// order to get the very next departures. The third argument determines how
// many results will be returned by EFA.
func (efa *EFAProvider) Departures(stopID int, due time.Time, results int) ([]*EFADeparture, error) {
	var rt string

	if efa.EnableRealtime {
		rt = "1"
	} else {
		rt = "0"
	}

	params := url.Values{
		"type_dm":              {"any"},
		"name_dm":              {strconv.Itoa(stopID)},
		"locationServerActive": {"1"},
		"useRealtime":          {rt},
		"dmLineSelection":      {"all"}, //FIXME enable line selection
		"limit":                {strconv.Itoa(results)},
		"mode":                 {"direct"},
		"stateless":            {"1"},
		"itdDate":              {due.Format("20060102")},
		"itdTime":              {due.Format("1504")},
	}

	var result departureMonitorResult

	if err := efa.postRequest(&result, params); err != nil {
		return nil, err
	}

	return result.Departures, nil

}
<commit_msg>small fix<commit_after>/*
 * Copyright (C) 2014 Michael Wendland
 *
 * This program is free software: you can redistribute it and/or modify
 * it under the terms of the GNU Affero General Public License as published
 * by the Free Software Foundation, either version 3 of the License, or
 * (at your option) any later version.
 *
 * This program is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
 * GNU Affero General Public License for more details.
 *
 * You should have received a copy of the GNU Affero General Public License
 * along with this program.  If not, see <http://www.gnu.org/licenses/>.
 *
 * Authors:
 *   Michael Wendland <michael@michiwend.com>
 */

package goefa

import (
	"net/url"
	"strconv"
	"time"
)

// EFADepartureArrival represents either an arrival or a departure and is not
// directly used but embedded by EFAArrival and EFADeparture.
type EFADepartureArrival struct {
	Area         int    `xml:"area,attr"`
	Countdown    int    `xml:"countdown,attr"`
	MapName      string `xml:"mapName,attr"`
	Platform     string `xml:"platform,attr"`
	PlatformName string `xml:"platformName,attr"`
	StopID       int    `xml:"displayName,attr"`
	StopName     string `xml:"stopName,attr"`
	Lat          int64  `xml:"x,attr"`
	Lng          int64  `xml:"y,attr"`

	DateTime    EFATime        `xml:"itdDateTime"`
	ServingLine EFAServingLine `xml:"itdServingLine"`
}

// EFAArrival represents a single arrival for a specific stop.
type EFAArrival struct {
	EFADepartureArrival `xml:"itdArrival"`
}

// EFADeparture represents a single departure for a specific stop.
type EFADeparture struct {
	EFADepartureArrival `xml:"itdDeparture"`
}

type departureMonitorResult struct {
	efaResponse
	Odv struct {
		OdvPlace struct {
		}
		OdvName struct {
			State string `xml:"state,attr"`
		} `xml:"itdOdvName"`
	} `xml:"itdDepartureMonitorRequest>itdOdv"`
	Departures []*EFADeparture `xml:"itdDepartureMonitorRequest>itdDepartureList>itdDeparture"`
}

func (d *departureMonitorResult) endpoint() string {
	return "XML_DM_REQUEST"
}

// Departures performs a stateless dm_request for the corresponding stopID and
// returns an array of EFADepartures. Use time.Now() as the second argument in
// order to get the very next departures. The third argument determines how
// many results will be returned by EFA.
func (efa *EFAProvider) Departures(stopID int, due time.Time, results int) ([]*EFADeparture, error) {
	var rt string

	if efa.EnableRealtime {
		rt = "1"
	} else {
		rt = "0"
	}

	params := url.Values{
		"type_dm":              {"any"},
		"name_dm":              {strconv.Itoa(stopID)},
		"locationServerActive": {"1"},
		"useRealtime":          {rt},
		"dmLineSelection":      {"all"}, //FIXME enable line selection
		"limit":                {strconv.Itoa(results)},
		"mode":                 {"direct"},
		"stateless":            {"1"},
		"itdDate":              {due.Format("20060102")},
		"itdTime":              {due.Format("1504")},
	}

	var result departureMonitorResult

	if err := efa.postRequest(&result, params); err != nil {
		return nil, err
	}

	return result.Departures, nil

}
<|endoftext|>


<commit_before>package hostess

import (
	"errors"
	"fmt"
	"net"
	"sort"
	"strings"
)

// ErrInvalidVersionArg is raised when a function expects IPv 4 or 6 but is
// passed a value not 4 or 6.
var ErrInvalidVersionArg = errors.New("Version argument must be 4 or 6")

// Hostlist is a sortable set of Hostnames. When in a Hostlist, Hostnames must
// follow some rules:
//
// 	- Hostlist may contain IPv4 AND IPv6 ("IP version" or "IPv") Hostnames.
// 	- Names are only allowed to overlap if IP version is different.
// 	- Adding a Hostname for an existing name will replace the old one.
//
// The Hostlist uses a deterministic Sort order designed to make a hostfile
// output look a particular way. Generally you don't need to worry about this
// as Sort will be called automatically before Format. However, the Hostlist
// may or may not be sorted at any particular time during runtime.
//
// See the docs and implementation in Sort and Add for more details.
type Hostlist []*Hostname

// NewHostlist initializes a new Hostlist
func NewHostlist() *Hostlist {
	return &Hostlist{}
}

// Len returns the number of Hostnames in the list, part of sort.Interface
func (h Hostlist) Len() int {
	return len(h)
}

// MakeSurrogateIP takes an IP like 127.0.0.1 and munges it to 0.0.0.1 so we can
// sort it more easily. Note that we don't actually want to change the value,
// so we use value copies here (not pointers).
func MakeSurrogateIP(IP net.IP) net.IP {
	if len(IP.String()) > 3 && IP.String()[0:3] == "127" {
		return net.ParseIP("0" + IP.String()[3:])
	}
	return IP
}

// Less determines the sort order of two Hostnames, part of sort.Interface
func (h Hostlist) Less(A, B int) bool {
	// Sort "localhost" at the top
	if h[A].Domain == "localhost" {
		return true
	}
	if h[B].Domain == "localhost" {
		return false
	}

	// Sort IPv4 before IPv6
	// A is IPv4 and B is IPv6. A wins!
	if !h[A].IPv6 && h[B].IPv6 {
		return true
	}
	// A is IPv6 but B is IPv4. A loses!
	if h[A].IPv6 && !h[B].IPv6 {
		return false
	}

	// Compare the the IP addresses (byte array)
	// We want to push 127. to the top so we're going to mark it zero.
	surrogateA := MakeSurrogateIP(h[A].IP)
	surrogateB := MakeSurrogateIP(h[B].IP)
	if !surrogateA.Equal(surrogateB) {
		for charIndex := range surrogateA {
			// A and B's IPs differ at this index, and A is less. A wins!
			if surrogateA[charIndex] < surrogateB[charIndex] {
				return true
			}
			// A and B's IPs differ at this index, and B is less. A loses!
			if surrogateA[charIndex] > surrogateB[charIndex] {
				return false
			}
		}
		// If we got here then the IPs are the same and we want to continue on
		// to the domain sorting section.
	}

	// Prep for sorting by domain name
	aLength := len(h[A].Domain)
	bLength := len(h[B].Domain)
	max := aLength
	if bLength > max {
		max = bLength
	}

	// Sort domains alphabetically
	// TODO: This works best if domains are lowercased. However, we do not
	// enforce lowercase because of UTF-8 domain names, which may be broken by
	// case folding. There is a way to do this correctly but it's complicated
	// so I'm not going to do it right now.
	for charIndex := 0; charIndex < max; charIndex++ {
		// This index is longer than A, so A is shorter. A wins!
		if charIndex >= aLength {
			return true
		}
		// This index is longer than B, so B is shorter. A loses!
		if charIndex >= bLength {
			return false
		}
		// A and B differ at this index and A is less. A wins!
		if h[A].Domain[charIndex] < h[B].Domain[charIndex] {
			return true
		}
		// A and B differ at this index and B is less. A loses!
		if h[A].Domain[charIndex] > h[B].Domain[charIndex] {
			return false
		}
	}

	// If we got here then A and B are the same -- by definition A is not Less
	// than B so we return false. Technically we shouldn't get here since Add
	// should not allow duplicates, but we'll guard anyway.
	return false
}

// Swap changes the position of two Hostnames, part of sort.Interface
func (h Hostlist) Swap(i, j int) {
	h[i], h[j] = h[j], h[i]
}

// Sort this list of Hostnames, according to Hostlist sorting rules:
//
// 	1. localhost comes before other domains
// 	2. IPv4 comes before IPv6
// 	3. IPs are sorted in numerical order
// 	4. domains are sorted in alphabetical
func (h *Hostlist) Sort() {
	sort.Sort(*h)
}

// Contains returns true if this Hostlist has the specified Hostname
func (h *Hostlist) Contains(b *Hostname) bool {
	for _, a := range *h {
		if a.Equal(b) {
			return true
		}
	}
	return false
}

// ContainsDomain returns true if a Hostname in this Hostlist matches domain
func (h *Hostlist) ContainsDomain(domain string) bool {
	for _, hostname := range *h {
		if hostname.Domain == domain {
			return true
		}
	}
	return false
}

// ContainsIP returns true if a Hostname in this Hostlist matches IP
func (h *Hostlist) ContainsIP(IP net.IP) bool {
	for _, hostname := range *h {
		if hostname.EqualIP(IP) {
			return true
		}
	}
	return false
}

// Add a new Hostname to this hostlist. Add uses some merging logic in the
// event it finds duplicated hostnames. In the case of a conflict (incompatible
// entries) the last write wins. In the case of duplicates, duplicates will be
// removed and the remaining entry will be enabled if any of the duplicates was
// enabled.
//
// Both duplicate and conflicts return errors so you are aware of them, but you
// don't necessarily need to do anything about the error.
func (h *Hostlist) Add(hostnamev *Hostname) error {
	hostname := NewHostname(hostnamev.Domain, hostnamev.IP.String(), hostnamev.Enabled)
	for index, found := range *h {
		if found.Equal(hostname) {
			// If either hostname is enabled we will set the existing one to
			// enabled state. That way if we add a hostname from the end of a
			// hosts file it will take over, and if we later add a disabled one
			// the original one will stick. We still error in this case so the
			// user can see that there is a duplicate.
			(*h)[index].Enabled = found.Enabled || hostname.Enabled
			return fmt.Errorf("Duplicate hostname entry for %s -> %s",
				hostname.Domain, hostname.IP)
		} else if found.Domain == hostname.Domain && found.IPv6 == hostname.IPv6 {
			(*h)[index] = hostname
			return fmt.Errorf("Conflicting hostname entries for %s -> %s and -> %s",
				hostname.Domain, hostname.IP, found.IP)
		}
	}
	*h = append(*h, hostname)
	return nil
}

// IndexOf will indicate the index of a Hostname in Hostlist, or -1 if it is
// not found.
func (h *Hostlist) IndexOf(host *Hostname) int {
	for index, found := range *h {
		if found.Equal(host) {
			return index
		}
	}
	return -1
}

// IndexOfDomainV will indicate the index of a Hostname in Hostlist that has
// the same domain and IP version, or -1 if it is not found.
//
// This function will panic if IP version is not 4 or 6.
func (h *Hostlist) IndexOfDomainV(domain string, version int) int {
	if version != 4 && version != 6 {
		panic(ErrInvalidVersionArg)
	}
	for index, hostname := range *h {
		if hostname.Domain == domain && hostname.IPv6 == (version == 6) {
			return index
		}
	}
	return -1
}

// Remove will delete the Hostname at the specified index. If index is out of
// bounds (i.e. -1), Remove silently no-ops.
func (h *Hostlist) Remove(index int) {
	if index > -1 && index < len(*h) {
		*h = append((*h)[:index], (*h)[index+1:]...)
	}
}

// RemoveDomain removes both IPv4 and IPv6 Hostname entries matching domain.
func (h *Hostlist) RemoveDomain(domain string) {
	h.Remove(h.IndexOfDomainV(domain, 4))
	h.Remove(h.IndexOfDomainV(domain, 6))
}

// RemoveDomainV removes a Hostname entry matching the domain and IP version.
//
// This function will panic if IP version is not 4 or 6.
func (h *Hostlist) RemoveDomainV(domain string, version int) {
	if version != 4 && version != 6 {
		panic(ErrInvalidVersionArg)
	}
	h.Remove(h.IndexOfDomainV(domain, version))
}

// Enable will change any Hostnames matching domain to be enabled.
func (h *Hostlist) Enable(domain string) {
	for _, hostname := range *h {
		if hostname.Domain == domain {
			hostname.Enabled = true
		}
	}
}

// EnableV will change a Hostname matching domain and IP version to be enabled.
//
// This function will panic if IP version is not 4 or 6.
func (h *Hostlist) EnableV(domain string, version int) {
	if version != 4 && version != 6 {
		panic(ErrInvalidVersionArg)
	}
	for _, hostname := range *h {
		if hostname.Domain == domain && hostname.IPv6 == (version == 6) {
			hostname.Enabled = true
		}
	}
}

// Disable will change any Hostnames matching domain to be disabled.
func (h *Hostlist) Disable(domain string) {
	for _, hostname := range *h {
		if hostname.Domain == domain {
			hostname.Enabled = false
		}
	}
}

// DisableV will change any Hostnames matching domain and IP version to be disabled.
//
// This function will panic if IP version is not 4 or 6.
func (h *Hostlist) DisableV(domain string, version int) {
	if version != 4 && version != 6 {
		panic(ErrInvalidVersionArg)
	}
	for _, hostname := range *h {
		if hostname.Domain == domain && hostname.IPv6 == (version == 6) {
			hostname.Enabled = false
		}
	}
}

// FilterByIP filters the list of hostnames by IP address.
func (h *Hostlist) FilterByIP(IP net.IP) (hostnames []*Hostname) {
	for _, hostname := range *h {
		if hostname.IP.Equal(IP) {
			hostnames = append(hostnames, hostname)
		}
	}
	return
}

// FilterByDomain filters the list of hostnames by Domain.
func (h *Hostlist) FilterByDomain(domain string) (hostnames []*Hostname) {
	for _, hostname := range *h {
		if hostname.Domain == domain {
			hostnames = append(hostnames, hostname)
		}
	}
	return
}

// FilterByDomainV filters the list of hostnames by domain and IPv4 or IPv6.
// This should never contain more than one item, but returns a list for
// consistency with other filter functions.
func (h *Hostlist) FilterByDomainV(domain string, version int) (hostnames []*Hostname) {
	for _, hostname := range *h {
		if hostname.Domain == domain && hostname.IPv6 == (version == 6) {
			hostnames = append(hostnames, hostname)
		}
	}
	return
}

// GetUniqueIPs extracts an ordered list of unique IPs from the Hostlist.
// This calls Sort() internally.
func (h *Hostlist) GetUniqueIPs() []net.IP {
	h.Sort()
	// A map doesn't preserve order so we're going to use the map to check
	// whether we've seen something and use the list to keep track of the
	// order.
	seen := make(map[string]bool)
	inOrder := []net.IP{}

	for _, hostname := range *h {
		key := (*hostname).IP.String()
		if !seen[key] {
			seen[key] = true
			inOrder = append(inOrder, (*hostname).IP)
		}
	}
	return inOrder
}

// Format takes the current list of Hostnames in this Hostfile and turns it
// into a string suitable for use as an /etc/hosts file.
// Sorting uses the following logic:
//
// 1. List is sorted by IP address
// 2. Commented items are sorted displayed
// 3. 127.* appears at the top of the list (so boot resolvers don't break)
// 4. When present, "localhost" will always appear first in the domain list
func (h *Hostlist) Format() []byte {
	h.Sort()
	out := []byte{}

	// We want to output one line of hostnames per IP, so first we get that
	// list of IPs and iterate.
	for _, IP := range h.GetUniqueIPs() {
		// Technically if an IP has some disabled hostnames we'll show two
		// lines, one starting with a comment (#).
		enabledIPs := []string{}
		disabledIPs := []string{}

		// For this IP, get all hostnames that match and iterate over them.
		for _, hostname := range h.FilterByIP(IP) {
			// If it's enabled, put it in the enabled bucket (likewise for
			// disabled hostnames)
			if hostname.Enabled {
				enabledIPs = append(enabledIPs, hostname.Domain)
			} else {
				disabledIPs = append(disabledIPs, hostname.Domain)
			}
		}

		// Finally, if the bucket contains anything, concatenate it all
		// together and append it to the output. Also add a newline.
		if len(enabledIPs) > 0 {
			concat := fmt.Sprintf("%s %s", IP.String(), strings.Join(enabledIPs, " "))
			out = append(out, []byte(concat)...)
			out = append(out, []byte("\n")...)
		}

		if len(disabledIPs) > 0 {
			concat := fmt.Sprintf("# %s %s", IP.String(), strings.Join(disabledIPs, " "))
			out = append(out, []byte(concat)...)
			out = append(out, []byte("\n")...)
		}
	}

	return out
}
<commit_msg>Remove returns the number of items it acted upon<commit_after>package hostess

import (
	"errors"
	"fmt"
	"net"
	"sort"
	"strings"
)

// ErrInvalidVersionArg is raised when a function expects IPv 4 or 6 but is
// passed a value not 4 or 6.
var ErrInvalidVersionArg = errors.New("Version argument must be 4 or 6")

// Hostlist is a sortable set of Hostnames. When in a Hostlist, Hostnames must
// follow some rules:
//
// 	- Hostlist may contain IPv4 AND IPv6 ("IP version" or "IPv") Hostnames.
// 	- Names are only allowed to overlap if IP version is different.
// 	- Adding a Hostname for an existing name will replace the old one.
//
// The Hostlist uses a deterministic Sort order designed to make a hostfile
// output look a particular way. Generally you don't need to worry about this
// as Sort will be called automatically before Format. However, the Hostlist
// may or may not be sorted at any particular time during runtime.
//
// See the docs and implementation in Sort and Add for more details.
type Hostlist []*Hostname

// NewHostlist initializes a new Hostlist
func NewHostlist() *Hostlist {
	return &Hostlist{}
}

// Len returns the number of Hostnames in the list, part of sort.Interface
func (h Hostlist) Len() int {
	return len(h)
}

// MakeSurrogateIP takes an IP like 127.0.0.1 and munges it to 0.0.0.1 so we can
// sort it more easily. Note that we don't actually want to change the value,
// so we use value copies here (not pointers).
func MakeSurrogateIP(IP net.IP) net.IP {
	if len(IP.String()) > 3 && IP.String()[0:3] == "127" {
		return net.ParseIP("0" + IP.String()[3:])
	}
	return IP
}

// Less determines the sort order of two Hostnames, part of sort.Interface
func (h Hostlist) Less(A, B int) bool {
	// Sort "localhost" at the top
	if h[A].Domain == "localhost" {
		return true
	}
	if h[B].Domain == "localhost" {
		return false
	}

	// Sort IPv4 before IPv6
	// A is IPv4 and B is IPv6. A wins!
	if !h[A].IPv6 && h[B].IPv6 {
		return true
	}
	// A is IPv6 but B is IPv4. A loses!
	if h[A].IPv6 && !h[B].IPv6 {
		return false
	}

	// Compare the the IP addresses (byte array)
	// We want to push 127. to the top so we're going to mark it zero.
	surrogateA := MakeSurrogateIP(h[A].IP)
	surrogateB := MakeSurrogateIP(h[B].IP)
	if !surrogateA.Equal(surrogateB) {
		for charIndex := range surrogateA {
			// A and B's IPs differ at this index, and A is less. A wins!
			if surrogateA[charIndex] < surrogateB[charIndex] {
				return true
			}
			// A and B's IPs differ at this index, and B is less. A loses!
			if surrogateA[charIndex] > surrogateB[charIndex] {
				return false
			}
		}
		// If we got here then the IPs are the same and we want to continue on
		// to the domain sorting section.
	}

	// Prep for sorting by domain name
	aLength := len(h[A].Domain)
	bLength := len(h[B].Domain)
	max := aLength
	if bLength > max {
		max = bLength
	}

	// Sort domains alphabetically
	// TODO: This works best if domains are lowercased. However, we do not
	// enforce lowercase because of UTF-8 domain names, which may be broken by
	// case folding. There is a way to do this correctly but it's complicated
	// so I'm not going to do it right now.
	for charIndex := 0; charIndex < max; charIndex++ {
		// This index is longer than A, so A is shorter. A wins!
		if charIndex >= aLength {
			return true
		}
		// This index is longer than B, so B is shorter. A loses!
		if charIndex >= bLength {
			return false
		}
		// A and B differ at this index and A is less. A wins!
		if h[A].Domain[charIndex] < h[B].Domain[charIndex] {
			return true
		}
		// A and B differ at this index and B is less. A loses!
		if h[A].Domain[charIndex] > h[B].Domain[charIndex] {
			return false
		}
	}

	// If we got here then A and B are the same -- by definition A is not Less
	// than B so we return false. Technically we shouldn't get here since Add
	// should not allow duplicates, but we'll guard anyway.
	return false
}

// Swap changes the position of two Hostnames, part of sort.Interface
func (h Hostlist) Swap(i, j int) {
	h[i], h[j] = h[j], h[i]
}

// Sort this list of Hostnames, according to Hostlist sorting rules:
//
// 	1. localhost comes before other domains
// 	2. IPv4 comes before IPv6
// 	3. IPs are sorted in numerical order
// 	4. domains are sorted in alphabetical
func (h *Hostlist) Sort() {
	sort.Sort(*h)
}

// Contains returns true if this Hostlist has the specified Hostname
func (h *Hostlist) Contains(b *Hostname) bool {
	for _, a := range *h {
		if a.Equal(b) {
			return true
		}
	}
	return false
}

// ContainsDomain returns true if a Hostname in this Hostlist matches domain
func (h *Hostlist) ContainsDomain(domain string) bool {
	for _, hostname := range *h {
		if hostname.Domain == domain {
			return true
		}
	}
	return false
}

// ContainsIP returns true if a Hostname in this Hostlist matches IP
func (h *Hostlist) ContainsIP(IP net.IP) bool {
	for _, hostname := range *h {
		if hostname.EqualIP(IP) {
			return true
		}
	}
	return false
}

// Add a new Hostname to this hostlist. Add uses some merging logic in the
// event it finds duplicated hostnames. In the case of a conflict (incompatible
// entries) the last write wins. In the case of duplicates, duplicates will be
// removed and the remaining entry will be enabled if any of the duplicates was
// enabled.
//
// Both duplicate and conflicts return errors so you are aware of them, but you
// don't necessarily need to do anything about the error.
func (h *Hostlist) Add(hostnamev *Hostname) error {
	hostname := NewHostname(hostnamev.Domain, hostnamev.IP.String(), hostnamev.Enabled)
	for index, found := range *h {
		if found.Equal(hostname) {
			// If either hostname is enabled we will set the existing one to
			// enabled state. That way if we add a hostname from the end of a
			// hosts file it will take over, and if we later add a disabled one
			// the original one will stick. We still error in this case so the
			// user can see that there is a duplicate.
			(*h)[index].Enabled = found.Enabled || hostname.Enabled
			return fmt.Errorf("Duplicate hostname entry for %s -> %s",
				hostname.Domain, hostname.IP)
		} else if found.Domain == hostname.Domain && found.IPv6 == hostname.IPv6 {
			(*h)[index] = hostname
			return fmt.Errorf("Conflicting hostname entries for %s -> %s and -> %s",
				hostname.Domain, hostname.IP, found.IP)
		}
	}
	*h = append(*h, hostname)
	return nil
}

// IndexOf will indicate the index of a Hostname in Hostlist, or -1 if it is
// not found.
func (h *Hostlist) IndexOf(host *Hostname) int {
	for index, found := range *h {
		if found.Equal(host) {
			return index
		}
	}
	return -1
}

// IndexOfDomainV will indicate the index of a Hostname in Hostlist that has
// the same domain and IP version, or -1 if it is not found.
//
// This function will panic if IP version is not 4 or 6.
func (h *Hostlist) IndexOfDomainV(domain string, version int) int {
	if version != 4 && version != 6 {
		panic(ErrInvalidVersionArg)
	}
	for index, hostname := range *h {
		if hostname.Domain == domain && hostname.IPv6 == (version == 6) {
			return index
		}
	}
	return -1
}

// Remove will delete the Hostname at the specified index. If index is out of
// bounds (i.e. -1), Remove silently no-ops. Remove returns the number of items
// removed (0 or 1).
func (h *Hostlist) Remove(index int) int {
	if index > -1 && index < len(*h) {
		*h = append((*h)[:index], (*h)[index+1:]...)
		return 1
	}
	return 0
}

// RemoveDomain removes both IPv4 and IPv6 Hostname entries matching domain.
// Returns the number of entries removed.
func (h *Hostlist) RemoveDomain(domain string) int {
	return h.RemoveDomainV(domain, 4) + h.RemoveDomainV(domain, 6)
}

// RemoveDomainV removes a Hostname entry matching the domain and IP version.
func (h *Hostlist) RemoveDomainV(domain string, version int) int {
	return h.Remove(h.IndexOfDomainV(domain, version))
}

// Enable will change any Hostnames matching domain to be enabled.
func (h *Hostlist) Enable(domain string) {
	for _, hostname := range *h {
		if hostname.Domain == domain {
			hostname.Enabled = true
		}
	}
}

// EnableV will change a Hostname matching domain and IP version to be enabled.
//
// This function will panic if IP version is not 4 or 6.
func (h *Hostlist) EnableV(domain string, version int) {
	if version != 4 && version != 6 {
		panic(ErrInvalidVersionArg)
	}
	for _, hostname := range *h {
		if hostname.Domain == domain && hostname.IPv6 == (version == 6) {
			hostname.Enabled = true
		}
	}
}

// Disable will change any Hostnames matching domain to be disabled.
func (h *Hostlist) Disable(domain string) {
	for _, hostname := range *h {
		if hostname.Domain == domain {
			hostname.Enabled = false
		}
	}
}

// DisableV will change any Hostnames matching domain and IP version to be disabled.
//
// This function will panic if IP version is not 4 or 6.
func (h *Hostlist) DisableV(domain string, version int) {
	if version != 4 && version != 6 {
		panic(ErrInvalidVersionArg)
	}
	for _, hostname := range *h {
		if hostname.Domain == domain && hostname.IPv6 == (version == 6) {
			hostname.Enabled = false
		}
	}
}

// FilterByIP filters the list of hostnames by IP address.
func (h *Hostlist) FilterByIP(IP net.IP) (hostnames []*Hostname) {
	for _, hostname := range *h {
		if hostname.IP.Equal(IP) {
			hostnames = append(hostnames, hostname)
		}
	}
	return
}

// FilterByDomain filters the list of hostnames by Domain.
func (h *Hostlist) FilterByDomain(domain string) (hostnames []*Hostname) {
	for _, hostname := range *h {
		if hostname.Domain == domain {
			hostnames = append(hostnames, hostname)
		}
	}
	return
}

// FilterByDomainV filters the list of hostnames by domain and IPv4 or IPv6.
// This should never contain more than one item, but returns a list for
// consistency with other filter functions.
//
// This function will panic if IP version is not 4 or 6.
func (h *Hostlist) FilterByDomainV(domain string, version int) (hostnames []*Hostname) {
	if version != 4 && version != 6 {
		panic(ErrInvalidVersionArg)
	}
	for _, hostname := range *h {
		if hostname.Domain == domain && hostname.IPv6 == (version == 6) {
			hostnames = append(hostnames, hostname)
		}
	}
	return
}

// GetUniqueIPs extracts an ordered list of unique IPs from the Hostlist.
// This calls Sort() internally.
func (h *Hostlist) GetUniqueIPs() []net.IP {
	h.Sort()
	// A map doesn't preserve order so we're going to use the map to check
	// whether we've seen something and use the list to keep track of the
	// order.
	seen := make(map[string]bool)
	inOrder := []net.IP{}

	for _, hostname := range *h {
		key := (*hostname).IP.String()
		if !seen[key] {
			seen[key] = true
			inOrder = append(inOrder, (*hostname).IP)
		}
	}
	return inOrder
}

// Format takes the current list of Hostnames in this Hostfile and turns it
// into a string suitable for use as an /etc/hosts file.
// Sorting uses the following logic:
//
// 1. List is sorted by IP address
// 2. Commented items are sorted displayed
// 3. 127.* appears at the top of the list (so boot resolvers don't break)
// 4. When present, "localhost" will always appear first in the domain list
func (h *Hostlist) Format() []byte {
	h.Sort()
	out := []byte{}

	// We want to output one line of hostnames per IP, so first we get that
	// list of IPs and iterate.
	for _, IP := range h.GetUniqueIPs() {
		// Technically if an IP has some disabled hostnames we'll show two
		// lines, one starting with a comment (#).
		enabledIPs := []string{}
		disabledIPs := []string{}

		// For this IP, get all hostnames that match and iterate over them.
		for _, hostname := range h.FilterByIP(IP) {
			// If it's enabled, put it in the enabled bucket (likewise for
			// disabled hostnames)
			if hostname.Enabled {
				enabledIPs = append(enabledIPs, hostname.Domain)
			} else {
				disabledIPs = append(disabledIPs, hostname.Domain)
			}
		}

		// Finally, if the bucket contains anything, concatenate it all
		// together and append it to the output. Also add a newline.
		if len(enabledIPs) > 0 {
			concat := fmt.Sprintf("%s %s", IP.String(), strings.Join(enabledIPs, " "))
			out = append(out, []byte(concat)...)
			out = append(out, []byte("\n")...)
		}

		if len(disabledIPs) > 0 {
			concat := fmt.Sprintf("# %s %s", IP.String(), strings.Join(disabledIPs, " "))
			out = append(out, []byte(concat)...)
			out = append(out, []byte("\n")...)
		}
	}

	return out
}
<|endoftext|>


<commit_before>/*
Copyright 2019 kubeflow.org.

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 knative

import (
	"fmt"
	"strconv"

	"sigs.k8s.io/controller-runtime/pkg/client"

	"github.com/kubeflow/kfserving/pkg/apis/serving/v1alpha2"
	"github.com/kubeflow/kfserving/pkg/constants"
	"github.com/kubeflow/kfserving/pkg/controller/inferenceservice/resources/credentials"
	"github.com/kubeflow/kfserving/pkg/utils"
	v1 "k8s.io/api/core/v1"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	"knative.dev/serving/pkg/apis/autoscaling"
	"knative.dev/serving/pkg/apis/serving"
	knservingv1 "knative.dev/serving/pkg/apis/serving/v1"
)

var serviceAnnotationDisallowedList = []string{
	autoscaling.MinScaleAnnotationKey,
	autoscaling.MaxScaleAnnotationKey,
	constants.StorageInitializerSourceUriInternalAnnotationKey,
	"kubectl.kubernetes.io/last-applied-configuration",
}

const (
	// Set to 20% of the resource for main container, InferenceService defaults to 1CPU which is 200m for queue-proxy
	// https://github.com/knative/serving/blob/1d263950f9f2fea85a4dd394948a029c328af9d9/pkg/reconciler/revision/resources/resourceboundary.go#L30
	DefaultQueueSideCarResourcePercentage = "20"
)

type ServiceBuilder struct {
	inferenceServiceConfig *v1alpha2.InferenceServicesConfig
	credentialBuilder      *credentials.CredentialBuilder
}

func NewServiceBuilder(client client.Client, config *v1.ConfigMap) *ServiceBuilder {
	inferenceServiceConfig, err := v1alpha2.NewInferenceServicesConfig(config)
	if err != nil {
		fmt.Printf("Failed to get inference service config %s", err)
		panic("Failed to get inference service config")

	}
	return &ServiceBuilder{
		inferenceServiceConfig: inferenceServiceConfig,
		credentialBuilder:      credentials.NewCredentialBulder(client, config),
	}
}

func (c *ServiceBuilder) CreateInferenceServiceComponent(isvc *v1alpha2.InferenceService, component constants.InferenceServiceComponent, isCanary bool) (*knservingv1.Service, error) {
	serviceName := constants.DefaultServiceName(isvc.Name, component)
	if isCanary {
		serviceName = constants.CanaryServiceName(isvc.Name, component)
	}
	switch component {
	case constants.Predictor:
		predictorSpec := &isvc.Spec.Default.Predictor
		if isCanary {
			predictorSpec = &isvc.Spec.Canary.Predictor
		}
		return c.CreatePredictorService(serviceName, isvc.ObjectMeta, predictorSpec, isCanary)
	case constants.Transformer:
		transformerSpec := isvc.Spec.Default.Transformer
		if isCanary {
			transformerSpec = isvc.Spec.Canary.Transformer
		}
		if transformerSpec == nil {
			return nil, nil
		}
		return c.CreateTransformerService(serviceName, isvc.ObjectMeta, transformerSpec, isCanary)
	case constants.Explainer:
		explainerSpec := isvc.Spec.Default.Explainer
		predictorService := constants.PredictorURL(isvc.ObjectMeta, isCanary)
		if isvc.Spec.Default.Transformer != nil {
			predictorService = constants.TransformerURL(isvc.ObjectMeta, isCanary)
		}
		if explainerSpec == nil {
			return nil, nil
		}
		return c.CreateExplainerService(serviceName, isvc.ObjectMeta, explainerSpec, predictorService)
	}
	return nil, fmt.Errorf("Invalid Component")
}

func addLoggerAnnotations(logger *v1alpha2.Logger, annotations map[string]string) bool {
	if logger != nil {
		annotations[constants.LoggerInternalAnnotationKey] = "true"
		if logger.Url != nil {
			annotations[constants.LoggerSinkUrlInternalAnnotationKey] = *logger.Url
		}
		annotations[constants.LoggerModeInternalAnnotationKey] = string(logger.Mode)
		return true
	}
	return false
}

func addLoggerContainerPort(container *v1.Container) {
	if container != nil {
		if container.Ports == nil {
			port, _ := strconv.Atoi(constants.InferenceServiceDefaultLoggerPort)
			container.Ports = []v1.ContainerPort{
				v1.ContainerPort{
					ContainerPort: int32(port),
				},
			}
		}
	}
}

func (c *ServiceBuilder) CreatePredictorService(name string, metadata metav1.ObjectMeta, predictorSpec *v1alpha2.PredictorSpec, isCanary bool) (*knservingv1.Service, error) {
	annotations, err := c.buildAnnotations(metadata, predictorSpec.MinReplicas, predictorSpec.MaxReplicas, predictorSpec.Parallelism)
	if err != nil {
		return nil, err
	}

	// KNative does not support INIT containers or mounting, so we add annotations that trigger the
	// StorageInitializer injector to mutate the underlying deployment to provision model data
	if sourceURI := predictorSpec.GetStorageUri(); sourceURI != "" {
		annotations[constants.StorageInitializerSourceUriInternalAnnotationKey] = sourceURI
	}

	// Knative does not support multiple containers so we add an annotation that triggers pod
	// mutator to add it
	hasInferenceLogging := addLoggerAnnotations(predictorSpec.Logger, annotations)
	container := predictorSpec.GetContainer(metadata.Name, predictorSpec.Parallelism, c.inferenceServiceConfig)
	if hasInferenceLogging {
		addLoggerContainerPort(container)
	}

	endpoint := constants.InferenceServiceDefault
	if isCanary {
		endpoint = constants.InferenceServiceCanary
	}

	service := &knservingv1.Service{
		ObjectMeta: metav1.ObjectMeta{
			Name:      name,
			Namespace: metadata.Namespace,
			Labels:    metadata.Labels,
		},
		Spec: knservingv1.ServiceSpec{
			ConfigurationSpec: knservingv1.ConfigurationSpec{
				Template: knservingv1.RevisionTemplateSpec{
					ObjectMeta: metav1.ObjectMeta{
						Labels: utils.Union(metadata.Labels, map[string]string{
							constants.InferenceServicePodLabelKey: metadata.Name,
							constants.KServiceComponentLabel:      constants.Predictor.String(),
							constants.KServiceModelLabel:          metadata.Name,
							constants.KServiceEndpointLabel:       endpoint,
						}),
						Annotations: annotations,
					},
					Spec: knservingv1.RevisionSpec{
						// Defaulting here since this always shows a diff with nil vs 300s(knative default)
						// we may need to expose this field in future
						TimeoutSeconds: &constants.DefaultPredictorTimeout,
						PodSpec: v1.PodSpec{
							ServiceAccountName: predictorSpec.ServiceAccountName,
							Containers: []v1.Container{
								*container,
							},
						},
					},
				},
			},
		},
	}

	if err := c.credentialBuilder.CreateSecretVolumeAndEnv(
		metadata.Namespace,
		predictorSpec.ServiceAccountName,
		&service.Spec.Template.Spec.Containers[0],
		&service.Spec.Template.Spec.Volumes,
	); err != nil {
		return nil, err
	}

	return service, nil
}

func (c *ServiceBuilder) CreateTransformerService(name string, metadata metav1.ObjectMeta, transformerSpec *v1alpha2.TransformerSpec, isCanary bool) (*knservingv1.Service, error) {
	annotations, err := c.buildAnnotations(metadata, transformerSpec.MinReplicas, transformerSpec.MaxReplicas, transformerSpec.Parallelism)
	if err != nil {
		return nil, err
	}

	// KNative does not support INIT containers or mounting, so we add annotations that trigger the
	// StorageInitializer injector to mutate the underlying deployment to provision model data
	if sourceURI := transformerSpec.GetStorageUri(); sourceURI != "" {
		annotations[constants.StorageInitializerSourceUriInternalAnnotationKey] = sourceURI
	}
	// Knative does not support multiple containers so we add an annotation that triggers pod
	// mutator to add it
	hasInferenceLogging := addLoggerAnnotations(transformerSpec.Logger, annotations)
	container := transformerSpec.GetContainerSpec(metadata, isCanary)
	if hasInferenceLogging {
		addLoggerContainerPort(container)
	}

	endpoint := constants.InferenceServiceDefault
	if isCanary {
		endpoint = constants.InferenceServiceCanary
	}

	service := &knservingv1.Service{
		ObjectMeta: metav1.ObjectMeta{
			Name:      name,
			Namespace: metadata.Namespace,
			Labels:    metadata.Labels,
		},
		Spec: knservingv1.ServiceSpec{
			ConfigurationSpec: knservingv1.ConfigurationSpec{
				Template: knservingv1.RevisionTemplateSpec{
					ObjectMeta: metav1.ObjectMeta{
						Labels: utils.Union(metadata.Labels, map[string]string{
							constants.InferenceServicePodLabelKey: metadata.Name,
							constants.KServiceComponentLabel:      constants.Transformer.String(),
							constants.KServiceModelLabel:          metadata.Name,
							constants.KServiceEndpointLabel:       endpoint,
						}),
						Annotations: annotations,
					},
					Spec: knservingv1.RevisionSpec{
						// Defaulting here since this always shows a diff with nil vs 300s(knative default)
						// we may need to expose this field in future
						TimeoutSeconds: &constants.DefaultTransformerTimeout,
						PodSpec: v1.PodSpec{
							ServiceAccountName: transformerSpec.ServiceAccountName,
							Containers: []v1.Container{
								*container,
							},
						},
					},
				},
			},
		},
	}

	if err := c.credentialBuilder.CreateSecretVolumeAndEnv(
		metadata.Namespace,
		transformerSpec.ServiceAccountName,
		&service.Spec.Template.Spec.Containers[0],
		&service.Spec.Template.Spec.Volumes,
	); err != nil {
		return nil, err
	}

	return service, nil
}

func (c *ServiceBuilder) CreateExplainerService(name string, metadata metav1.ObjectMeta, explainerSpec *v1alpha2.ExplainerSpec, predictorService string) (*knservingv1.Service, error) {
	annotations, err := c.buildAnnotations(metadata, explainerSpec.MinReplicas, explainerSpec.MaxReplicas, explainerSpec.Parallelism)
	if err != nil {
		return nil, err
	}

	// KNative does not support INIT containers or mounting, so we add annotations that trigger the
	// ModelInitializer injector to mutate the underlying deployment to provision model data
	if sourceURI := explainerSpec.GetStorageUri(); sourceURI != "" {
		annotations[constants.StorageInitializerSourceUriInternalAnnotationKey] = sourceURI
	}

	// Knative does not support multiple containers so we add an annotation that triggers pod
	// mutator to add it
	hasInferenceLogging := addLoggerAnnotations(explainerSpec.Logger, annotations)
	container := explainerSpec.CreateExplainerContainer(metadata.Name, explainerSpec.Parallelism, predictorService, c.inferenceServiceConfig)
	if hasInferenceLogging {
		addLoggerContainerPort(container)
	}

	service := &knservingv1.Service{
		ObjectMeta: metav1.ObjectMeta{
			Name:      name,
			Namespace: metadata.Namespace,
			Labels:    metadata.Labels,
		},
		Spec: knservingv1.ServiceSpec{
			ConfigurationSpec: knservingv1.ConfigurationSpec{
				Template: knservingv1.RevisionTemplateSpec{
					ObjectMeta: metav1.ObjectMeta{
						Labels: utils.Union(metadata.Labels, map[string]string{
							constants.InferenceServicePodLabelKey: metadata.Name,
							constants.KServiceComponentLabel:      constants.Explainer.String(),
							constants.KServiceModelLabel:          metadata.Name,
						}),
						Annotations: annotations,
					},
					Spec: knservingv1.RevisionSpec{
						// Defaulting here since this always shows a diff with nil vs 300s(knative default)
						// we may need to expose this field in future
						TimeoutSeconds: &constants.DefaultExplainerTimeout,
						PodSpec: v1.PodSpec{
							ServiceAccountName: explainerSpec.ServiceAccountName,
							Containers: []v1.Container{
								*container,
							},
						},
					},
				},
			},
		},
	}

	if err := c.credentialBuilder.CreateSecretVolumeAndEnv(
		metadata.Namespace,
		explainerSpec.ServiceAccountName,
		&service.Spec.Template.Spec.Containers[0],
		&service.Spec.Template.Spec.Volumes,
	); err != nil {
		return nil, err
	}

	return service, nil
}

func (c *ServiceBuilder) buildAnnotations(metadata metav1.ObjectMeta, minReplicas *int, maxReplicas int, parallelism int) (map[string]string, error) {
	annotations := utils.Filter(metadata.Annotations, func(key string) bool {
		return !utils.Includes(serviceAnnotationDisallowedList, key)
	})

	if minReplicas == nil {
		annotations[autoscaling.MinScaleAnnotationKey] = fmt.Sprint(constants.DefaultMinReplicas)
	} else if *minReplicas != 0 {
		annotations[autoscaling.MinScaleAnnotationKey] = fmt.Sprint(*minReplicas)
	}

	if maxReplicas != 0 {
		annotations[autoscaling.MaxScaleAnnotationKey] = fmt.Sprint(maxReplicas)
	}

	if _, ok := annotations[serving.QueueSideCarResourcePercentageAnnotation]; !ok {
		annotations[serving.QueueSideCarResourcePercentageAnnotation] = DefaultQueueSideCarResourcePercentage
	}
	// User can pass down scaling target annotation to overwrite the target default 1
	if _, ok := annotations[autoscaling.TargetAnnotationKey]; !ok {
		if parallelism == 0 {
			annotations[autoscaling.TargetAnnotationKey] = constants.DefaultScalingTarget
		} else {
			annotations[autoscaling.TargetAnnotationKey] = strconv.Itoa(parallelism)
		}
	}
	// User can pass down scaling class annotation to overwrite the default scaling KPA
	if _, ok := annotations[autoscaling.ClassAnnotationKey]; !ok {
		annotations[autoscaling.ClassAnnotationKey] = autoscaling.KPA
	}
	return annotations, nil
}
<commit_msg>Set ContainerConcurrency to Parallelism (#806)<commit_after>/*
Copyright 2019 kubeflow.org.

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 knative

import (
	"fmt"
	"strconv"

	"sigs.k8s.io/controller-runtime/pkg/client"

	"github.com/kubeflow/kfserving/pkg/apis/serving/v1alpha2"
	"github.com/kubeflow/kfserving/pkg/constants"
	"github.com/kubeflow/kfserving/pkg/controller/inferenceservice/resources/credentials"
	"github.com/kubeflow/kfserving/pkg/utils"
	v1 "k8s.io/api/core/v1"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	"knative.dev/serving/pkg/apis/autoscaling"
	"knative.dev/serving/pkg/apis/serving"
	knservingv1 "knative.dev/serving/pkg/apis/serving/v1"
)

var serviceAnnotationDisallowedList = []string{
	autoscaling.MinScaleAnnotationKey,
	autoscaling.MaxScaleAnnotationKey,
	constants.StorageInitializerSourceUriInternalAnnotationKey,
	"kubectl.kubernetes.io/last-applied-configuration",
}

const (
	// Set to 20% of the resource for main container, InferenceService defaults to 1CPU which is 200m for queue-proxy
	// https://github.com/knative/serving/blob/1d263950f9f2fea85a4dd394948a029c328af9d9/pkg/reconciler/revision/resources/resourceboundary.go#L30
	DefaultQueueSideCarResourcePercentage = "20"
)

type ServiceBuilder struct {
	inferenceServiceConfig *v1alpha2.InferenceServicesConfig
	credentialBuilder      *credentials.CredentialBuilder
}

func NewServiceBuilder(client client.Client, config *v1.ConfigMap) *ServiceBuilder {
	inferenceServiceConfig, err := v1alpha2.NewInferenceServicesConfig(config)
	if err != nil {
		fmt.Printf("Failed to get inference service config %s", err)
		panic("Failed to get inference service config")

	}
	return &ServiceBuilder{
		inferenceServiceConfig: inferenceServiceConfig,
		credentialBuilder:      credentials.NewCredentialBulder(client, config),
	}
}

func (c *ServiceBuilder) CreateInferenceServiceComponent(isvc *v1alpha2.InferenceService, component constants.InferenceServiceComponent, isCanary bool) (*knservingv1.Service, error) {
	serviceName := constants.DefaultServiceName(isvc.Name, component)
	if isCanary {
		serviceName = constants.CanaryServiceName(isvc.Name, component)
	}
	switch component {
	case constants.Predictor:
		predictorSpec := &isvc.Spec.Default.Predictor
		if isCanary {
			predictorSpec = &isvc.Spec.Canary.Predictor
		}
		return c.CreatePredictorService(serviceName, isvc.ObjectMeta, predictorSpec, isCanary)
	case constants.Transformer:
		transformerSpec := isvc.Spec.Default.Transformer
		if isCanary {
			transformerSpec = isvc.Spec.Canary.Transformer
		}
		if transformerSpec == nil {
			return nil, nil
		}
		return c.CreateTransformerService(serviceName, isvc.ObjectMeta, transformerSpec, isCanary)
	case constants.Explainer:
		explainerSpec := isvc.Spec.Default.Explainer
		predictorService := constants.PredictorURL(isvc.ObjectMeta, isCanary)
		if isvc.Spec.Default.Transformer != nil {
			predictorService = constants.TransformerURL(isvc.ObjectMeta, isCanary)
		}
		if explainerSpec == nil {
			return nil, nil
		}
		return c.CreateExplainerService(serviceName, isvc.ObjectMeta, explainerSpec, predictorService)
	}
	return nil, fmt.Errorf("Invalid Component")
}

func addLoggerAnnotations(logger *v1alpha2.Logger, annotations map[string]string) bool {
	if logger != nil {
		annotations[constants.LoggerInternalAnnotationKey] = "true"
		if logger.Url != nil {
			annotations[constants.LoggerSinkUrlInternalAnnotationKey] = *logger.Url
		}
		annotations[constants.LoggerModeInternalAnnotationKey] = string(logger.Mode)
		return true
	}
	return false
}

func addLoggerContainerPort(container *v1.Container) {
	if container != nil {
		if container.Ports == nil {
			port, _ := strconv.Atoi(constants.InferenceServiceDefaultLoggerPort)
			container.Ports = []v1.ContainerPort{
				v1.ContainerPort{
					ContainerPort: int32(port),
				},
			}
		}
	}
}

func (c *ServiceBuilder) CreatePredictorService(name string, metadata metav1.ObjectMeta, predictorSpec *v1alpha2.PredictorSpec, isCanary bool) (*knservingv1.Service, error) {
	annotations, err := c.buildAnnotations(metadata, predictorSpec.MinReplicas, predictorSpec.MaxReplicas, predictorSpec.Parallelism)
	if err != nil {
		return nil, err
	}

	// KNative does not support INIT containers or mounting, so we add annotations that trigger the
	// StorageInitializer injector to mutate the underlying deployment to provision model data
	if sourceURI := predictorSpec.GetStorageUri(); sourceURI != "" {
		annotations[constants.StorageInitializerSourceUriInternalAnnotationKey] = sourceURI
	}

	// Knative does not support multiple containers so we add an annotation that triggers pod
	// mutator to add it
	hasInferenceLogging := addLoggerAnnotations(predictorSpec.Logger, annotations)
	container := predictorSpec.GetContainer(metadata.Name, predictorSpec.Parallelism, c.inferenceServiceConfig)
	if hasInferenceLogging {
		addLoggerContainerPort(container)
	}

	endpoint := constants.InferenceServiceDefault
	if isCanary {
		endpoint = constants.InferenceServiceCanary
	}
	concurrency := int64(predictorSpec.Parallelism)
	service := &knservingv1.Service{
		ObjectMeta: metav1.ObjectMeta{
			Name:      name,
			Namespace: metadata.Namespace,
			Labels:    metadata.Labels,
		},
		Spec: knservingv1.ServiceSpec{
			ConfigurationSpec: knservingv1.ConfigurationSpec{
				Template: knservingv1.RevisionTemplateSpec{
					ObjectMeta: metav1.ObjectMeta{
						Labels: utils.Union(metadata.Labels, map[string]string{
							constants.InferenceServicePodLabelKey: metadata.Name,
							constants.KServiceComponentLabel:      constants.Predictor.String(),
							constants.KServiceModelLabel:          metadata.Name,
							constants.KServiceEndpointLabel:       endpoint,
						}),
						Annotations: annotations,
					},
					Spec: knservingv1.RevisionSpec{
						// Defaulting here since this always shows a diff with nil vs 300s(knative default)
						// we may need to expose this field in future
						TimeoutSeconds:       &constants.DefaultPredictorTimeout,
						ContainerConcurrency: &concurrency,
						PodSpec: v1.PodSpec{
							ServiceAccountName: predictorSpec.ServiceAccountName,
							Containers: []v1.Container{
								*container,
							},
						},
					},
				},
			},
		},
	}

	if err := c.credentialBuilder.CreateSecretVolumeAndEnv(
		metadata.Namespace,
		predictorSpec.ServiceAccountName,
		&service.Spec.Template.Spec.Containers[0],
		&service.Spec.Template.Spec.Volumes,
	); err != nil {
		return nil, err
	}

	return service, nil
}

func (c *ServiceBuilder) CreateTransformerService(name string, metadata metav1.ObjectMeta, transformerSpec *v1alpha2.TransformerSpec, isCanary bool) (*knservingv1.Service, error) {
	annotations, err := c.buildAnnotations(metadata, transformerSpec.MinReplicas, transformerSpec.MaxReplicas, transformerSpec.Parallelism)
	if err != nil {
		return nil, err
	}

	// KNative does not support INIT containers or mounting, so we add annotations that trigger the
	// StorageInitializer injector to mutate the underlying deployment to provision model data
	if sourceURI := transformerSpec.GetStorageUri(); sourceURI != "" {
		annotations[constants.StorageInitializerSourceUriInternalAnnotationKey] = sourceURI
	}
	// Knative does not support multiple containers so we add an annotation that triggers pod
	// mutator to add it
	hasInferenceLogging := addLoggerAnnotations(transformerSpec.Logger, annotations)
	container := transformerSpec.GetContainerSpec(metadata, isCanary)
	if hasInferenceLogging {
		addLoggerContainerPort(container)
	}

	endpoint := constants.InferenceServiceDefault
	if isCanary {
		endpoint = constants.InferenceServiceCanary
	}

	concurrency := int64(transformerSpec.Parallelism)
	service := &knservingv1.Service{
		ObjectMeta: metav1.ObjectMeta{
			Name:      name,
			Namespace: metadata.Namespace,
			Labels:    metadata.Labels,
		},
		Spec: knservingv1.ServiceSpec{
			ConfigurationSpec: knservingv1.ConfigurationSpec{
				Template: knservingv1.RevisionTemplateSpec{
					ObjectMeta: metav1.ObjectMeta{
						Labels: utils.Union(metadata.Labels, map[string]string{
							constants.InferenceServicePodLabelKey: metadata.Name,
							constants.KServiceComponentLabel:      constants.Transformer.String(),
							constants.KServiceModelLabel:          metadata.Name,
							constants.KServiceEndpointLabel:       endpoint,
						}),
						Annotations: annotations,
					},
					Spec: knservingv1.RevisionSpec{
						// Defaulting here since this always shows a diff with nil vs 300s(knative default)
						// we may need to expose this field in future
						TimeoutSeconds:       &constants.DefaultTransformerTimeout,
						ContainerConcurrency: &concurrency,
						PodSpec: v1.PodSpec{
							ServiceAccountName: transformerSpec.ServiceAccountName,
							Containers: []v1.Container{
								*container,
							},
						},
					},
				},
			},
		},
	}

	if err := c.credentialBuilder.CreateSecretVolumeAndEnv(
		metadata.Namespace,
		transformerSpec.ServiceAccountName,
		&service.Spec.Template.Spec.Containers[0],
		&service.Spec.Template.Spec.Volumes,
	); err != nil {
		return nil, err
	}

	return service, nil
}

func (c *ServiceBuilder) CreateExplainerService(name string, metadata metav1.ObjectMeta, explainerSpec *v1alpha2.ExplainerSpec, predictorService string) (*knservingv1.Service, error) {
	annotations, err := c.buildAnnotations(metadata, explainerSpec.MinReplicas, explainerSpec.MaxReplicas, explainerSpec.Parallelism)
	if err != nil {
		return nil, err
	}

	// KNative does not support INIT containers or mounting, so we add annotations that trigger the
	// ModelInitializer injector to mutate the underlying deployment to provision model data
	if sourceURI := explainerSpec.GetStorageUri(); sourceURI != "" {
		annotations[constants.StorageInitializerSourceUriInternalAnnotationKey] = sourceURI
	}

	// Knative does not support multiple containers so we add an annotation that triggers pod
	// mutator to add it
	hasInferenceLogging := addLoggerAnnotations(explainerSpec.Logger, annotations)
	container := explainerSpec.CreateExplainerContainer(metadata.Name, explainerSpec.Parallelism, predictorService, c.inferenceServiceConfig)
	if hasInferenceLogging {
		addLoggerContainerPort(container)
	}

	concurrency := int64(explainerSpec.Parallelism)
	service := &knservingv1.Service{
		ObjectMeta: metav1.ObjectMeta{
			Name:      name,
			Namespace: metadata.Namespace,
			Labels:    metadata.Labels,
		},
		Spec: knservingv1.ServiceSpec{
			ConfigurationSpec: knservingv1.ConfigurationSpec{
				Template: knservingv1.RevisionTemplateSpec{
					ObjectMeta: metav1.ObjectMeta{
						Labels: utils.Union(metadata.Labels, map[string]string{
							constants.InferenceServicePodLabelKey: metadata.Name,
							constants.KServiceComponentLabel:      constants.Explainer.String(),
							constants.KServiceModelLabel:          metadata.Name,
						}),
						Annotations: annotations,
					},
					Spec: knservingv1.RevisionSpec{
						// Defaulting here since this always shows a diff with nil vs 300s(knative default)
						// we may need to expose this field in future
						TimeoutSeconds: &constants.DefaultExplainerTimeout,
						ContainerConcurrency: &concurrency,
						PodSpec: v1.PodSpec{
							ServiceAccountName: explainerSpec.ServiceAccountName,
							Containers: []v1.Container{
								*container,
							},
						},
					},
				},
			},
		},
	}

	if err := c.credentialBuilder.CreateSecretVolumeAndEnv(
		metadata.Namespace,
		explainerSpec.ServiceAccountName,
		&service.Spec.Template.Spec.Containers[0],
		&service.Spec.Template.Spec.Volumes,
	); err != nil {
		return nil, err
	}

	return service, nil
}

func (c *ServiceBuilder) buildAnnotations(metadata metav1.ObjectMeta, minReplicas *int, maxReplicas int, parallelism int) (map[string]string, error) {
	annotations := utils.Filter(metadata.Annotations, func(key string) bool {
		return !utils.Includes(serviceAnnotationDisallowedList, key)
	})

	if minReplicas == nil {
		annotations[autoscaling.MinScaleAnnotationKey] = fmt.Sprint(constants.DefaultMinReplicas)
	} else if *minReplicas != 0 {
		annotations[autoscaling.MinScaleAnnotationKey] = fmt.Sprint(*minReplicas)
	}

	if maxReplicas != 0 {
		annotations[autoscaling.MaxScaleAnnotationKey] = fmt.Sprint(maxReplicas)
	}

	if _, ok := annotations[serving.QueueSideCarResourcePercentageAnnotation]; !ok {
		annotations[serving.QueueSideCarResourcePercentageAnnotation] = DefaultQueueSideCarResourcePercentage
	}
	// User can pass down scaling target annotation to overwrite the target default 1
	if _, ok := annotations[autoscaling.TargetAnnotationKey]; !ok {
		if parallelism == 0 {
			annotations[autoscaling.TargetAnnotationKey] = constants.DefaultScalingTarget
		} else {
			annotations[autoscaling.TargetAnnotationKey] = strconv.Itoa(parallelism)
		}
	}
	// User can pass down scaling class annotation to overwrite the default scaling KPA
	if _, ok := annotations[autoscaling.ClassAnnotationKey]; !ok {
		annotations[autoscaling.ClassAnnotationKey] = autoscaling.KPA
	}
	return annotations, nil
}
<|endoftext|>


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

package clnt

import (
	"net"
	"syscall"
	"go9p.googlecode.com/hg/p"
)

// Creates an authentication fid for the specified user. Returns the fid, if
// successful, or an Error.
func (clnt *Clnt) Auth(user p.User, aname string) (*Fid, *p.Error) {
	fid := clnt.FidAlloc()
	tc := clnt.NewFcall()
	err := p.PackTauth(tc, fid.Fid, user.Name(), aname, uint32(user.Id()), clnt.Dotu)
	if err != nil {
		return nil, err
	}

	_, err = clnt.Rpc(tc)
	if err != nil {
		return nil, err
	}

	fid.User = user
	fid.walked = true
	return fid, nil
}

// Creates a fid for the specified user that points to the root
// of the file server's file tree. Returns a Fid pointing to the root,
// if successful, or an Error.
func (clnt *Clnt) Attach(afid *Fid, user p.User, aname string) (*Fid, *p.Error) {
	var afno uint32

	if afid != nil {
		afno = afid.Fid
	} else {
		afno = p.NOFID
	}

	fid := clnt.FidAlloc()
	tc := clnt.NewFcall()
	err := p.PackTattach(tc, fid.Fid, afno, user.Name(), aname, uint32(user.Id()), clnt.Dotu)
	if err != nil {
		return nil, err
	}

	rc, err := clnt.Rpc(tc)
	if err != nil {
		return nil, err
	}
	if rc.Type == p.Rerror {
		return nil, &p.Error{rc.Error, int(rc.Errornum)}
	}

	fid.Qid = rc.Qid
	fid.User = user
	fid.walked = true
	return fid, nil
}

// Connects to a file server and attaches to it as the specified user.
func Mount(ntype, addr, aname string, user p.User) (*Clnt, *p.Error) {
	c, e := net.Dial(ntype, "", addr)
	if e != nil {
		return nil, &p.Error{e.String(), syscall.EIO}
	}

	return MountConn(c, aname, user)
}

func MountConn(c net.Conn, aname string, user p.User) (*Clnt, *p.Error) {
	clnt, err := Connect(c, 8192+p.IOHDRSZ, true)
	if err != nil {
		return nil, err
	}

	fid, err := clnt.Attach(nil, user, aname)
	if err != nil {
		clnt.Unmount()
		return nil, err
	}

	clnt.Root = fid
	return clnt, nil
}

// Closes the connection to the file sever.
func (clnt *Clnt) Unmount() {
	clnt.Lock()
	clnt.err = &p.Error{"connection closed", syscall.ECONNRESET}
	clnt.conn.Close()
	clnt.Unlock()
}
<commit_msg>gofix for netdial<commit_after>// Copyright 2009 The Go Authors.  All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package clnt

import (
	"net"
	"syscall"
	"go9p.googlecode.com/hg/p"
)

// Creates an authentication fid for the specified user. Returns the fid, if
// successful, or an Error.
func (clnt *Clnt) Auth(user p.User, aname string) (*Fid, *p.Error) {
	fid := clnt.FidAlloc()
	tc := clnt.NewFcall()
	err := p.PackTauth(tc, fid.Fid, user.Name(), aname, uint32(user.Id()), clnt.Dotu)
	if err != nil {
		return nil, err
	}

	_, err = clnt.Rpc(tc)
	if err != nil {
		return nil, err
	}

	fid.User = user
	fid.walked = true
	return fid, nil
}

// Creates a fid for the specified user that points to the root
// of the file server's file tree. Returns a Fid pointing to the root,
// if successful, or an Error.
func (clnt *Clnt) Attach(afid *Fid, user p.User, aname string) (*Fid, *p.Error) {
	var afno uint32

	if afid != nil {
		afno = afid.Fid
	} else {
		afno = p.NOFID
	}

	fid := clnt.FidAlloc()
	tc := clnt.NewFcall()
	err := p.PackTattach(tc, fid.Fid, afno, user.Name(), aname, uint32(user.Id()), clnt.Dotu)
	if err != nil {
		return nil, err
	}

	rc, err := clnt.Rpc(tc)
	if err != nil {
		return nil, err
	}
	if rc.Type == p.Rerror {
		return nil, &p.Error{rc.Error, int(rc.Errornum)}
	}

	fid.Qid = rc.Qid
	fid.User = user
	fid.walked = true
	return fid, nil
}

// Connects to a file server and attaches to it as the specified user.
func Mount(ntype, addr, aname string, user p.User) (*Clnt, *p.Error) {
	c, e := net.Dial(ntype, addr)
	if e != nil {
		return nil, &p.Error{e.String(), syscall.EIO}
	}

	return MountConn(c, aname, user)
}

func MountConn(c net.Conn, aname string, user p.User) (*Clnt, *p.Error) {
	clnt, err := Connect(c, 8192+p.IOHDRSZ, true)
	if err != nil {
		return nil, err
	}

	fid, err := clnt.Attach(nil, user, aname)
	if err != nil {
		clnt.Unmount()
		return nil, err
	}

	clnt.Root = fid
	return clnt, nil
}

// Closes the connection to the file sever.
func (clnt *Clnt) Unmount() {
	clnt.Lock()
	clnt.err = &p.Error{"connection closed", syscall.ECONNRESET}
	clnt.conn.Close()
	clnt.Unlock()
}
<|endoftext|>


<commit_before>package main

import (
	"compress/gzip"
	"encoding/json"
	"fmt"
	"io"
	"io/ioutil"
	"log"
	"math/rand"
	"net/http"
	"net/url"
	"os"
	"path/filepath"
	"strconv"
	"strings"
	"time"

	"github.com/couchbaselabs/cbfs/config"
	"github.com/dustin/gomemcached"
	"github.com/dustin/gomemcached/client"
)

func doGetConfig(w http.ResponseWriter, req *http.Request) {
	err := updateConfig()
	if err != nil && !gomemcached.IsNotFound(err) {
		log.Printf("Error updating config: %v", err)
	}

	w.Header().Set("Content-Type", "application/json")
	w.WriteHeader(200)

	e := json.NewEncoder(w)
	err = e.Encode(&globalConfig)
	if err != nil {
		log.Printf("Error sending config: %v", err)
	}
}

func putConfig(w http.ResponseWriter, req *http.Request) {
	d := json.NewDecoder(req.Body)
	conf := cbfsconfig.CBFSConfig{}

	err := d.Decode(&conf)
	if err != nil {
		w.WriteHeader(500)
		fmt.Fprintf(w, "Error reading config: %v", err)
		return
	}

	err = StoreConfig(conf)
	if err != nil {
		w.WriteHeader(500)
		fmt.Fprintf(w, "Error writing config: %v", err)
		return
	}

	err = updateConfig()
	if err != nil {
		log.Printf("Error fetching newly stored config: %v", err)
	}

	w.WriteHeader(204)
}

func doList(w http.ResponseWriter, req *http.Request) {
	w.WriteHeader(200)
	explen := getHash().Size() * 2
	filepath.Walk(*root, func(path string, info os.FileInfo, err error) error {
		if err != nil {
			return err
		}
		if !info.IsDir() && !strings.HasPrefix(info.Name(), "tmp") &&
			len(info.Name()) == explen {
			_, e := w.Write([]byte(info.Name() + "\n"))
			return e
		}
		return nil
	})
}

func doListTasks(w http.ResponseWriter, req *http.Request) {
	tasks, err := listRunningTasks()
	if err != nil {
		w.WriteHeader(500)
		fmt.Fprintf(w, "Error listing tasks:  %v", err)
		return
	}

	w.Header().Set("Content-Type", "application/json")
	w.WriteHeader(200)

	// Reformat for more APIish output.
	output := map[string]map[string]TaskState{}

	for _, tl := range tasks {
		// Remove node prefix from local task names.
		npre := tl.Node + "/"

		for k, v := range tl.Tasks {
			if strings.HasPrefix(k, npre) {
				delete(tl.Tasks, k)
				tl.Tasks[k[len(npre):]] = v
			}
		}
		output[tl.Node] = tl.Tasks
	}

	e := json.NewEncoder(w)
	err = e.Encode(output)
	if err != nil {
		log.Printf("Error encoding running tasks list: %v", err)
	}
}

func doGetMeta(w http.ResponseWriter, req *http.Request, path string) {
	got := fileMeta{}
	err := couchbase.Get(path, &got)
	if err != nil {
		log.Printf("Error getting file %#v: %v", path, err)
		w.WriteHeader(404)
		w.Write([]byte(err.Error()))
		return
	}

	w.Header().Set("Content-Type", "application/json")
	w.WriteHeader(200)
	if got.Userdata == nil {
		w.Write([]byte("{}"))
	} else {
		w.Write(*got.Userdata)
	}
}

func putMeta(w http.ResponseWriter, req *http.Request, path string) {
	got := fileMeta{}
	casid := uint64(0)
	err := couchbase.Gets(path, &got, &casid)
	if err != nil {
		log.Printf("Error getting file %#v: %v", path, err)
		w.WriteHeader(404)
		w.Write([]byte(err.Error()))
		return
	}

	r := json.RawMessage{}
	err = json.NewDecoder(req.Body).Decode(&r)
	if err != nil {
		w.WriteHeader(400)
		w.Write([]byte(err.Error()))
		return
	}

	got.Userdata = &r
	b := mustEncode(&got)

	err = couchbase.Do(path, func(mc *memcached.Client, vb uint16) error {
		req := &gomemcached.MCRequest{
			Opcode:  gomemcached.SET,
			VBucket: vb,
			Key:     []byte(path),
			Cas:     casid,
			Opaque:  0,
			Extras:  []byte{0, 0, 0, 0, 0, 0, 0, 0},
			Body:    b}
		resp, err := mc.Send(req)
		if err != nil {
			return err
		}
		if resp.Status != gomemcached.SUCCESS {
			return resp
		}
		return nil
	})

	if err == nil {
		w.WriteHeader(201)
	} else {
		w.WriteHeader(500)
		w.Write([]byte(err.Error()))
	}
}

func doListNodes(w http.ResponseWriter, req *http.Request) {
	nl, err := findAllNodes()
	if err != nil {
		log.Printf("Error executing nodes view: %v", err)
		w.WriteHeader(500)
		fmt.Fprintf(w, "Error generating node list: %v", err)
		return
	}

	respob := map[string]map[string]interface{}{}
	for _, node := range nl {
		age := time.Since(node.Time)
		respob[node.name] = map[string]interface{}{
			"size":       node.storageSize,
			"addr":       node.Address(),
			"starttime":  node.Started,
			"hbtime":     node.Time,
			"hbage_ms":   age.Nanoseconds() / 1e6,
			"hbage_str":  age.String(),
			"used":       node.Used,
			"free":       node.Free,
			"addr_raw":   node.Addr,
			"bindaddr":   node.BindAddr,
			"framesbind": node.FrameBind,
		}
		// Grandfathering these in.
		if !node.Started.IsZero() {
			uptime := time.Since(node.Started)
			respob[node.name]["uptime_ms"] = uptime.Nanoseconds() / 1e6
			respob[node.name]["uptime_str"] = uptime.String()
		}

	}

	w.Header().Set("Content-Type", "application/json")
	w.Write(mustEncode(respob))
}

func doGetFramesData(w http.ResponseWriter, req *http.Request) {
	w.Header().Set("Content-Type", "application/json")
	w.WriteHeader(200)
	w.Write(mustEncode(getFramesInfos()))
}

func proxyViewRequest(w http.ResponseWriter, req *http.Request,
	path string) {

	node := couchbase.Nodes[rand.Intn(len(couchbase.Nodes))]
	u, err := url.Parse(node.CouchAPIBase)
	if err != nil {
		w.WriteHeader(http.StatusBadGateway)
		return
	}

	u.Path = "/" + path
	u.RawQuery = req.URL.RawQuery

	client := &http.Client{
		Transport: TimeoutTransport(*viewTimeout),
	}

	res, err := client.Get(u.String())
	if err != nil {
		w.WriteHeader(http.StatusBadGateway)
		return
	}
	defer res.Body.Close()

	for k, vs := range res.Header {
		w.Header()[k] = vs
	}

	output := io.Writer(w)

	if canGzip(req) {
		w.Header().Set("Content-Encoding", "gzip")
		gz := gzip.NewWriter(w)
		defer gz.Close()
		output = gz
	}
	w.WriteHeader(res.StatusCode)

	io.Copy(output, res.Body)
}

func proxyCRUDGet(w http.ResponseWriter, req *http.Request,
	path string) {

	val, err := couchbase.GetRaw(path)
	if err != nil {
		w.WriteHeader(404)
		fmt.Fprintf(w, "Error getting value: %v", err)
		return
	}
	w.WriteHeader(200)
	w.Write(val)
}

func proxyCRUDPut(w http.ResponseWriter, req *http.Request,
	path string) {

	data, err := ioutil.ReadAll(req.Body)
	if err != nil {
		w.WriteHeader(500)
		fmt.Fprintf(w, "Error reading data: %v", err)
		return
	}

	err = couchbase.SetRaw(path, 0, data)
	if err != nil {
		w.WriteHeader(500)
		fmt.Fprintf(w, "Error storing value: %v", err)
		return
	}

	w.WriteHeader(204)
}

func proxyCRUDDelete(w http.ResponseWriter, req *http.Request,
	path string) {

	err := couchbase.Delete(path)
	if err != nil {
		w.WriteHeader(500)
		fmt.Fprintf(w, "Error deleting value: %v", err)
		return
	}

	w.WriteHeader(204)
}

func doListDocs(w http.ResponseWriter, req *http.Request,
	path string) {

	// trim off trailing slash early so we handle them consistently
	if strings.HasSuffix(path, "/") {
		path = path[0 : len(path)-1]
	}

	includeMeta := req.FormValue("includeMeta")
	depthString := req.FormValue("depth")
	depth := 1
	if depthString != "" {
		i, err := strconv.Atoi(depthString)
		if err != nil {
			w.WriteHeader(400)
			fmt.Fprintf(w, "Error processing depth parameter: %v", err)
			return
		}
		depth = i
	}

	fl, err := listFiles(path, includeMeta == "true", depth)
	if err != nil {
		log.Printf("Error executing file browse view: %v", err)
		w.WriteHeader(500)
		fmt.Fprintf(w, "Error generating file list: %v", err)
		return
	}

	if len(fl.Dirs) == 0 && len(fl.Files) == 0 {
		w.WriteHeader(404)
		return
	}

	w.Header().Set("Content-Type", "application/json")
	w.WriteHeader(200)

	e := json.NewEncoder(w)
	err = e.Encode(fl)
	if err != nil {
		log.Printf("Error writing json stream: %v", err)
	}
}

func doPing(w http.ResponseWriter, req *http.Request) {
	w.WriteHeader(204)
}
<commit_msg>When compressing, don't keep the original content-length.<commit_after>package main

import (
	"compress/gzip"
	"encoding/json"
	"fmt"
	"io"
	"io/ioutil"
	"log"
	"math/rand"
	"net/http"
	"net/url"
	"os"
	"path/filepath"
	"strconv"
	"strings"
	"time"

	"github.com/couchbaselabs/cbfs/config"
	"github.com/dustin/gomemcached"
	"github.com/dustin/gomemcached/client"
)

func doGetConfig(w http.ResponseWriter, req *http.Request) {
	err := updateConfig()
	if err != nil && !gomemcached.IsNotFound(err) {
		log.Printf("Error updating config: %v", err)
	}

	w.Header().Set("Content-Type", "application/json")
	w.WriteHeader(200)

	e := json.NewEncoder(w)
	err = e.Encode(&globalConfig)
	if err != nil {
		log.Printf("Error sending config: %v", err)
	}
}

func putConfig(w http.ResponseWriter, req *http.Request) {
	d := json.NewDecoder(req.Body)
	conf := cbfsconfig.CBFSConfig{}

	err := d.Decode(&conf)
	if err != nil {
		w.WriteHeader(500)
		fmt.Fprintf(w, "Error reading config: %v", err)
		return
	}

	err = StoreConfig(conf)
	if err != nil {
		w.WriteHeader(500)
		fmt.Fprintf(w, "Error writing config: %v", err)
		return
	}

	err = updateConfig()
	if err != nil {
		log.Printf("Error fetching newly stored config: %v", err)
	}

	w.WriteHeader(204)
}

func doList(w http.ResponseWriter, req *http.Request) {
	w.WriteHeader(200)
	explen := getHash().Size() * 2
	filepath.Walk(*root, func(path string, info os.FileInfo, err error) error {
		if err != nil {
			return err
		}
		if !info.IsDir() && !strings.HasPrefix(info.Name(), "tmp") &&
			len(info.Name()) == explen {
			_, e := w.Write([]byte(info.Name() + "\n"))
			return e
		}
		return nil
	})
}

func doListTasks(w http.ResponseWriter, req *http.Request) {
	tasks, err := listRunningTasks()
	if err != nil {
		w.WriteHeader(500)
		fmt.Fprintf(w, "Error listing tasks:  %v", err)
		return
	}

	w.Header().Set("Content-Type", "application/json")
	w.WriteHeader(200)

	// Reformat for more APIish output.
	output := map[string]map[string]TaskState{}

	for _, tl := range tasks {
		// Remove node prefix from local task names.
		npre := tl.Node + "/"

		for k, v := range tl.Tasks {
			if strings.HasPrefix(k, npre) {
				delete(tl.Tasks, k)
				tl.Tasks[k[len(npre):]] = v
			}
		}
		output[tl.Node] = tl.Tasks
	}

	e := json.NewEncoder(w)
	err = e.Encode(output)
	if err != nil {
		log.Printf("Error encoding running tasks list: %v", err)
	}
}

func doGetMeta(w http.ResponseWriter, req *http.Request, path string) {
	got := fileMeta{}
	err := couchbase.Get(path, &got)
	if err != nil {
		log.Printf("Error getting file %#v: %v", path, err)
		w.WriteHeader(404)
		w.Write([]byte(err.Error()))
		return
	}

	w.Header().Set("Content-Type", "application/json")
	w.WriteHeader(200)
	if got.Userdata == nil {
		w.Write([]byte("{}"))
	} else {
		w.Write(*got.Userdata)
	}
}

func putMeta(w http.ResponseWriter, req *http.Request, path string) {
	got := fileMeta{}
	casid := uint64(0)
	err := couchbase.Gets(path, &got, &casid)
	if err != nil {
		log.Printf("Error getting file %#v: %v", path, err)
		w.WriteHeader(404)
		w.Write([]byte(err.Error()))
		return
	}

	r := json.RawMessage{}
	err = json.NewDecoder(req.Body).Decode(&r)
	if err != nil {
		w.WriteHeader(400)
		w.Write([]byte(err.Error()))
		return
	}

	got.Userdata = &r
	b := mustEncode(&got)

	err = couchbase.Do(path, func(mc *memcached.Client, vb uint16) error {
		req := &gomemcached.MCRequest{
			Opcode:  gomemcached.SET,
			VBucket: vb,
			Key:     []byte(path),
			Cas:     casid,
			Opaque:  0,
			Extras:  []byte{0, 0, 0, 0, 0, 0, 0, 0},
			Body:    b}
		resp, err := mc.Send(req)
		if err != nil {
			return err
		}
		if resp.Status != gomemcached.SUCCESS {
			return resp
		}
		return nil
	})

	if err == nil {
		w.WriteHeader(201)
	} else {
		w.WriteHeader(500)
		w.Write([]byte(err.Error()))
	}
}

func doListNodes(w http.ResponseWriter, req *http.Request) {
	nl, err := findAllNodes()
	if err != nil {
		log.Printf("Error executing nodes view: %v", err)
		w.WriteHeader(500)
		fmt.Fprintf(w, "Error generating node list: %v", err)
		return
	}

	respob := map[string]map[string]interface{}{}
	for _, node := range nl {
		age := time.Since(node.Time)
		respob[node.name] = map[string]interface{}{
			"size":       node.storageSize,
			"addr":       node.Address(),
			"starttime":  node.Started,
			"hbtime":     node.Time,
			"hbage_ms":   age.Nanoseconds() / 1e6,
			"hbage_str":  age.String(),
			"used":       node.Used,
			"free":       node.Free,
			"addr_raw":   node.Addr,
			"bindaddr":   node.BindAddr,
			"framesbind": node.FrameBind,
		}
		// Grandfathering these in.
		if !node.Started.IsZero() {
			uptime := time.Since(node.Started)
			respob[node.name]["uptime_ms"] = uptime.Nanoseconds() / 1e6
			respob[node.name]["uptime_str"] = uptime.String()
		}

	}

	w.Header().Set("Content-Type", "application/json")
	w.Write(mustEncode(respob))
}

func doGetFramesData(w http.ResponseWriter, req *http.Request) {
	w.Header().Set("Content-Type", "application/json")
	w.WriteHeader(200)
	w.Write(mustEncode(getFramesInfos()))
}

func proxyViewRequest(w http.ResponseWriter, req *http.Request,
	path string) {

	node := couchbase.Nodes[rand.Intn(len(couchbase.Nodes))]
	u, err := url.Parse(node.CouchAPIBase)
	if err != nil {
		w.WriteHeader(http.StatusBadGateway)
		return
	}

	u.Path = "/" + path
	u.RawQuery = req.URL.RawQuery

	client := &http.Client{
		Transport: TimeoutTransport(*viewTimeout),
	}

	res, err := client.Get(u.String())
	if err != nil {
		w.WriteHeader(http.StatusBadGateway)
		return
	}
	defer res.Body.Close()

	for k, vs := range res.Header {
		w.Header()[k] = vs
	}

	output := io.Writer(w)

	if canGzip(req) {
		w.Header().Set("Content-Encoding", "gzip")
		w.Header().Del("Content-Length")
		gz := gzip.NewWriter(w)
		defer gz.Close()
		output = gz
	}
	w.WriteHeader(res.StatusCode)

	io.Copy(output, res.Body)
}

func proxyCRUDGet(w http.ResponseWriter, req *http.Request,
	path string) {

	val, err := couchbase.GetRaw(path)
	if err != nil {
		w.WriteHeader(404)
		fmt.Fprintf(w, "Error getting value: %v", err)
		return
	}
	w.WriteHeader(200)
	w.Write(val)
}

func proxyCRUDPut(w http.ResponseWriter, req *http.Request,
	path string) {

	data, err := ioutil.ReadAll(req.Body)
	if err != nil {
		w.WriteHeader(500)
		fmt.Fprintf(w, "Error reading data: %v", err)
		return
	}

	err = couchbase.SetRaw(path, 0, data)
	if err != nil {
		w.WriteHeader(500)
		fmt.Fprintf(w, "Error storing value: %v", err)
		return
	}

	w.WriteHeader(204)
}

func proxyCRUDDelete(w http.ResponseWriter, req *http.Request,
	path string) {

	err := couchbase.Delete(path)
	if err != nil {
		w.WriteHeader(500)
		fmt.Fprintf(w, "Error deleting value: %v", err)
		return
	}

	w.WriteHeader(204)
}

func doListDocs(w http.ResponseWriter, req *http.Request,
	path string) {

	// trim off trailing slash early so we handle them consistently
	if strings.HasSuffix(path, "/") {
		path = path[0 : len(path)-1]
	}

	includeMeta := req.FormValue("includeMeta")
	depthString := req.FormValue("depth")
	depth := 1
	if depthString != "" {
		i, err := strconv.Atoi(depthString)
		if err != nil {
			w.WriteHeader(400)
			fmt.Fprintf(w, "Error processing depth parameter: %v", err)
			return
		}
		depth = i
	}

	fl, err := listFiles(path, includeMeta == "true", depth)
	if err != nil {
		log.Printf("Error executing file browse view: %v", err)
		w.WriteHeader(500)
		fmt.Fprintf(w, "Error generating file list: %v", err)
		return
	}

	if len(fl.Dirs) == 0 && len(fl.Files) == 0 {
		w.WriteHeader(404)
		return
	}

	w.Header().Set("Content-Type", "application/json")
	w.WriteHeader(200)

	e := json.NewEncoder(w)
	err = e.Encode(fl)
	if err != nil {
		log.Printf("Error writing json stream: %v", err)
	}
}

func doPing(w http.ResponseWriter, req *http.Request) {
	w.WriteHeader(204)
}
<|endoftext|>