<commit_before>package gatt

import (
	"encoding/binary"
	"fmt"
	"log"
	"sync"

	"github.com/currantlabs/ble"
	"github.com/currantlabs/ble/linux/att"
)

const (
	cccNotify   = 0x0001
	cccIndicate = 0x0002
)

// NewClient returns a GATT Client.
func NewClient(conn ble.Conn) (*Client, error) {
	p := &Client{
		subs: make(map[uint16]*sub),
		conn: conn,
	}
	p.ac = att.NewClient(conn, p)
	go p.ac.Loop()
	return p, nil
}

// A Client is a GATT Client.
type Client struct {
	sync.RWMutex

	profile *ble.Profile
	name    string
	subs    map[uint16]*sub

	ac   *att.Client
	conn ble.Conn
}

// Address returns the address of the client.
func (p *Client) Address() ble.Addr {
	p.RLock()
	defer p.RUnlock()
	return p.conn.RemoteAddr()
}

// Name returns the name of the client.
func (p *Client) Name() string {
	p.RLock()
	defer p.RUnlock()
	return p.name
}

// Profile returns the discovered profile.
func (p *Client) Profile() *ble.Profile {
	p.RLock()
	defer p.RUnlock()
	return p.profile
}

// DiscoverProfile discovers the whole hierachy of a server.
func (p *Client) DiscoverProfile(force bool) (*ble.Profile, error) {
	if p.profile != nil && !force {
		return p.profile, nil
	}
	ss, err := p.DiscoverServices(nil)
	if err != nil {
		return nil, fmt.Errorf("can't discover services: %s\n", err)
	}
	for _, s := range ss {
		cs, err := p.DiscoverCharacteristics(nil, s)
		if err != nil {
			return nil, fmt.Errorf("can't discover characteristics: %s\n", err)
		}
		for _, c := range cs {
			_, err := p.DiscoverDescriptors(nil, c)
			if err != nil {
				return nil, fmt.Errorf("can't discover descriptors: %s\n", err)
			}
		}
	}
	p.profile = &ble.Profile{Services: ss}
	return p.profile, nil
}

// DiscoverServices finds all the primary services on a server. [Vol 3, Part G, 4.4.1]
// If filter is specified, only filtered services are returned.
func (p *Client) DiscoverServices(filter []ble.UUID) ([]*ble.Service, error) {
	p.Lock()
	defer p.Unlock()
	if p.profile == nil {
		p.profile = &ble.Profile{}
	}
	start := uint16(0x0001)
	for {
		length, b, err := p.ac.ReadByGroupType(start, 0xFFFF, ble.PrimaryServiceUUID)
		if err == ble.ErrAttrNotFound {
			return p.profile.Services, nil
		}
		if err != nil {
			return nil, err
		}
		for len(b) != 0 {
			h := binary.LittleEndian.Uint16(b[:2])
			endh := binary.LittleEndian.Uint16(b[2:4])
			u := ble.UUID(b[4:length])
			if filter == nil || ble.Contains(filter, u) {
				s := &ble.Service{
					UUID:      u,
					Handle:    h,
					EndHandle: endh,
				}
				p.profile.Services = append(p.profile.Services, s)
			}
			if endh == 0xFFFF {
				return p.profile.Services, nil
			}
			start = endh + 1
			b = b[length:]
		}
	}
}

// DiscoverIncludedServices finds the included services of a service. [Vol 3, Part G, 4.5.1]
// If filter is specified, only filtered services are returned.
func (p *Client) DiscoverIncludedServices(ss []ble.UUID, s *ble.Service) ([]*ble.Service, error) {
	p.Lock()
	defer p.Unlock()
	return nil, nil
}

// DiscoverCharacteristics finds all the characteristics within a service. [Vol 3, Part G, 4.6.1]
// If filter is specified, only filtered characteristics are returned.
func (p *Client) DiscoverCharacteristics(filter []ble.UUID, s *ble.Service) ([]*ble.Characteristic, error) {
	p.Lock()
	defer p.Unlock()
	start := s.Handle
	var lastChar *ble.Characteristic
	for start <= s.EndHandle {
		length, b, err := p.ac.ReadByType(start, s.EndHandle, ble.CharacteristicUUID)
		if err == ble.ErrAttrNotFound {
			break
		} else if err != nil {
			return nil, err
		}
		for len(b) != 0 {
			h := binary.LittleEndian.Uint16(b[:2])
			p := ble.Property(b[2])
			vh := binary.LittleEndian.Uint16(b[3:5])
			u := ble.UUID(b[5:length])
			c := &ble.Characteristic{
				UUID:        u,
				Property:    p,
				Handle:      h,
				ValueHandle: vh,
				EndHandle:   s.EndHandle,
			}
			if filter == nil || ble.Contains(filter, u) {
				s.Characteristics = append(s.Characteristics, c)
			}
			if lastChar != nil {
				lastChar.EndHandle = c.Handle - 1
			}
			lastChar = c
			start = vh + 1
			b = b[length:]
		}
	}
	return s.Characteristics, nil
}

// DiscoverDescriptors finds all the descriptors within a characteristic. [Vol 3, Part G, 4.7.1]
// If filter is specified, only filtered descriptors are returned.
func (p *Client) DiscoverDescriptors(filter []ble.UUID, c *ble.Characteristic) ([]*ble.Descriptor, error) {
	p.Lock()
	defer p.Unlock()
	start := c.ValueHandle + 1
	for start <= c.EndHandle {
		fmt, b, err := p.ac.FindInformation(start, c.EndHandle)
		if err == ble.ErrAttrNotFound {
			break
		} else if err != nil {
			return nil, err
		}
		length := 2 + 2
		if fmt == 0x02 {
			length = 2 + 16
		}
		for len(b) != 0 {
			h := binary.LittleEndian.Uint16(b[:2])
			u := ble.UUID(b[2:length])
			d := &ble.Descriptor{UUID: u, Handle: h}
			if filter == nil || ble.Contains(filter, u) {
				c.Descriptors = append(c.Descriptors, d)
			}
			if u.Equal(ble.ClientCharacteristicConfigUUID) {
				c.CCCD = d
			}
			start = h + 1
			b = b[length:]
		}
	}
	return c.Descriptors, nil
}

// ReadCharacteristic reads a characteristic value from a server. [Vol 3, Part G, 4.8.1]
func (p *Client) ReadCharacteristic(c *ble.Characteristic) ([]byte, error) {
	p.Lock()
	defer p.Unlock()
	return p.ac.Read(c.ValueHandle)
}

// ReadLongCharacteristic reads a characteristic value which is longer than the MTU. [Vol 3, Part G, 4.8.3]
func (p *Client) ReadLongCharacteristic(c *ble.Characteristic) ([]byte, error) {
	p.Lock()
	defer p.Unlock()
	return nil, nil
}

// WriteCharacteristic writes a characteristic value to a server. [Vol 3, Part G, 4.9.3]
func (p *Client) WriteCharacteristic(c *ble.Characteristic, v []byte, noRsp bool) error {
	p.Lock()
	defer p.Unlock()
	if noRsp {
		p.ac.WriteCommand(c.ValueHandle, v)
		return nil
	}
	return p.ac.Write(c.ValueHandle, v)
}

// ReadDescriptor reads a characteristic descriptor from a server. [Vol 3, Part G, 4.12.1]
func (p *Client) ReadDescriptor(d *ble.Descriptor) ([]byte, error) {
	p.Lock()
	defer p.Unlock()
	return p.ac.Read(d.Handle)
}

// WriteDescriptor writes a characteristic descriptor to a server. [Vol 3, Part G, 4.12.3]
func (p *Client) WriteDescriptor(d *ble.Descriptor, v []byte) error {
	p.Lock()
	defer p.Unlock()
	return p.ac.Write(d.Handle, v)
}

// ReadRSSI retrieves the current RSSI value of remote peripheral. [Vol 2, Part E, 7.5.4]
func (p *Client) ReadRSSI() int {
	p.Lock()
	defer p.Unlock()
	// TODO:
	return 0
}

// ExchangeMTU informs the server of the client’s maximum receive MTU size and
// request the server to respond with its maximum receive MTU size. [Vol 3, Part F, 3.4.2.1]
func (p *Client) ExchangeMTU(mtu int) (int, error) {
	p.Lock()
	defer p.Unlock()
	return p.ac.ExchangeMTU(mtu)
}

// Subscribe subscribes to indication (if ind is set true), or notification of a
// characteristic value. [Vol 3, Part G, 4.10 & 4.11]
func (p *Client) Subscribe(c *ble.Characteristic, ind bool, h ble.NotificationHandler) error {
	p.Lock()
	defer p.Unlock()
	if ind {
		return p.setHandlers(c.CCCD.Handle, c.ValueHandle, cccIndicate, h)
	}
	return p.setHandlers(c.CCCD.Handle, c.ValueHandle, cccNotify, h)
}

// Unsubscribe unsubscribes to indication (if ind is set true), or notification
// of a specified characteristic value. [Vol 3, Part G, 4.10 & 4.11]
func (p *Client) Unsubscribe(c *ble.Characteristic, ind bool) error {
	p.Lock()
	defer p.Unlock()
	if ind {
		return p.setHandlers(c.CCCD.Handle, c.ValueHandle, cccIndicate, nil)
	}
	return p.setHandlers(c.CCCD.Handle, c.ValueHandle, cccNotify, nil)
}

func (p *Client) setHandlers(cccdh, vh, flag uint16, h ble.NotificationHandler) error {
	s, ok := p.subs[vh]
	if !ok {
		s = &sub{cccdh, 0x0000, nil, nil}
		p.subs[vh] = s
	}
	switch {
	case h == nil && (s.ccc&flag) == 0:
		return nil
	case h != nil && (s.ccc&flag) != 0:
		return nil
	case h == nil && (s.ccc&flag) != 0:
		s.ccc &= ^uint16(flag)
	case h != nil && (s.ccc&flag) == 0:
		s.ccc |= flag
	}

	v := make([]byte, 2)
	binary.LittleEndian.PutUint16(v, s.ccc)
	if flag == cccNotify {
		s.nHandler = h
	} else {
		s.iHandler = h
	}
	return p.ac.Write(s.cccdh, v)
}

// ClearSubscriptions clears all subscriptions to notifications and indications.
func (p *Client) ClearSubscriptions() error {
	p.Lock()
	defer p.Unlock()
	zero := make([]byte, 2)
	for vh, s := range p.subs {
		if err := p.ac.Write(s.cccdh, zero); err != nil {
			return err
		}
		delete(p.subs, vh)
	}
	return nil
}

// CancelConnection disconnects the connection.
func (p *Client) CancelConnection() error {
	p.Lock()
	defer p.Unlock()
	return p.conn.Close()
}

// HandleNotification ...
func (p *Client) HandleNotification(req []byte) {
	p.Lock()
	defer p.Unlock()
	vh := att.HandleValueIndication(req).AttributeHandle()
	sub, ok := p.subs[vh]
	if !ok {
		// FIXME: disconnects and propagate an error to the user.
		log.Printf("Got an unregistered notification")
		return
	}
	fn := sub.nHandler
	if req[0] == att.HandleValueIndicationCode {
		fn = sub.iHandler
	}
	if fn != nil {
		fn(req[3:])
	}
}

type sub struct {
	cccdh    uint16
	ccc      uint16
	nHandler ble.NotificationHandler
	iHandler ble.NotificationHandler
}
<commit_msg>gatt: check CCCD before subscription<commit_after>package gatt

import (
	"encoding/binary"
	"fmt"
	"log"
	"sync"

	"github.com/currantlabs/ble"
	"github.com/currantlabs/ble/linux/att"
)

const (
	cccNotify   = 0x0001
	cccIndicate = 0x0002
)

// NewClient returns a GATT Client.
func NewClient(conn ble.Conn) (*Client, error) {
	p := &Client{
		subs: make(map[uint16]*sub),
		conn: conn,
	}
	p.ac = att.NewClient(conn, p)
	go p.ac.Loop()
	return p, nil
}

// A Client is a GATT Client.
type Client struct {
	sync.RWMutex

	profile *ble.Profile
	name    string
	subs    map[uint16]*sub

	ac   *att.Client
	conn ble.Conn
}

// Address returns the address of the client.
func (p *Client) Address() ble.Addr {
	p.RLock()
	defer p.RUnlock()
	return p.conn.RemoteAddr()
}

// Name returns the name of the client.
func (p *Client) Name() string {
	p.RLock()
	defer p.RUnlock()
	return p.name
}

// Profile returns the discovered profile.
func (p *Client) Profile() *ble.Profile {
	p.RLock()
	defer p.RUnlock()
	return p.profile
}

// DiscoverProfile discovers the whole hierachy of a server.
func (p *Client) DiscoverProfile(force bool) (*ble.Profile, error) {
	if p.profile != nil && !force {
		return p.profile, nil
	}
	ss, err := p.DiscoverServices(nil)
	if err != nil {
		return nil, fmt.Errorf("can't discover services: %s\n", err)
	}
	for _, s := range ss {
		cs, err := p.DiscoverCharacteristics(nil, s)
		if err != nil {
			return nil, fmt.Errorf("can't discover characteristics: %s\n", err)
		}
		for _, c := range cs {
			_, err := p.DiscoverDescriptors(nil, c)
			if err != nil {
				return nil, fmt.Errorf("can't discover descriptors: %s\n", err)
			}
		}
	}
	p.profile = &ble.Profile{Services: ss}
	return p.profile, nil
}

// DiscoverServices finds all the primary services on a server. [Vol 3, Part G, 4.4.1]
// If filter is specified, only filtered services are returned.
func (p *Client) DiscoverServices(filter []ble.UUID) ([]*ble.Service, error) {
	p.Lock()
	defer p.Unlock()
	if p.profile == nil {
		p.profile = &ble.Profile{}
	}
	start := uint16(0x0001)
	for {
		length, b, err := p.ac.ReadByGroupType(start, 0xFFFF, ble.PrimaryServiceUUID)
		if err == ble.ErrAttrNotFound {
			return p.profile.Services, nil
		}
		if err != nil {
			return nil, err
		}
		for len(b) != 0 {
			h := binary.LittleEndian.Uint16(b[:2])
			endh := binary.LittleEndian.Uint16(b[2:4])
			u := ble.UUID(b[4:length])
			if filter == nil || ble.Contains(filter, u) {
				s := &ble.Service{
					UUID:      u,
					Handle:    h,
					EndHandle: endh,
				}
				p.profile.Services = append(p.profile.Services, s)
			}
			if endh == 0xFFFF {
				return p.profile.Services, nil
			}
			start = endh + 1
			b = b[length:]
		}
	}
}

// DiscoverIncludedServices finds the included services of a service. [Vol 3, Part G, 4.5.1]
// If filter is specified, only filtered services are returned.
func (p *Client) DiscoverIncludedServices(ss []ble.UUID, s *ble.Service) ([]*ble.Service, error) {
	p.Lock()
	defer p.Unlock()
	return nil, nil
}

// DiscoverCharacteristics finds all the characteristics within a service. [Vol 3, Part G, 4.6.1]
// If filter is specified, only filtered characteristics are returned.
func (p *Client) DiscoverCharacteristics(filter []ble.UUID, s *ble.Service) ([]*ble.Characteristic, error) {
	p.Lock()
	defer p.Unlock()
	start := s.Handle
	var lastChar *ble.Characteristic
	for start <= s.EndHandle {
		length, b, err := p.ac.ReadByType(start, s.EndHandle, ble.CharacteristicUUID)
		if err == ble.ErrAttrNotFound {
			break
		} else if err != nil {
			return nil, err
		}
		for len(b) != 0 {
			h := binary.LittleEndian.Uint16(b[:2])
			p := ble.Property(b[2])
			vh := binary.LittleEndian.Uint16(b[3:5])
			u := ble.UUID(b[5:length])
			c := &ble.Characteristic{
				UUID:        u,
				Property:    p,
				Handle:      h,
				ValueHandle: vh,
				EndHandle:   s.EndHandle,
			}
			if filter == nil || ble.Contains(filter, u) {
				s.Characteristics = append(s.Characteristics, c)
			}
			if lastChar != nil {
				lastChar.EndHandle = c.Handle - 1
			}
			lastChar = c
			start = vh + 1
			b = b[length:]
		}
	}
	return s.Characteristics, nil
}

// DiscoverDescriptors finds all the descriptors within a characteristic. [Vol 3, Part G, 4.7.1]
// If filter is specified, only filtered descriptors are returned.
func (p *Client) DiscoverDescriptors(filter []ble.UUID, c *ble.Characteristic) ([]*ble.Descriptor, error) {
	p.Lock()
	defer p.Unlock()
	start := c.ValueHandle + 1
	for start <= c.EndHandle {
		fmt, b, err := p.ac.FindInformation(start, c.EndHandle)
		if err == ble.ErrAttrNotFound {
			break
		} else if err != nil {
			return nil, err
		}
		length := 2 + 2
		if fmt == 0x02 {
			length = 2 + 16
		}
		for len(b) != 0 {
			h := binary.LittleEndian.Uint16(b[:2])
			u := ble.UUID(b[2:length])
			d := &ble.Descriptor{UUID: u, Handle: h}
			if filter == nil || ble.Contains(filter, u) {
				c.Descriptors = append(c.Descriptors, d)
			}
			if u.Equal(ble.ClientCharacteristicConfigUUID) {
				c.CCCD = d
			}
			start = h + 1
			b = b[length:]
		}
	}
	return c.Descriptors, nil
}

// ReadCharacteristic reads a characteristic value from a server. [Vol 3, Part G, 4.8.1]
func (p *Client) ReadCharacteristic(c *ble.Characteristic) ([]byte, error) {
	p.Lock()
	defer p.Unlock()
	return p.ac.Read(c.ValueHandle)
}

// ReadLongCharacteristic reads a characteristic value which is longer than the MTU. [Vol 3, Part G, 4.8.3]
func (p *Client) ReadLongCharacteristic(c *ble.Characteristic) ([]byte, error) {
	p.Lock()
	defer p.Unlock()
	return nil, nil
}

// WriteCharacteristic writes a characteristic value to a server. [Vol 3, Part G, 4.9.3]
func (p *Client) WriteCharacteristic(c *ble.Characteristic, v []byte, noRsp bool) error {
	p.Lock()
	defer p.Unlock()
	if noRsp {
		p.ac.WriteCommand(c.ValueHandle, v)
		return nil
	}
	return p.ac.Write(c.ValueHandle, v)
}

// ReadDescriptor reads a characteristic descriptor from a server. [Vol 3, Part G, 4.12.1]
func (p *Client) ReadDescriptor(d *ble.Descriptor) ([]byte, error) {
	p.Lock()
	defer p.Unlock()
	return p.ac.Read(d.Handle)
}

// WriteDescriptor writes a characteristic descriptor to a server. [Vol 3, Part G, 4.12.3]
func (p *Client) WriteDescriptor(d *ble.Descriptor, v []byte) error {
	p.Lock()
	defer p.Unlock()
	return p.ac.Write(d.Handle, v)
}

// ReadRSSI retrieves the current RSSI value of remote peripheral. [Vol 2, Part E, 7.5.4]
func (p *Client) ReadRSSI() int {
	p.Lock()
	defer p.Unlock()
	// TODO:
	return 0
}

// ExchangeMTU informs the server of the client’s maximum receive MTU size and
// request the server to respond with its maximum receive MTU size. [Vol 3, Part F, 3.4.2.1]
func (p *Client) ExchangeMTU(mtu int) (int, error) {
	p.Lock()
	defer p.Unlock()
	return p.ac.ExchangeMTU(mtu)
}

// Subscribe subscribes to indication (if ind is set true), or notification of a
// characteristic value. [Vol 3, Part G, 4.10 & 4.11]
func (p *Client) Subscribe(c *ble.Characteristic, ind bool, h ble.NotificationHandler) error {
	p.Lock()
	defer p.Unlock()
	if c.CCCD == nil {
		return fmt.Errorf("CCCD not found")
	}
	if ind {
		return p.setHandlers(c.CCCD.Handle, c.ValueHandle, cccIndicate, h)
	}
	return p.setHandlers(c.CCCD.Handle, c.ValueHandle, cccNotify, h)
}

// Unsubscribe unsubscribes to indication (if ind is set true), or notification
// of a specified characteristic value. [Vol 3, Part G, 4.10 & 4.11]
func (p *Client) Unsubscribe(c *ble.Characteristic, ind bool) error {
	p.Lock()
	defer p.Unlock()
	if c.CCCD == nil {
		return fmt.Errorf("CCCD not found")
	}
	if ind {
		return p.setHandlers(c.CCCD.Handle, c.ValueHandle, cccIndicate, nil)
	}
	return p.setHandlers(c.CCCD.Handle, c.ValueHandle, cccNotify, nil)
}

func (p *Client) setHandlers(cccdh, vh, flag uint16, h ble.NotificationHandler) error {
	s, ok := p.subs[vh]
	if !ok {
		s = &sub{cccdh, 0x0000, nil, nil}
		p.subs[vh] = s
	}
	switch {
	case h == nil && (s.ccc&flag) == 0:
		return nil
	case h != nil && (s.ccc&flag) != 0:
		return nil
	case h == nil && (s.ccc&flag) != 0:
		s.ccc &= ^uint16(flag)
	case h != nil && (s.ccc&flag) == 0:
		s.ccc |= flag
	}

	v := make([]byte, 2)
	binary.LittleEndian.PutUint16(v, s.ccc)
	if flag == cccNotify {
		s.nHandler = h
	} else {
		s.iHandler = h
	}
	return p.ac.Write(s.cccdh, v)
}

// ClearSubscriptions clears all subscriptions to notifications and indications.
func (p *Client) ClearSubscriptions() error {
	p.Lock()
	defer p.Unlock()
	zero := make([]byte, 2)
	for vh, s := range p.subs {
		if err := p.ac.Write(s.cccdh, zero); err != nil {
			return err
		}
		delete(p.subs, vh)
	}
	return nil
}

// CancelConnection disconnects the connection.
func (p *Client) CancelConnection() error {
	p.Lock()
	defer p.Unlock()
	return p.conn.Close()
}

// HandleNotification ...
func (p *Client) HandleNotification(req []byte) {
	p.Lock()
	defer p.Unlock()
	vh := att.HandleValueIndication(req).AttributeHandle()
	sub, ok := p.subs[vh]
	if !ok {
		// FIXME: disconnects and propagate an error to the user.
		log.Printf("Got an unregistered notification")
		return
	}
	fn := sub.nHandler
	if req[0] == att.HandleValueIndicationCode {
		fn = sub.iHandler
	}
	if fn != nil {
		fn(req[3:])
	}
}

type sub struct {
	cccdh    uint16
	ccc      uint16
	nHandler ble.NotificationHandler
	iHandler ble.NotificationHandler
}
<|endoftext|>


<commit_before>package main

import (
	"bufio"
	"bytes"
	"fmt"
	"io/ioutil"
	"net/http"
	"path/filepath"
	"regexp"
	"strconv"
	"strings"

	"github.com/lxc/lxd/lxd/metrics"
	"github.com/lxc/lxd/lxd/response"
	"github.com/lxc/lxd/lxd/storage/filesystem"
	"github.com/lxc/lxd/shared"
)

// These mountpoints are excluded as they are irrelevant for metrics.
// /var/lib/docker/* subdirectories are excluded for this reason: https://github.com/prometheus/node_exporter/pull/1003
var defMountPointsExcluded = regexp.MustCompile("^/(dev|proc|sys|var/lib/docker/.+)($|/)")
var defFSTypesExcluded = []string{
	"autofs", "binfmt_misc", "bpf", "cgroup", "cgroup2", "configfs", "debugfs", "devpts", "devtmpfs", "fusectl", "hugetlbfs", "iso9660", "mqueue", "nsfs", "overlay", "proc", "procfs", "pstore", "rpc_pipefs", "securityfs", "selinuxfs", "squashfs", "sysfs", "tracefs"}

var metricsCmd = APIEndpoint{
	Path: "metrics",

	Get: APIEndpointAction{Handler: metricsGet},
}

func metricsGet(d *Daemon, r *http.Request) response.Response {
	out := metrics.Metrics{}

	diskStats, err := getDiskMetrics(d)
	if err != nil {
		return response.SmartError(err)
	}

	out.Disk = diskStats

	filesystemStats, err := getFilesystemMetrics(d)
	if err != nil {
		return response.SmartError(err)
	}

	out.Filesystem = filesystemStats

	memStats, err := getMemoryMetrics(d)
	if err != nil {
		return response.SmartError(err)
	}

	out.Memory = memStats

	netStats, err := getNetworkMetrics(d)
	if err != nil {
		return response.SmartError(err)
	}

	out.Network = netStats

	out.ProcessesTotal, err = getTotalProcesses(d)
	if err != nil {
		return response.SmartError(err)
	}

	cpuStats, err := getCPUMetrics(d)
	if err != nil {
		return response.SmartError(err)
	}

	out.CPU = cpuStats

	return response.SyncResponse(true, &out)
}

func getCPUMetrics(d *Daemon) (map[string]metrics.CPUMetrics, error) {
	stats, err := ioutil.ReadFile("/proc/stat")
	if err != nil {
		return nil, fmt.Errorf("Failed to read /proc/stat: %w", err)
	}

	out := map[string]metrics.CPUMetrics{}
	scanner := bufio.NewScanner(bytes.NewReader(stats))

	for scanner.Scan() {
		fields := strings.Fields(scanner.Text())

		// Only consider CPU info, skip everything else.
		if !strings.HasPrefix(fields[0], "cpu") {
			continue
		}

		stats := metrics.CPUMetrics{}

		stats.SecondsUser, err = strconv.ParseUint(fields[1], 10, 64)
		if err != nil {
			return nil, fmt.Errorf("Failed to parse %q: %w", fields[1], err)
		}

		stats.SecondsUser *= 10

		stats.SecondsNice, err = strconv.ParseUint(fields[2], 10, 64)
		if err != nil {
			return nil, fmt.Errorf("Failed to parse %q: %w", fields[2], err)
		}

		stats.SecondsNice *= 10

		stats.SecondsSystem, err = strconv.ParseUint(fields[3], 10, 64)
		if err != nil {
			return nil, fmt.Errorf("Failed to parse %q: %w", fields[3], err)
		}

		stats.SecondsSystem *= 10

		stats.SecondsIdle, err = strconv.ParseUint(fields[4], 10, 64)
		if err != nil {
			return nil, fmt.Errorf("Failed to parse %q: %w", fields[4], err)
		}

		stats.SecondsIdle *= 10

		stats.SecondsIOWait, err = strconv.ParseUint(fields[5], 10, 64)
		if err != nil {
			return nil, fmt.Errorf("Failed to parse %q: %w", fields[5], err)
		}

		stats.SecondsIdle *= 10

		stats.SecondsIRQ, err = strconv.ParseUint(fields[6], 10, 64)
		if err != nil {
			return nil, fmt.Errorf("Failed to parse %q: %w", fields[6], err)
		}

		stats.SecondsIRQ *= 10

		stats.SecondsSoftIRQ, err = strconv.ParseUint(fields[7], 10, 64)
		if err != nil {
			return nil, fmt.Errorf("Failed to parse %q: %w", fields[7], err)
		}

		stats.SecondsSoftIRQ *= 10

		stats.SecondsSteal, err = strconv.ParseUint(fields[8], 10, 64)
		if err != nil {
			return nil, fmt.Errorf("Failed to parse %q: %w", fields[8], err)
		}

		stats.SecondsSteal *= 10

		out[fields[0]] = stats
	}

	return out, nil
}

func getTotalProcesses(d *Daemon) (uint64, error) {
	entries, err := ioutil.ReadDir("/proc")
	if err != nil {
		return 0, fmt.Errorf("Failed to read dir %q: %w", "/proc", err)
	}

	re := regexp.MustCompile(`^[[:digit:]]+$`)
	pidCount := uint64(0)

	for _, entry := range entries {
		// Skip everything which isn't a directory
		if !entry.IsDir() {
			continue
		}

		// Skip all non-PID directories
		if !re.MatchString(entry.Name()) {
			continue
		}

		cmdlinePath := filepath.Join("/proc", entry.Name(), "cmdline")

		cmdline, err := ioutil.ReadFile(cmdlinePath)
		if err != nil {
			continue
		}

		if string(cmdline) == "" {
			continue
		}

		pidCount++
	}

	return pidCount, nil
}

func getDiskMetrics(d *Daemon) (map[string]metrics.DiskMetrics, error) {
	diskStats, err := ioutil.ReadFile("/proc/diskstats")
	if err != nil {
		return nil, fmt.Errorf("Failed to read /proc/diskstats: %w", err)
	}

	out := map[string]metrics.DiskMetrics{}
	scanner := bufio.NewScanner(bytes.NewReader(diskStats))

	for scanner.Scan() {
		line := scanner.Text()

		if line == "" {
			continue
		}

		fields := strings.Fields(line)
		stats := metrics.DiskMetrics{}

		stats.ReadsCompleted, err = strconv.ParseUint(fields[3], 10, 64)
		if err != nil {
			return nil, fmt.Errorf("Failed to parse %q: %w", fields[3], err)
		}

		sectorsRead, err := strconv.ParseUint(fields[5], 10, 64)
		if err != nil {
			return nil, fmt.Errorf("Failed to parse %q: %w", fields[3], err)
		}

		stats.ReadBytes = sectorsRead * 512

		stats.WritesCompleted, err = strconv.ParseUint(fields[7], 10, 64)
		if err != nil {
			return nil, fmt.Errorf("Failed to parse %q: %w", fields[3], err)
		}

		sectorsWritten, err := strconv.ParseUint(fields[9], 10, 64)
		if err != nil {
			return nil, fmt.Errorf("Failed to parse %q: %w", fields[3], err)
		}

		stats.WrittenBytes = sectorsWritten * 512

		out[fields[2]] = stats
	}

	return out, nil
}

func getFilesystemMetrics(d *Daemon) (map[string]metrics.FilesystemMetrics, error) {
	mounts, err := ioutil.ReadFile("/proc/mounts")
	if err != nil {
		return nil, fmt.Errorf("Failed to read /proc/mounts: %w", err)
	}

	out := map[string]metrics.FilesystemMetrics{}
	scanner := bufio.NewScanner(bytes.NewReader(mounts))

	for scanner.Scan() {
		fields := strings.Fields(scanner.Text())

		// Skip uninteresting mounts
		if shared.StringInSlice(fields[2], defFSTypesExcluded) || defMountPointsExcluded.MatchString(fields[1]) {
			continue
		}

		stats := metrics.FilesystemMetrics{}

		stats.Mountpoint = fields[1]

		statfs, err := filesystem.StatVFS(stats.Mountpoint)
		if err != nil {
			return nil, fmt.Errorf("Failed to stat %s: %w", stats.Mountpoint, err)
		}

		fsType, err := filesystem.FSTypeToName(int32(statfs.Type))
		if err == nil {
			stats.FSType = fsType
		}

		stats.SizeBytes = statfs.Blocks * uint64(statfs.Bsize)
		stats.SizeBytes = statfs.Blocks * uint64(statfs.Bsize)
		stats.SizeBytes = statfs.Blocks * uint64(statfs.Bsize)
		stats.SizeBytes = statfs.Blocks * uint64(statfs.Bsize)

		out[fields[0]] = stats
	}

	return out, nil
}

func getMemoryMetrics(d *Daemon) (metrics.MemoryMetrics, error) {
	content, err := ioutil.ReadFile("/proc/meminfo")
	if err != nil {
		return metrics.MemoryMetrics{}, fmt.Errorf("Failed to read /proc/meminfo: %w", err)
	}

	out := metrics.MemoryMetrics{}
	scanner := bufio.NewScanner(bytes.NewReader(content))

	for scanner.Scan() {
		fields := strings.Fields(scanner.Text())

		fields[0] = strings.TrimRight(fields[0], ":")

		value, err := strconv.ParseUint(fields[1], 10, 64)
		if err != nil {
			return metrics.MemoryMetrics{}, fmt.Errorf("Failed to parse %q: %w", fields[1], err)
		}

		// Multiply suffix (kB)
		if len(fields) == 3 {
			value *= 1024
		}

		// FIXME: Missing RSS
		switch fields[0] {
		case "Active":
			out.ActiveBytes = value
		case "Active(anon)":
			out.ActiveAnonBytes = value
		case "Active(file)":
			out.ActiveFileBytes = value
		case "Cached":
			out.CachedBytes = value
		case "Dirty":
			out.DirtyBytes = value
		case "HugePages_Free":
			out.HugepagesFreeBytes = value
		case "HugePages_Total":
			out.HugepagesTotalBytes = value
		case "Inactive":
			out.InactiveBytes = value
		case "Inactive(anon)":
			out.InactiveAnonBytes = value
		case "Inactive(file)":
			out.InactiveFileBytes = value
		case "Mapped":
			out.MappedBytes = value
		case "MemAvailable":
			out.MemAvailableBytes = value
		case "MemFree":
			out.MemFreeBytes = value
		case "MemTotal":
			out.MemTotalBytes = value
		case "Shmem":
			out.ShmemBytes = value
		case "SwapCached":
			out.SwapBytes = value
		case "Unevictable":
			out.UnevictableBytes = value
		case "Writeback":
			out.WritebackBytes = value
		}
	}

	return out, nil
}

func getNetworkMetrics(d *Daemon) (map[string]metrics.NetworkMetrics, error) {
	out := map[string]metrics.NetworkMetrics{}

	for dev, state := range networkState() {
		stats := metrics.NetworkMetrics{}

		stats.ReceiveBytes = uint64(state.Counters.BytesReceived)
		stats.ReceiveDrop = uint64(state.Counters.PacketsDroppedInbound)
		stats.ReceiveErrors = uint64(state.Counters.ErrorsReceived)
		stats.ReceivePackets = uint64(state.Counters.PacketsReceived)
		stats.TransmitBytes = uint64(state.Counters.BytesReceived)
		stats.TransmitDrop = uint64(state.Counters.PacketsDroppedOutbound)
		stats.TransmitErrors = uint64(state.Counters.ErrorsSent)
		stats.TransmitPackets = uint64(state.Counters.PacketsSent)

		out[dev] = stats
	}

	return out, nil
}
<commit_msg>lxd-agent: Log metrics failures<commit_after>package main

import (
	"bufio"
	"bytes"
	"fmt"
	"io/ioutil"
	"net/http"
	"path/filepath"
	"regexp"
	"strconv"
	"strings"

	"github.com/lxc/lxd/lxd/metrics"
	"github.com/lxc/lxd/lxd/response"
	"github.com/lxc/lxd/lxd/storage/filesystem"
	"github.com/lxc/lxd/shared"
	log "github.com/lxc/lxd/shared/log15"
	"github.com/lxc/lxd/shared/logger"
)

// These mountpoints are excluded as they are irrelevant for metrics.
// /var/lib/docker/* subdirectories are excluded for this reason: https://github.com/prometheus/node_exporter/pull/1003
var defMountPointsExcluded = regexp.MustCompile("^/(dev|proc|sys|var/lib/docker/.+)($|/)")
var defFSTypesExcluded = []string{
	"autofs", "binfmt_misc", "bpf", "cgroup", "cgroup2", "configfs", "debugfs", "devpts", "devtmpfs", "fusectl", "hugetlbfs", "iso9660", "mqueue", "nsfs", "overlay", "proc", "procfs", "pstore", "rpc_pipefs", "securityfs", "selinuxfs", "squashfs", "sysfs", "tracefs"}

var metricsCmd = APIEndpoint{
	Path: "metrics",

	Get: APIEndpointAction{Handler: metricsGet},
}

func metricsGet(d *Daemon, r *http.Request) response.Response {
	out := metrics.Metrics{}

	diskStats, err := getDiskMetrics(d)
	if err != nil {
		logger.Warn("Failed to get disk metrics", log.Ctx{"err": err})
	} else {
		out.Disk = diskStats
	}

	filesystemStats, err := getFilesystemMetrics(d)
	if err != nil {
		logger.Warn("Failed to get filesystem metrics", log.Ctx{"err": err})
	} else {
		out.Filesystem = filesystemStats
	}

	memStats, err := getMemoryMetrics(d)
	if err != nil {
		logger.Warn("Failed to get memory metrics", log.Ctx{"err": err})
	} else {
		out.Memory = memStats
	}

	netStats, err := getNetworkMetrics(d)
	if err != nil {
		logger.Warn("Failed to get network metrics", log.Ctx{"err": err})
	} else {
		out.Network = netStats
	}

	out.ProcessesTotal, err = getTotalProcesses(d)
	if err != nil {
		logger.Warn("Failed to get total processes", log.Ctx{"err": err})
	}

	cpuStats, err := getCPUMetrics(d)
	if err != nil {
		logger.Warn("Failed to get CPU metrics", log.Ctx{"err": err})
	} else {
		out.CPU = cpuStats
	}

	return response.SyncResponse(true, &out)
}

func getCPUMetrics(d *Daemon) (map[string]metrics.CPUMetrics, error) {
	stats, err := ioutil.ReadFile("/proc/stat")
	if err != nil {
		return nil, fmt.Errorf("Failed to read /proc/stat: %w", err)
	}

	out := map[string]metrics.CPUMetrics{}
	scanner := bufio.NewScanner(bytes.NewReader(stats))

	for scanner.Scan() {
		fields := strings.Fields(scanner.Text())

		// Only consider CPU info, skip everything else.
		if !strings.HasPrefix(fields[0], "cpu") {
			continue
		}

		stats := metrics.CPUMetrics{}

		stats.SecondsUser, err = strconv.ParseUint(fields[1], 10, 64)
		if err != nil {
			return nil, fmt.Errorf("Failed to parse %q: %w", fields[1], err)
		}

		stats.SecondsUser *= 10

		stats.SecondsNice, err = strconv.ParseUint(fields[2], 10, 64)
		if err != nil {
			return nil, fmt.Errorf("Failed to parse %q: %w", fields[2], err)
		}

		stats.SecondsNice *= 10

		stats.SecondsSystem, err = strconv.ParseUint(fields[3], 10, 64)
		if err != nil {
			return nil, fmt.Errorf("Failed to parse %q: %w", fields[3], err)
		}

		stats.SecondsSystem *= 10

		stats.SecondsIdle, err = strconv.ParseUint(fields[4], 10, 64)
		if err != nil {
			return nil, fmt.Errorf("Failed to parse %q: %w", fields[4], err)
		}

		stats.SecondsIdle *= 10

		stats.SecondsIOWait, err = strconv.ParseUint(fields[5], 10, 64)
		if err != nil {
			return nil, fmt.Errorf("Failed to parse %q: %w", fields[5], err)
		}

		stats.SecondsIdle *= 10

		stats.SecondsIRQ, err = strconv.ParseUint(fields[6], 10, 64)
		if err != nil {
			return nil, fmt.Errorf("Failed to parse %q: %w", fields[6], err)
		}

		stats.SecondsIRQ *= 10

		stats.SecondsSoftIRQ, err = strconv.ParseUint(fields[7], 10, 64)
		if err != nil {
			return nil, fmt.Errorf("Failed to parse %q: %w", fields[7], err)
		}

		stats.SecondsSoftIRQ *= 10

		stats.SecondsSteal, err = strconv.ParseUint(fields[8], 10, 64)
		if err != nil {
			return nil, fmt.Errorf("Failed to parse %q: %w", fields[8], err)
		}

		stats.SecondsSteal *= 10

		out[fields[0]] = stats
	}

	return out, nil
}

func getTotalProcesses(d *Daemon) (uint64, error) {
	entries, err := ioutil.ReadDir("/proc")
	if err != nil {
		return 0, fmt.Errorf("Failed to read dir %q: %w", "/proc", err)
	}

	re := regexp.MustCompile(`^[[:digit:]]+$`)
	pidCount := uint64(0)

	for _, entry := range entries {
		// Skip everything which isn't a directory
		if !entry.IsDir() {
			continue
		}

		// Skip all non-PID directories
		if !re.MatchString(entry.Name()) {
			continue
		}

		cmdlinePath := filepath.Join("/proc", entry.Name(), "cmdline")

		cmdline, err := ioutil.ReadFile(cmdlinePath)
		if err != nil {
			continue
		}

		if string(cmdline) == "" {
			continue
		}

		pidCount++
	}

	return pidCount, nil
}

func getDiskMetrics(d *Daemon) (map[string]metrics.DiskMetrics, error) {
	diskStats, err := ioutil.ReadFile("/proc/diskstats")
	if err != nil {
		return nil, fmt.Errorf("Failed to read /proc/diskstats: %w", err)
	}

	out := map[string]metrics.DiskMetrics{}
	scanner := bufio.NewScanner(bytes.NewReader(diskStats))

	for scanner.Scan() {
		line := scanner.Text()

		if line == "" {
			continue
		}

		fields := strings.Fields(line)
		stats := metrics.DiskMetrics{}

		stats.ReadsCompleted, err = strconv.ParseUint(fields[3], 10, 64)
		if err != nil {
			return nil, fmt.Errorf("Failed to parse %q: %w", fields[3], err)
		}

		sectorsRead, err := strconv.ParseUint(fields[5], 10, 64)
		if err != nil {
			return nil, fmt.Errorf("Failed to parse %q: %w", fields[3], err)
		}

		stats.ReadBytes = sectorsRead * 512

		stats.WritesCompleted, err = strconv.ParseUint(fields[7], 10, 64)
		if err != nil {
			return nil, fmt.Errorf("Failed to parse %q: %w", fields[3], err)
		}

		sectorsWritten, err := strconv.ParseUint(fields[9], 10, 64)
		if err != nil {
			return nil, fmt.Errorf("Failed to parse %q: %w", fields[3], err)
		}

		stats.WrittenBytes = sectorsWritten * 512

		out[fields[2]] = stats
	}

	return out, nil
}

func getFilesystemMetrics(d *Daemon) (map[string]metrics.FilesystemMetrics, error) {
	mounts, err := ioutil.ReadFile("/proc/mounts")
	if err != nil {
		return nil, fmt.Errorf("Failed to read /proc/mounts: %w", err)
	}

	out := map[string]metrics.FilesystemMetrics{}
	scanner := bufio.NewScanner(bytes.NewReader(mounts))

	for scanner.Scan() {
		fields := strings.Fields(scanner.Text())

		// Skip uninteresting mounts
		if shared.StringInSlice(fields[2], defFSTypesExcluded) || defMountPointsExcluded.MatchString(fields[1]) {
			continue
		}

		stats := metrics.FilesystemMetrics{}

		stats.Mountpoint = fields[1]

		statfs, err := filesystem.StatVFS(stats.Mountpoint)
		if err != nil {
			return nil, fmt.Errorf("Failed to stat %s: %w", stats.Mountpoint, err)
		}

		fsType, err := filesystem.FSTypeToName(int32(statfs.Type))
		if err == nil {
			stats.FSType = fsType
		}

		stats.SizeBytes = statfs.Blocks * uint64(statfs.Bsize)
		stats.SizeBytes = statfs.Blocks * uint64(statfs.Bsize)
		stats.SizeBytes = statfs.Blocks * uint64(statfs.Bsize)
		stats.SizeBytes = statfs.Blocks * uint64(statfs.Bsize)

		out[fields[0]] = stats
	}

	return out, nil
}

func getMemoryMetrics(d *Daemon) (metrics.MemoryMetrics, error) {
	content, err := ioutil.ReadFile("/proc/meminfo")
	if err != nil {
		return metrics.MemoryMetrics{}, fmt.Errorf("Failed to read /proc/meminfo: %w", err)
	}

	out := metrics.MemoryMetrics{}
	scanner := bufio.NewScanner(bytes.NewReader(content))

	for scanner.Scan() {
		fields := strings.Fields(scanner.Text())

		fields[0] = strings.TrimRight(fields[0], ":")

		value, err := strconv.ParseUint(fields[1], 10, 64)
		if err != nil {
			return metrics.MemoryMetrics{}, fmt.Errorf("Failed to parse %q: %w", fields[1], err)
		}

		// Multiply suffix (kB)
		if len(fields) == 3 {
			value *= 1024
		}

		// FIXME: Missing RSS
		switch fields[0] {
		case "Active":
			out.ActiveBytes = value
		case "Active(anon)":
			out.ActiveAnonBytes = value
		case "Active(file)":
			out.ActiveFileBytes = value
		case "Cached":
			out.CachedBytes = value
		case "Dirty":
			out.DirtyBytes = value
		case "HugePages_Free":
			out.HugepagesFreeBytes = value
		case "HugePages_Total":
			out.HugepagesTotalBytes = value
		case "Inactive":
			out.InactiveBytes = value
		case "Inactive(anon)":
			out.InactiveAnonBytes = value
		case "Inactive(file)":
			out.InactiveFileBytes = value
		case "Mapped":
			out.MappedBytes = value
		case "MemAvailable":
			out.MemAvailableBytes = value
		case "MemFree":
			out.MemFreeBytes = value
		case "MemTotal":
			out.MemTotalBytes = value
		case "Shmem":
			out.ShmemBytes = value
		case "SwapCached":
			out.SwapBytes = value
		case "Unevictable":
			out.UnevictableBytes = value
		case "Writeback":
			out.WritebackBytes = value
		}
	}

	return out, nil
}

func getNetworkMetrics(d *Daemon) (map[string]metrics.NetworkMetrics, error) {
	out := map[string]metrics.NetworkMetrics{}

	for dev, state := range networkState() {
		stats := metrics.NetworkMetrics{}

		stats.ReceiveBytes = uint64(state.Counters.BytesReceived)
		stats.ReceiveDrop = uint64(state.Counters.PacketsDroppedInbound)
		stats.ReceiveErrors = uint64(state.Counters.ErrorsReceived)
		stats.ReceivePackets = uint64(state.Counters.PacketsReceived)
		stats.TransmitBytes = uint64(state.Counters.BytesReceived)
		stats.TransmitDrop = uint64(state.Counters.PacketsDroppedOutbound)
		stats.TransmitErrors = uint64(state.Counters.ErrorsSent)
		stats.TransmitPackets = uint64(state.Counters.PacketsSent)

		out[dev] = stats
	}

	return out, nil
}
<|endoftext|>


<commit_before>package main

import (
	"encoding/json"
	"fmt"
	"io"
	"mime"
	"mime/multipart"
	"os"
	"path/filepath"
	"sync"
	"time"

	"github.com/lxc/lxd/shared"

	log "gopkg.in/inconshreveable/log15.v2"
)

type imageStreamCacheEntry struct {
	ss     *shared.SimpleStreams
	expiry time.Time
}

var imageStreamCache = map[string]*imageStreamCacheEntry{}
var imageStreamCacheLock sync.Mutex

// ImageDownload checks if we have that Image Fingerprint else
// downloads the image from a remote server.
func (d *Daemon) ImageDownload(op *operation, server string, protocol string, certificate string, secret string, alias string, forContainer bool, autoUpdate bool) (string, error) {
	var err error
	var ss *shared.SimpleStreams

	if protocol == "" {
		protocol = "lxd"
	}

	fp := alias

	// Expand aliases
	if protocol == "simplestreams" {
		imageStreamCacheLock.Lock()
		entry, _ := imageStreamCache[server]
		if entry == nil || entry.expiry.Before(time.Now()) {
			ss, err = shared.SimpleStreamsClient(server, d.proxy)
			if err != nil {
				imageStreamCacheLock.Unlock()
				return "", err
			}

			entry = &imageStreamCacheEntry{ss: ss, expiry: time.Now().Add(time.Hour)}
			imageStreamCache[server] = entry
		} else {
			shared.LogDebugf("Using SimpleStreams cache entry for %s, expires at %s", server, entry.expiry)
			ss = entry.ss
		}
		imageStreamCacheLock.Unlock()

		target := ss.GetAlias(fp)
		if target != "" {
			fp = target
		}

		image, err := ss.GetImageInfo(fp)
		if err != nil {
			return "", err
		}

		if fp == alias {
			alias = image.Fingerprint
		}
		fp = image.Fingerprint
	} else if protocol == "lxd" {
		target, err := remoteGetImageFingerprint(d, server, certificate, fp)
		if err == nil && target != "" {
			fp = target
		}
	}

	if _, _, err := dbImageGet(d.db, fp, false, false); err == nil {
		shared.LogDebug("Image already exists in the db", log.Ctx{"image": fp})
		// already have it
		return fp, nil
	}

	shared.LogInfo(
		"Image not in the db, downloading it",
		log.Ctx{"image": fp, "server": server})

	// Now check if we already downloading the image
	d.imagesDownloadingLock.RLock()
	if waitChannel, ok := d.imagesDownloading[fp]; ok {
		// We already download the image
		d.imagesDownloadingLock.RUnlock()

		shared.LogDebug(
			"Already downloading the image, waiting for it to succeed",
			log.Ctx{"image": fp})

		// Wait until the download finishes (channel closes)
		if _, ok := <-waitChannel; ok {
			shared.LogWarnf("Value transmitted over image lock semaphore?")
		}

		if _, _, err := dbImageGet(d.db, fp, false, true); err != nil {
			shared.LogError(
				"Previous download didn't succeed",
				log.Ctx{"image": fp})

			return "", fmt.Errorf("Previous download didn't succeed")
		}

		shared.LogDebug(
			"Previous download succeeded",
			log.Ctx{"image": fp})

		return fp, nil
	}

	d.imagesDownloadingLock.RUnlock()

	shared.LogInfo(
		"Downloading the image",
		log.Ctx{"image": fp})

	// Add the download to the queue
	d.imagesDownloadingLock.Lock()
	d.imagesDownloading[fp] = make(chan bool)
	d.imagesDownloadingLock.Unlock()

	// Unlock once this func ends.
	defer func() {
		d.imagesDownloadingLock.Lock()
		if waitChannel, ok := d.imagesDownloading[fp]; ok {
			close(waitChannel)
			delete(d.imagesDownloading, fp)
		}
		d.imagesDownloadingLock.Unlock()
	}()

	exporturl := server

	var info shared.ImageInfo
	info.Fingerprint = fp

	destDir := shared.VarPath("images")
	destName := filepath.Join(destDir, fp)
	if shared.PathExists(destName) {
		d.Storage.ImageDelete(fp)
	}

	progress := func(progressInt int) {
		if op == nil {
			return
		}

		meta := op.metadata
		if meta == nil {
			meta = make(map[string]interface{})
		}

		progress := fmt.Sprintf("%d%%", progressInt)

		if meta["download_progress"] != progress {
			meta["download_progress"] = progress
			op.UpdateMetadata(meta)
		}
	}

	if protocol == "lxd" {
		/* grab the metadata from /1.0/images/%s */
		var url string
		if secret != "" {
			url = fmt.Sprintf(
				"%s/%s/images/%s?secret=%s",
				server, shared.APIVersion, fp, secret)
		} else {
			url = fmt.Sprintf("%s/%s/images/%s", server, shared.APIVersion, fp)
		}

		resp, err := d.httpGetSync(url, certificate)
		if err != nil {
			shared.LogError(
				"Failed to download image metadata",
				log.Ctx{"image": fp, "err": err})

			return "", err
		}

		if err := json.Unmarshal(resp.Metadata, &info); err != nil {
			return "", err
		}

		/* now grab the actual file from /1.0/images/%s/export */
		if secret != "" {
			exporturl = fmt.Sprintf(
				"%s/%s/images/%s/export?secret=%s",
				server, shared.APIVersion, fp, secret)

		} else {
			exporturl = fmt.Sprintf(
				"%s/%s/images/%s/export",
				server, shared.APIVersion, fp)
		}
	} else if protocol == "simplestreams" {
		err := ss.Download(fp, "meta", destName, nil)
		if err != nil {
			return "", err
		}

		err = ss.Download(fp, "root", destName+".rootfs", progress)
		if err != nil {
			return "", err
		}

		info, err := ss.GetImageInfo(fp)
		if err != nil {
			return "", err
		}

		info.Public = false
		info.AutoUpdate = autoUpdate

		_, err = imageBuildFromInfo(d, *info)
		if err != nil {
			return "", err
		}

		if alias != fp {
			id, _, err := dbImageGet(d.db, fp, false, true)
			if err != nil {
				return "", err
			}

			err = dbImageSourceInsert(d.db, id, server, protocol, "", alias)
			if err != nil {
				return "", err
			}
		}

		if forContainer {
			return fp, dbImageLastAccessInit(d.db, fp)
		}

		return fp, nil
	}

	raw, err := d.httpGetFile(exporturl, certificate)
	if err != nil {
		shared.LogError(
			"Failed to download image",
			log.Ctx{"image": fp, "err": err})
		return "", err
	}
	info.Size = raw.ContentLength

	ctype, ctypeParams, err := mime.ParseMediaType(raw.Header.Get("Content-Type"))
	if err != nil {
		ctype = "application/octet-stream"
	}

	body := &shared.TransferProgress{Reader: raw.Body, Length: raw.ContentLength, Handler: progress}

	if ctype == "multipart/form-data" {
		// Parse the POST data
		mr := multipart.NewReader(body, ctypeParams["boundary"])

		// Get the metadata tarball
		part, err := mr.NextPart()
		if err != nil {
			shared.LogError(
				"Invalid multipart image",
				log.Ctx{"image": fp, "err": err})

			return "", err
		}

		if part.FormName() != "metadata" {
			shared.LogError(
				"Invalid multipart image",
				log.Ctx{"image": fp, "err": err})

			return "", fmt.Errorf("Invalid multipart image")
		}

		destName = filepath.Join(destDir, info.Fingerprint)
		f, err := os.Create(destName)
		if err != nil {
			shared.LogError(
				"Failed to save image",
				log.Ctx{"image": fp, "err": err})

			return "", err
		}

		_, err = io.Copy(f, part)
		f.Close()

		if err != nil {
			shared.LogError(
				"Failed to save image",
				log.Ctx{"image": fp, "err": err})

			return "", err
		}

		// Get the rootfs tarball
		part, err = mr.NextPart()
		if err != nil {
			shared.LogError(
				"Invalid multipart image",
				log.Ctx{"image": fp, "err": err})

			return "", err
		}

		if part.FormName() != "rootfs" {
			shared.LogError(
				"Invalid multipart image",
				log.Ctx{"image": fp})
			return "", fmt.Errorf("Invalid multipart image")
		}

		destName = filepath.Join(destDir, info.Fingerprint+".rootfs")
		f, err = os.Create(destName)
		if err != nil {
			shared.LogError(
				"Failed to save image",
				log.Ctx{"image": fp, "err": err})
			return "", err
		}

		_, err = io.Copy(f, part)
		f.Close()

		if err != nil {
			shared.LogError(
				"Failed to save image",
				log.Ctx{"image": fp, "err": err})
			return "", err
		}
	} else {
		destName = filepath.Join(destDir, info.Fingerprint)

		f, err := os.Create(destName)
		if err != nil {
			shared.LogError(
				"Failed to save image",
				log.Ctx{"image": fp, "err": err})

			return "", err
		}

		_, err = io.Copy(f, body)
		f.Close()

		if err != nil {
			shared.LogError(
				"Failed to save image",
				log.Ctx{"image": fp, "err": err})
			return "", err
		}
	}

	if protocol == "direct" {
		imageMeta, err := getImageMetadata(destName)
		if err != nil {
			return "", err
		}

		info.Architecture = imageMeta.Architecture
		info.CreationDate = time.Unix(imageMeta.CreationDate, 0)
		info.ExpiryDate = time.Unix(imageMeta.ExpiryDate, 0)
		info.Properties = imageMeta.Properties
	}

	// By default, make all downloaded images private
	info.Public = false

	if alias != fp && secret == "" {
		info.AutoUpdate = autoUpdate
	}

	_, err = imageBuildFromInfo(d, info)
	if err != nil {
		shared.LogError(
			"Failed to create image",
			log.Ctx{"image": fp, "err": err})

		return "", err
	}

	if alias != fp {
		id, _, err := dbImageGet(d.db, fp, false, true)
		if err != nil {
			return "", err
		}

		err = dbImageSourceInsert(d.db, id, server, protocol, "", alias)
		if err != nil {
			return "", err
		}
	}

	shared.LogDebug(
		"Download succeeded",
		log.Ctx{"image": fp})

	if forContainer {
		return fp, dbImageLastAccessInit(d.db, fp)
	}

	return fp, nil
}
<commit_msg>lxd/daemon_images: improve log for image downloads<commit_after>package main

import (
	"encoding/json"
	"fmt"
	"io"
	"mime"
	"mime/multipart"
	"os"
	"path/filepath"
	"sync"
	"time"

	"github.com/lxc/lxd/shared"

	log "gopkg.in/inconshreveable/log15.v2"
)

type imageStreamCacheEntry struct {
	ss     *shared.SimpleStreams
	expiry time.Time
}

var imageStreamCache = map[string]*imageStreamCacheEntry{}
var imageStreamCacheLock sync.Mutex

// ImageDownload checks if we have that Image Fingerprint else
// downloads the image from a remote server.
func (d *Daemon) ImageDownload(op *operation, server string, protocol string, certificate string, secret string, alias string, forContainer bool, autoUpdate bool) (string, error) {
	var err error
	var ss *shared.SimpleStreams

	if protocol == "" {
		protocol = "lxd"
	}

	fp := alias

	// Expand aliases
	if protocol == "simplestreams" {
		imageStreamCacheLock.Lock()
		entry, _ := imageStreamCache[server]
		if entry == nil || entry.expiry.Before(time.Now()) {
			ss, err = shared.SimpleStreamsClient(server, d.proxy)
			if err != nil {
				imageStreamCacheLock.Unlock()
				return "", err
			}

			entry = &imageStreamCacheEntry{ss: ss, expiry: time.Now().Add(time.Hour)}
			imageStreamCache[server] = entry
		} else {
			shared.LogDebugf("Using SimpleStreams cache entry for %s, expires at %s", server, entry.expiry)
			ss = entry.ss
		}
		imageStreamCacheLock.Unlock()

		target := ss.GetAlias(fp)
		if target != "" {
			fp = target
		}

		image, err := ss.GetImageInfo(fp)
		if err != nil {
			return "", err
		}

		if fp == alias {
			alias = image.Fingerprint
		}
		fp = image.Fingerprint
	} else if protocol == "lxd" {
		target, err := remoteGetImageFingerprint(d, server, certificate, fp)
		if err == nil && target != "" {
			fp = target
		}
	}

	if _, _, err := dbImageGet(d.db, fp, false, false); err == nil {
		shared.LogDebug("Image already exists in the db", log.Ctx{"image": fp})
		// already have it
		return fp, nil
	}

	// Now check if we already downloading the image
	d.imagesDownloadingLock.RLock()
	if waitChannel, ok := d.imagesDownloading[fp]; ok {
		// We already download the image
		d.imagesDownloadingLock.RUnlock()

		shared.LogDebug(
			"Already downloading the image, waiting for it to succeed",
			log.Ctx{"image": fp})

		// Wait until the download finishes (channel closes)
		if _, ok := <-waitChannel; ok {
			shared.LogWarnf("Value transmitted over image lock semaphore?")
		}

		if _, _, err := dbImageGet(d.db, fp, false, true); err != nil {
			shared.LogError(
				"Previous download didn't succeed",
				log.Ctx{"image": fp})

			return "", fmt.Errorf("Previous download didn't succeed")
		}

		shared.LogDebug(
			"Previous download succeeded",
			log.Ctx{"image": fp})

		return fp, nil
	}

	d.imagesDownloadingLock.RUnlock()

	shared.LogInfo(
		"Downloading image",
		log.Ctx{"trigger": op.url, "image": fp, "operation": op.id, "alias": alias, "server": server})

	// Add the download to the queue
	d.imagesDownloadingLock.Lock()
	d.imagesDownloading[fp] = make(chan bool)
	d.imagesDownloadingLock.Unlock()

	// Unlock once this func ends.
	defer func() {
		d.imagesDownloadingLock.Lock()
		if waitChannel, ok := d.imagesDownloading[fp]; ok {
			close(waitChannel)
			delete(d.imagesDownloading, fp)
		}
		d.imagesDownloadingLock.Unlock()
	}()

	exporturl := server

	var info shared.ImageInfo
	info.Fingerprint = fp

	destDir := shared.VarPath("images")
	destName := filepath.Join(destDir, fp)
	if shared.PathExists(destName) {
		d.Storage.ImageDelete(fp)
	}

	progress := func(progressInt int) {
		if op == nil {
			return
		}

		meta := op.metadata
		if meta == nil {
			meta = make(map[string]interface{})
		}

		progress := fmt.Sprintf("%d%%", progressInt)

		if meta["download_progress"] != progress {
			meta["download_progress"] = progress
			op.UpdateMetadata(meta)
		}
	}

	if protocol == "lxd" {
		/* grab the metadata from /1.0/images/%s */
		var url string
		if secret != "" {
			url = fmt.Sprintf(
				"%s/%s/images/%s?secret=%s",
				server, shared.APIVersion, fp, secret)
		} else {
			url = fmt.Sprintf("%s/%s/images/%s", server, shared.APIVersion, fp)
		}

		resp, err := d.httpGetSync(url, certificate)
		if err != nil {
			shared.LogError(
				"Failed to download image metadata",
				log.Ctx{"image": fp, "err": err})

			return "", err
		}

		if err := json.Unmarshal(resp.Metadata, &info); err != nil {
			return "", err
		}

		/* now grab the actual file from /1.0/images/%s/export */
		if secret != "" {
			exporturl = fmt.Sprintf(
				"%s/%s/images/%s/export?secret=%s",
				server, shared.APIVersion, fp, secret)

		} else {
			exporturl = fmt.Sprintf(
				"%s/%s/images/%s/export",
				server, shared.APIVersion, fp)
		}
	} else if protocol == "simplestreams" {
		err := ss.Download(fp, "meta", destName, nil)
		if err != nil {
			return "", err
		}

		err = ss.Download(fp, "root", destName+".rootfs", progress)
		if err != nil {
			return "", err
		}

		info, err := ss.GetImageInfo(fp)
		if err != nil {
			return "", err
		}

		info.Public = false
		info.AutoUpdate = autoUpdate

		_, err = imageBuildFromInfo(d, *info)
		if err != nil {
			return "", err
		}

		if alias != fp {
			id, _, err := dbImageGet(d.db, fp, false, true)
			if err != nil {
				return "", err
			}

			err = dbImageSourceInsert(d.db, id, server, protocol, "", alias)
			if err != nil {
				return "", err
			}
		}

		shared.LogInfo(
			"Image downloaded",
			log.Ctx{"image": fp, "fingerprint": fp, "operation": op.id, "alias": alias, "server": server})

		if forContainer {
			return fp, dbImageLastAccessInit(d.db, fp)
		}

		return fp, nil
	}

	raw, err := d.httpGetFile(exporturl, certificate)
	if err != nil {
		shared.LogError(
			"Failed to download image",
			log.Ctx{"image": fp, "err": err})
		return "", err
	}
	info.Size = raw.ContentLength

	ctype, ctypeParams, err := mime.ParseMediaType(raw.Header.Get("Content-Type"))
	if err != nil {
		ctype = "application/octet-stream"
	}

	body := &shared.TransferProgress{Reader: raw.Body, Length: raw.ContentLength, Handler: progress}

	if ctype == "multipart/form-data" {
		// Parse the POST data
		mr := multipart.NewReader(body, ctypeParams["boundary"])

		// Get the metadata tarball
		part, err := mr.NextPart()
		if err != nil {
			shared.LogError(
				"Invalid multipart image",
				log.Ctx{"image": fp, "err": err})

			return "", err
		}

		if part.FormName() != "metadata" {
			shared.LogError(
				"Invalid multipart image",
				log.Ctx{"image": fp, "err": err})

			return "", fmt.Errorf("Invalid multipart image")
		}

		destName = filepath.Join(destDir, info.Fingerprint)
		f, err := os.Create(destName)
		if err != nil {
			shared.LogError(
				"Failed to save image",
				log.Ctx{"image": fp, "err": err})

			return "", err
		}

		_, err = io.Copy(f, part)
		f.Close()

		if err != nil {
			shared.LogError(
				"Failed to save image",
				log.Ctx{"image": fp, "err": err})

			return "", err
		}

		// Get the rootfs tarball
		part, err = mr.NextPart()
		if err != nil {
			shared.LogError(
				"Invalid multipart image",
				log.Ctx{"image": fp, "err": err})

			return "", err
		}

		if part.FormName() != "rootfs" {
			shared.LogError(
				"Invalid multipart image",
				log.Ctx{"image": fp})
			return "", fmt.Errorf("Invalid multipart image")
		}

		destName = filepath.Join(destDir, info.Fingerprint+".rootfs")
		f, err = os.Create(destName)
		if err != nil {
			shared.LogError(
				"Failed to save image",
				log.Ctx{"image": fp, "err": err})
			return "", err
		}

		_, err = io.Copy(f, part)
		f.Close()

		if err != nil {
			shared.LogError(
				"Failed to save image",
				log.Ctx{"image": fp, "err": err})
			return "", err
		}
	} else {
		destName = filepath.Join(destDir, info.Fingerprint)

		f, err := os.Create(destName)
		if err != nil {
			shared.LogError(
				"Failed to save image",
				log.Ctx{"image": fp, "err": err})

			return "", err
		}

		_, err = io.Copy(f, body)
		f.Close()

		if err != nil {
			shared.LogError(
				"Failed to save image",
				log.Ctx{"image": fp, "err": err})
			return "", err
		}
	}

	if protocol == "direct" {
		imageMeta, err := getImageMetadata(destName)
		if err != nil {
			return "", err
		}

		info.Architecture = imageMeta.Architecture
		info.CreationDate = time.Unix(imageMeta.CreationDate, 0)
		info.ExpiryDate = time.Unix(imageMeta.ExpiryDate, 0)
		info.Properties = imageMeta.Properties
	}

	// By default, make all downloaded images private
	info.Public = false

	if alias != fp && secret == "" {
		info.AutoUpdate = autoUpdate
	}

	_, err = imageBuildFromInfo(d, info)
	if err != nil {
		shared.LogError(
			"Failed to create image",
			log.Ctx{"image": fp, "err": err})

		return "", err
	}

	if alias != fp {
		id, _, err := dbImageGet(d.db, fp, false, true)
		if err != nil {
			return "", err
		}

		err = dbImageSourceInsert(d.db, id, server, protocol, "", alias)
		if err != nil {
			return "", err
		}
	}

	shared.LogInfo(
		"Image downloaded",
		log.Ctx{"image": fp, "operation": op.id, "alias": alias, "server": server})

	if forContainer {
		return fp, dbImageLastAccessInit(d.db, fp)
	}

	return fp, nil
}
<|endoftext|>


<commit_before>package main

import (
	"encoding/json"
	"fmt"
	"io"
	"mime"
	"mime/multipart"
	"os"
	"path/filepath"
	"time"

	"github.com/lxc/lxd/shared"

	log "gopkg.in/inconshreveable/log15.v2"
)

// ImageDownload checks if we have that Image Fingerprint else
// downloads the image from a remote server.
func (d *Daemon) ImageDownload(op *operation, server string, protocol string, certificate string, secret string, alias string, forContainer bool, autoUpdate bool) (string, error) {
	var err error
	var ss *shared.SimpleStreams
	fp := alias

	// Expand aliases
	if protocol == "simplestreams" {
		ss, err = shared.SimpleStreamsClient(server)
		if err != nil {
			return "", err
		}

		target := ss.GetAlias(fp)
		if target != "" {
			fp = target
		}

		image, err := ss.GetImageInfo(fp)
		if err != nil {
			return "", err
		}

		if fp == alias {
			alias = image.Fingerprint
		}
		fp = image.Fingerprint
	} else if protocol == "lxd" {
		target, err := remoteGetImageFingerprint(d, server, certificate, fp)
		if err == nil && target != "" {
			fp = target
		}
	}

	if _, _, err := dbImageGet(d.db, fp, false, false); err == nil {
		shared.Log.Debug("Image already exists in the db", log.Ctx{"image": fp})
		// already have it
		return fp, nil
	}

	shared.Log.Info(
		"Image not in the db, downloading it",
		log.Ctx{"image": fp, "server": server})

	// Now check if we already downloading the image
	d.imagesDownloadingLock.RLock()
	if waitChannel, ok := d.imagesDownloading[fp]; ok {
		// We already download the image
		d.imagesDownloadingLock.RUnlock()

		shared.Log.Info(
			"Already downloading the image, waiting for it to succeed",
			log.Ctx{"image": fp})

		// Wait until the download finishes (channel closes)
		if _, ok := <-waitChannel; ok {
			shared.Log.Warn("Value transmitted over image lock semaphore?")
		}

		if _, _, err := dbImageGet(d.db, fp, false, true); err != nil {
			shared.Log.Error(
				"Previous download didn't succeed",
				log.Ctx{"image": fp})

			return "", fmt.Errorf("Previous download didn't succeed")
		}

		shared.Log.Info(
			"Previous download succeeded",
			log.Ctx{"image": fp})

		return fp, nil
	}

	d.imagesDownloadingLock.RUnlock()

	shared.Log.Info(
		"Downloading the image",
		log.Ctx{"image": fp})

	// Add the download to the queue
	d.imagesDownloadingLock.Lock()
	d.imagesDownloading[fp] = make(chan bool)
	d.imagesDownloadingLock.Unlock()

	// Unlock once this func ends.
	defer func() {
		d.imagesDownloadingLock.Lock()
		if waitChannel, ok := d.imagesDownloading[fp]; ok {
			close(waitChannel)
			delete(d.imagesDownloading, fp)
		}
		d.imagesDownloadingLock.Unlock()
	}()

	exporturl := server

	var info shared.ImageInfo
	info.Fingerprint = fp

	destDir := shared.VarPath("images")
	destName := filepath.Join(destDir, fp)
	if shared.PathExists(destName) {
		d.Storage.ImageDelete(fp)
	}

	progress := func(progressInt int) {
		if op == nil {
			return
		}

		meta := op.metadata
		if meta == nil {
			meta = make(map[string]interface{})
		}

		progress := fmt.Sprintf("%d%%", progressInt)

		if meta["download_progress"] != progress {
			meta["download_progress"] = progress
			op.UpdateMetadata(meta)
		}
	}

	if protocol == "" || protocol == "lxc" {
		/* grab the metadata from /1.0/images/%s */
		var url string
		if secret != "" {
			url = fmt.Sprintf(
				"%s/%s/images/%s?secret=%s",
				server, shared.APIVersion, fp, secret)
		} else {
			url = fmt.Sprintf("%s/%s/images/%s", server, shared.APIVersion, fp)
		}

		resp, err := d.httpGetSync(url, certificate)
		if err != nil {
			shared.Log.Error(
				"Failed to download image metadata",
				log.Ctx{"image": fp, "err": err})

			return "", err
		}

		if err := json.Unmarshal(resp.Metadata, &info); err != nil {
			return "", err
		}

		/* now grab the actual file from /1.0/images/%s/export */
		if secret != "" {
			exporturl = fmt.Sprintf(
				"%s/%s/images/%s/export?secret=%s",
				server, shared.APIVersion, fp, secret)

		} else {
			exporturl = fmt.Sprintf(
				"%s/%s/images/%s/export",
				server, shared.APIVersion, fp)
		}
	} else if protocol == "simplestreams" {
		err := ss.Download(fp, "meta", destName, nil)
		if err != nil {
			return "", err
		}

		err = ss.Download(fp, "root", destName+".rootfs", progress)
		if err != nil {
			return "", err
		}

		info, err := ss.GetImageInfo(fp)
		if err != nil {
			return "", err
		}

		info.Public = false
		info.AutoUpdate = autoUpdate

		_, err = imageBuildFromInfo(d, *info)
		if err != nil {
			return "", err
		}

		if alias != fp {
			id, _, err := dbImageGet(d.db, fp, false, true)
			if err != nil {
				return "", err
			}

			err = dbImageSourceInsert(d.db, id, server, protocol, "", alias)
			if err != nil {
				return "", err
			}
		}

		if forContainer {
			return fp, dbImageLastAccessInit(d.db, fp)
		}

		return fp, nil
	}

	raw, err := d.httpGetFile(exporturl, certificate)
	if err != nil {
		shared.Log.Error(
			"Failed to download image",
			log.Ctx{"image": fp, "err": err})
		return "", err
	}
	info.Size = raw.ContentLength

	ctype, ctypeParams, err := mime.ParseMediaType(raw.Header.Get("Content-Type"))
	if err != nil {
		ctype = "application/octet-stream"
	}

	body := &shared.TransferProgress{Reader: raw.Body, Length: raw.ContentLength, Handler: progress}

	if ctype == "multipart/form-data" {
		// Parse the POST data
		mr := multipart.NewReader(body, ctypeParams["boundary"])

		// Get the metadata tarball
		part, err := mr.NextPart()
		if err != nil {
			shared.Log.Error(
				"Invalid multipart image",
				log.Ctx{"image": fp, "err": err})

			return "", err
		}

		if part.FormName() != "metadata" {
			shared.Log.Error(
				"Invalid multipart image",
				log.Ctx{"image": fp, "err": err})

			return "", fmt.Errorf("Invalid multipart image")
		}

		destName = filepath.Join(destDir, info.Fingerprint)
		f, err := os.Create(destName)
		if err != nil {
			shared.Log.Error(
				"Failed to save image",
				log.Ctx{"image": fp, "err": err})

			return "", err
		}

		_, err = io.Copy(f, part)
		f.Close()

		if err != nil {
			shared.Log.Error(
				"Failed to save image",
				log.Ctx{"image": fp, "err": err})

			return "", err
		}

		// Get the rootfs tarball
		part, err = mr.NextPart()
		if err != nil {
			shared.Log.Error(
				"Invalid multipart image",
				log.Ctx{"image": fp, "err": err})

			return "", err
		}

		if part.FormName() != "rootfs" {
			shared.Log.Error(
				"Invalid multipart image",
				log.Ctx{"image": fp})
			return "", fmt.Errorf("Invalid multipart image")
		}

		destName = filepath.Join(destDir, info.Fingerprint+".rootfs")
		f, err = os.Create(destName)
		if err != nil {
			shared.Log.Error(
				"Failed to save image",
				log.Ctx{"image": fp, "err": err})
			return "", err
		}

		_, err = io.Copy(f, part)
		f.Close()

		if err != nil {
			shared.Log.Error(
				"Failed to save image",
				log.Ctx{"image": fp, "err": err})
			return "", err
		}
	} else {
		destName = filepath.Join(destDir, info.Fingerprint)

		f, err := os.Create(destName)
		if err != nil {
			shared.Log.Error(
				"Failed to save image",
				log.Ctx{"image": fp, "err": err})

			return "", err
		}

		_, err = io.Copy(f, body)
		f.Close()

		if err != nil {
			shared.Log.Error(
				"Failed to save image",
				log.Ctx{"image": fp, "err": err})
			return "", err
		}
	}

	if protocol == "direct" {
		imageMeta, err := getImageMetadata(destName)
		if err != nil {
			return "", err
		}

		info.Architecture = imageMeta.Architecture
		info.CreationDate = time.Unix(imageMeta.CreationDate, 0)
		info.ExpiryDate = time.Unix(imageMeta.ExpiryDate, 0)
		info.Properties = imageMeta.Properties
	}

	// By default, make all downloaded images private
	info.Public = false

	if alias != fp {
		id, _, err := dbImageGet(d.db, fp, false, true)
		if err != nil {
			return "", err
		}

		err = dbImageSourceInsert(d.db, id, server, protocol, "", alias)
		if err != nil {
			return "", err
		}
	}

	_, err = imageBuildFromInfo(d, info)
	if err != nil {
		shared.Log.Error(
			"Failed to create image",
			log.Ctx{"image": fp, "err": err})

		return "", err
	}

	shared.Log.Info(
		"Download succeeded",
		log.Ctx{"image": fp})

	if forContainer {
		return fp, dbImageLastAccessInit(d.db, fp)
	}

	return fp, nil
}
<commit_msg>Fix default protocol in image download<commit_after>package main

import (
	"encoding/json"
	"fmt"
	"io"
	"mime"
	"mime/multipart"
	"os"
	"path/filepath"
	"time"

	"github.com/lxc/lxd/shared"

	log "gopkg.in/inconshreveable/log15.v2"
)

// ImageDownload checks if we have that Image Fingerprint else
// downloads the image from a remote server.
func (d *Daemon) ImageDownload(op *operation, server string, protocol string, certificate string, secret string, alias string, forContainer bool, autoUpdate bool) (string, error) {
	var err error
	var ss *shared.SimpleStreams

	if protocol == "" {
		protocol = "lxd"
	}

	fp := alias

	// Expand aliases
	if protocol == "simplestreams" {
		ss, err = shared.SimpleStreamsClient(server)
		if err != nil {
			return "", err
		}

		target := ss.GetAlias(fp)
		if target != "" {
			fp = target
		}

		image, err := ss.GetImageInfo(fp)
		if err != nil {
			return "", err
		}

		if fp == alias {
			alias = image.Fingerprint
		}
		fp = image.Fingerprint
	} else if protocol == "lxd" {
		target, err := remoteGetImageFingerprint(d, server, certificate, fp)
		if err == nil && target != "" {
			fp = target
		}
	}

	if _, _, err := dbImageGet(d.db, fp, false, false); err == nil {
		shared.Log.Debug("Image already exists in the db", log.Ctx{"image": fp})
		// already have it
		return fp, nil
	}

	shared.Log.Info(
		"Image not in the db, downloading it",
		log.Ctx{"image": fp, "server": server})

	// Now check if we already downloading the image
	d.imagesDownloadingLock.RLock()
	if waitChannel, ok := d.imagesDownloading[fp]; ok {
		// We already download the image
		d.imagesDownloadingLock.RUnlock()

		shared.Log.Info(
			"Already downloading the image, waiting for it to succeed",
			log.Ctx{"image": fp})

		// Wait until the download finishes (channel closes)
		if _, ok := <-waitChannel; ok {
			shared.Log.Warn("Value transmitted over image lock semaphore?")
		}

		if _, _, err := dbImageGet(d.db, fp, false, true); err != nil {
			shared.Log.Error(
				"Previous download didn't succeed",
				log.Ctx{"image": fp})

			return "", fmt.Errorf("Previous download didn't succeed")
		}

		shared.Log.Info(
			"Previous download succeeded",
			log.Ctx{"image": fp})

		return fp, nil
	}

	d.imagesDownloadingLock.RUnlock()

	shared.Log.Info(
		"Downloading the image",
		log.Ctx{"image": fp})

	// Add the download to the queue
	d.imagesDownloadingLock.Lock()
	d.imagesDownloading[fp] = make(chan bool)
	d.imagesDownloadingLock.Unlock()

	// Unlock once this func ends.
	defer func() {
		d.imagesDownloadingLock.Lock()
		if waitChannel, ok := d.imagesDownloading[fp]; ok {
			close(waitChannel)
			delete(d.imagesDownloading, fp)
		}
		d.imagesDownloadingLock.Unlock()
	}()

	exporturl := server

	var info shared.ImageInfo
	info.Fingerprint = fp

	destDir := shared.VarPath("images")
	destName := filepath.Join(destDir, fp)
	if shared.PathExists(destName) {
		d.Storage.ImageDelete(fp)
	}

	progress := func(progressInt int) {
		if op == nil {
			return
		}

		meta := op.metadata
		if meta == nil {
			meta = make(map[string]interface{})
		}

		progress := fmt.Sprintf("%d%%", progressInt)

		if meta["download_progress"] != progress {
			meta["download_progress"] = progress
			op.UpdateMetadata(meta)
		}
	}

	if protocol == "lxd" {
		/* grab the metadata from /1.0/images/%s */
		var url string
		if secret != "" {
			url = fmt.Sprintf(
				"%s/%s/images/%s?secret=%s",
				server, shared.APIVersion, fp, secret)
		} else {
			url = fmt.Sprintf("%s/%s/images/%s", server, shared.APIVersion, fp)
		}

		resp, err := d.httpGetSync(url, certificate)
		if err != nil {
			shared.Log.Error(
				"Failed to download image metadata",
				log.Ctx{"image": fp, "err": err})

			return "", err
		}

		if err := json.Unmarshal(resp.Metadata, &info); err != nil {
			return "", err
		}

		/* now grab the actual file from /1.0/images/%s/export */
		if secret != "" {
			exporturl = fmt.Sprintf(
				"%s/%s/images/%s/export?secret=%s",
				server, shared.APIVersion, fp, secret)

		} else {
			exporturl = fmt.Sprintf(
				"%s/%s/images/%s/export",
				server, shared.APIVersion, fp)
		}
	} else if protocol == "simplestreams" {
		err := ss.Download(fp, "meta", destName, nil)
		if err != nil {
			return "", err
		}

		err = ss.Download(fp, "root", destName+".rootfs", progress)
		if err != nil {
			return "", err
		}

		info, err := ss.GetImageInfo(fp)
		if err != nil {
			return "", err
		}

		info.Public = false
		info.AutoUpdate = autoUpdate

		_, err = imageBuildFromInfo(d, *info)
		if err != nil {
			return "", err
		}

		if alias != fp {
			id, _, err := dbImageGet(d.db, fp, false, true)
			if err != nil {
				return "", err
			}

			err = dbImageSourceInsert(d.db, id, server, protocol, "", alias)
			if err != nil {
				return "", err
			}
		}

		if forContainer {
			return fp, dbImageLastAccessInit(d.db, fp)
		}

		return fp, nil
	}

	raw, err := d.httpGetFile(exporturl, certificate)
	if err != nil {
		shared.Log.Error(
			"Failed to download image",
			log.Ctx{"image": fp, "err": err})
		return "", err
	}
	info.Size = raw.ContentLength

	ctype, ctypeParams, err := mime.ParseMediaType(raw.Header.Get("Content-Type"))
	if err != nil {
		ctype = "application/octet-stream"
	}

	body := &shared.TransferProgress{Reader: raw.Body, Length: raw.ContentLength, Handler: progress}

	if ctype == "multipart/form-data" {
		// Parse the POST data
		mr := multipart.NewReader(body, ctypeParams["boundary"])

		// Get the metadata tarball
		part, err := mr.NextPart()
		if err != nil {
			shared.Log.Error(
				"Invalid multipart image",
				log.Ctx{"image": fp, "err": err})

			return "", err
		}

		if part.FormName() != "metadata" {
			shared.Log.Error(
				"Invalid multipart image",
				log.Ctx{"image": fp, "err": err})

			return "", fmt.Errorf("Invalid multipart image")
		}

		destName = filepath.Join(destDir, info.Fingerprint)
		f, err := os.Create(destName)
		if err != nil {
			shared.Log.Error(
				"Failed to save image",
				log.Ctx{"image": fp, "err": err})

			return "", err
		}

		_, err = io.Copy(f, part)
		f.Close()

		if err != nil {
			shared.Log.Error(
				"Failed to save image",
				log.Ctx{"image": fp, "err": err})

			return "", err
		}

		// Get the rootfs tarball
		part, err = mr.NextPart()
		if err != nil {
			shared.Log.Error(
				"Invalid multipart image",
				log.Ctx{"image": fp, "err": err})

			return "", err
		}

		if part.FormName() != "rootfs" {
			shared.Log.Error(
				"Invalid multipart image",
				log.Ctx{"image": fp})
			return "", fmt.Errorf("Invalid multipart image")
		}

		destName = filepath.Join(destDir, info.Fingerprint+".rootfs")
		f, err = os.Create(destName)
		if err != nil {
			shared.Log.Error(
				"Failed to save image",
				log.Ctx{"image": fp, "err": err})
			return "", err
		}

		_, err = io.Copy(f, part)
		f.Close()

		if err != nil {
			shared.Log.Error(
				"Failed to save image",
				log.Ctx{"image": fp, "err": err})
			return "", err
		}
	} else {
		destName = filepath.Join(destDir, info.Fingerprint)

		f, err := os.Create(destName)
		if err != nil {
			shared.Log.Error(
				"Failed to save image",
				log.Ctx{"image": fp, "err": err})

			return "", err
		}

		_, err = io.Copy(f, body)
		f.Close()

		if err != nil {
			shared.Log.Error(
				"Failed to save image",
				log.Ctx{"image": fp, "err": err})
			return "", err
		}
	}

	if protocol == "direct" {
		imageMeta, err := getImageMetadata(destName)
		if err != nil {
			return "", err
		}

		info.Architecture = imageMeta.Architecture
		info.CreationDate = time.Unix(imageMeta.CreationDate, 0)
		info.ExpiryDate = time.Unix(imageMeta.ExpiryDate, 0)
		info.Properties = imageMeta.Properties
	}

	// By default, make all downloaded images private
	info.Public = false

	if alias != fp {
		id, _, err := dbImageGet(d.db, fp, false, true)
		if err != nil {
			return "", err
		}

		err = dbImageSourceInsert(d.db, id, server, protocol, "", alias)
		if err != nil {
			return "", err
		}
	}

	_, err = imageBuildFromInfo(d, info)
	if err != nil {
		shared.Log.Error(
			"Failed to create image",
			log.Ctx{"image": fp, "err": err})

		return "", err
	}

	shared.Log.Info(
		"Download succeeded",
		log.Ctx{"image": fp})

	if forContainer {
		return fp, dbImageLastAccessInit(d.db, fp)
	}

	return fp, nil
}
<|endoftext|>


<commit_before>//go:build linux && cgo && !agent
// +build linux,cgo,!agent

package db

import (
	"fmt"

	"github.com/lxc/lxd/lxd/db/query"
	"github.com/pkg/errors"
)

// Operation holds information about a single LXD operation running on a node
// in the cluster.
type Operation struct {
	ID          int64         // Stable database identifier
	UUID        string        // User-visible identifier
	NodeAddress string        // Address of the node the operation is running on
	Type        OperationType // Type of the operation
}

// GetLocalOperations returns all operations associated with this node.
func (c *ClusterTx) GetLocalOperations() ([]Operation, error) {
	return c.operations("node_id=?", c.nodeID)
}

// GetLocalOperationsUUIDs returns the UUIDs of all operations associated with this
// node.
func (c *ClusterTx) GetLocalOperationsUUIDs() ([]string, error) {
	stmt := "SELECT uuid FROM operations WHERE node_id=?"
	return query.SelectStrings(c.tx, stmt, c.nodeID)
}

// GetNodesWithRunningOperations returns a list of nodes that have running operations
func (c *ClusterTx) GetNodesWithRunningOperations(project string) ([]string, error) {
	stmt := `
SELECT DISTINCT nodes.address
  FROM operations
  LEFT OUTER JOIN projects ON projects.id = operations.project_id
  JOIN nodes ON nodes.id = operations.node_id
 WHERE projects.name = ? OR operations.project_id IS NULL
`
	return query.SelectStrings(c.tx, stmt, project)
}

// GetOperationWithID returns the operation with the given ID.
func (c *ClusterTx) GetOperationWithID(opID int) (Operation, error) {
	null := Operation{}
	operations, err := c.operations("id=?", opID)
	if err != nil {
		return null, err
	}
	switch len(operations) {
	case 0:
		return null, ErrNoSuchObject
	case 1:
		return operations[0], nil
	default:
		return null, fmt.Errorf("More than one operation matches")
	}
}

// GetOperationByUUID returns the operation with the given UUID.
func (c *ClusterTx) GetOperationByUUID(uuid string) (Operation, error) {
	null := Operation{}
	operations, err := c.operations("uuid=?", uuid)
	if err != nil {
		return null, err
	}
	switch len(operations) {
	case 0:
		return null, ErrNoSuchObject
	case 1:
		return operations[0], nil
	default:
		return null, fmt.Errorf("More than one operation matches")
	}
}

// CreateOperation adds a new operations to the table.
func (c *ClusterTx) CreateOperation(project, uuid string, typ OperationType) (int64, error) {
	var projectID interface{}

	if project != "" {
		var err error
		projectID, err = c.GetProjectID(project)
		if err != nil {
			return -1, errors.Wrap(err, "Fetch project ID")
		}
	} else {
		projectID = nil
	}

	columns := []string{"uuid", "node_id", "type", "project_id"}
	values := []interface{}{uuid, c.nodeID, typ, projectID}
	return query.UpsertObject(c.tx, "operations", columns, values)
}

// RemoveOperation removes the operation with the given UUID.
func (c *ClusterTx) RemoveOperation(uuid string) error {
	result, err := c.tx.Exec("DELETE FROM operations WHERE uuid=?", uuid)
	if err != nil {
		return err
	}
	n, err := result.RowsAffected()
	if err != nil {
		return err
	}
	if n != 1 {
		return fmt.Errorf("query deleted %d rows instead of 1", n)
	}
	return nil
}

// Remove all operations for the given node.
func (c *ClusterTx) removeNodeOperations(nodeID int64) error {
	_, err := c.tx.Exec("DELETE FROM operations WHERE node_id=?", nodeID)
	if err != nil {
		return err
	}

	return nil
}

// Operations returns all operations in the cluster, filtered by the given clause.
func (c *ClusterTx) operations(where string, args ...interface{}) ([]Operation, error) {
	operations := []Operation{}
	dest := func(i int) []interface{} {
		operations = append(operations, Operation{})
		return []interface{}{
			&operations[i].ID,
			&operations[i].UUID,
			&operations[i].NodeAddress,
			&operations[i].Type,
		}
	}
	sql := `
SELECT operations.id, uuid, nodes.address, type FROM operations JOIN nodes ON nodes.id = node_id `
	if where != "" {
		sql += fmt.Sprintf("WHERE %s ", where)
	}
	sql += "ORDER BY operations.id"
	stmt, err := c.tx.Prepare(sql)
	if err != nil {
		return nil, err
	}
	defer stmt.Close()
	err = query.SelectObjects(stmt, dest, args...)
	if err != nil {
		return nil, errors.Wrap(err, "Failed to fetch operations")
	}
	return operations, nil
}
<commit_msg>lxd/db/operations: Adds GetOnlineNodesWithRunningOperationsOfType function<commit_after>//go:build linux && cgo && !agent
// +build linux,cgo,!agent

package db

import (
	"fmt"
	"time"

	"github.com/lxc/lxd/lxd/db/query"
	"github.com/pkg/errors"
)

// Operation holds information about a single LXD operation running on a node
// in the cluster.
type Operation struct {
	ID          int64         // Stable database identifier
	UUID        string        // User-visible identifier
	NodeAddress string        // Address of the node the operation is running on
	Type        OperationType // Type of the operation
}

// GetLocalOperations returns all operations associated with this node.
func (c *ClusterTx) GetLocalOperations() ([]Operation, error) {
	return c.operations("node_id=?", c.nodeID)
}

// GetLocalOperationsUUIDs returns the UUIDs of all operations associated with this
// node.
func (c *ClusterTx) GetLocalOperationsUUIDs() ([]string, error) {
	stmt := "SELECT uuid FROM operations WHERE node_id=?"
	return query.SelectStrings(c.tx, stmt, c.nodeID)
}

// GetNodesWithRunningOperations returns a list of nodes that have running operations
func (c *ClusterTx) GetNodesWithRunningOperations(project string) ([]string, error) {
	stmt := `
SELECT DISTINCT nodes.address
  FROM operations
  LEFT OUTER JOIN projects ON projects.id = operations.project_id
  JOIN nodes ON nodes.id = operations.node_id
 WHERE projects.name = ? OR operations.project_id IS NULL
`
	return query.SelectStrings(c.tx, stmt, project)
}

// GetOnlineNodesWithRunningOperationsOfType returns a list of online nodes that have running operations of type.
func (c *ClusterTx) GetOnlineNodesWithRunningOperationsOfType(projectName string, opType OperationType) ([]string, error) {
	var addresses []string

	offlineThreshold, err := c.GetNodeOfflineThreshold()
	if err != nil {
		return nil, err
	}

	stmt := `
SELECT DISTINCT nodes.address, nodes.heartbeat
  FROM operations
  LEFT OUTER JOIN projects ON projects.id = operations.project_id
  JOIN nodes ON nodes.id = operations.node_id
 WHERE (projects.name = ? OR operations.project_id IS NULL) AND type = ?
`
	rows, err := c.tx.Query(stmt, projectName, opType)
	if err != nil {
		return nil, err
	}
	defer rows.Close()

	for rows.Next() {
		var address string
		var heartbeat time.Time

		err := rows.Scan(&address, &heartbeat)
		if err != nil {
			return nil, err
		}

		if nodeIsOffline(offlineThreshold, heartbeat) {
			continue
		}

		addresses = append(addresses, address)
	}
	if rows.Err() != nil {
		return nil, err
	}

	return query.SelectStrings(c.tx, stmt, projectName)
}

// GetOperationWithID returns the operation with the given ID.
func (c *ClusterTx) GetOperationWithID(opID int) (Operation, error) {
	null := Operation{}
	operations, err := c.operations("id=?", opID)
	if err != nil {
		return null, err
	}
	switch len(operations) {
	case 0:
		return null, ErrNoSuchObject
	case 1:
		return operations[0], nil
	default:
		return null, fmt.Errorf("More than one operation matches")
	}
}

// GetOperationByUUID returns the operation with the given UUID.
func (c *ClusterTx) GetOperationByUUID(uuid string) (Operation, error) {
	null := Operation{}
	operations, err := c.operations("uuid=?", uuid)
	if err != nil {
		return null, err
	}
	switch len(operations) {
	case 0:
		return null, ErrNoSuchObject
	case 1:
		return operations[0], nil
	default:
		return null, fmt.Errorf("More than one operation matches")
	}
}

// CreateOperation adds a new operations to the table.
func (c *ClusterTx) CreateOperation(project, uuid string, typ OperationType) (int64, error) {
	var projectID interface{}

	if project != "" {
		var err error
		projectID, err = c.GetProjectID(project)
		if err != nil {
			return -1, errors.Wrap(err, "Fetch project ID")
		}
	} else {
		projectID = nil
	}

	columns := []string{"uuid", "node_id", "type", "project_id"}
	values := []interface{}{uuid, c.nodeID, typ, projectID}
	return query.UpsertObject(c.tx, "operations", columns, values)
}

// RemoveOperation removes the operation with the given UUID.
func (c *ClusterTx) RemoveOperation(uuid string) error {
	result, err := c.tx.Exec("DELETE FROM operations WHERE uuid=?", uuid)
	if err != nil {
		return err
	}
	n, err := result.RowsAffected()
	if err != nil {
		return err
	}
	if n != 1 {
		return fmt.Errorf("query deleted %d rows instead of 1", n)
	}
	return nil
}

// Remove all operations for the given node.
func (c *ClusterTx) removeNodeOperations(nodeID int64) error {
	_, err := c.tx.Exec("DELETE FROM operations WHERE node_id=?", nodeID)
	if err != nil {
		return err
	}

	return nil
}

// Operations returns all operations in the cluster, filtered by the given clause.
func (c *ClusterTx) operations(where string, args ...interface{}) ([]Operation, error) {
	operations := []Operation{}
	dest := func(i int) []interface{} {
		operations = append(operations, Operation{})
		return []interface{}{
			&operations[i].ID,
			&operations[i].UUID,
			&operations[i].NodeAddress,
			&operations[i].Type,
		}
	}
	sql := `
SELECT operations.id, uuid, nodes.address, type FROM operations JOIN nodes ON nodes.id = node_id `
	if where != "" {
		sql += fmt.Sprintf("WHERE %s ", where)
	}
	sql += "ORDER BY operations.id"
	stmt, err := c.tx.Prepare(sql)
	if err != nil {
		return nil, err
	}
	defer stmt.Close()
	err = query.SelectObjects(stmt, dest, args...)
	if err != nil {
		return nil, errors.Wrap(err, "Failed to fetch operations")
	}
	return operations, nil
}
<|endoftext|>


<commit_before>package mesos

import (
	"encoding/json"
	"net/http"
	"net/url"
)

type State struct {
	ActivatedSlaves int64   `json:"activated_slaves"`
	BuildDate       string  `json:"build_date"`
	BuildTime       float64 `json:"build_time"`
	BuildUser       string  `json:"build_user"`
	// CompletedFrameworks    []Something       `json:"completed_frameworks"`
	DeactivatedSlaves int64             `json:"deactivated_slaves"`
	ElectedTime       float64           `json:"elected_time"`
	FailedTasks       int64             `json:"failed_tasks"`
	FinishedTasks     int64             `json:"finished_tasks"`
	Flags             map[string]string `json:"flags"`
	Frameworks        []Framework       `json:"frameworks"`
	GitSHA            string            `json:"git_sha"`
	GitTag            string            `json:"git_tag"`
	Hostname          string            `json:"hostname"`
	ID                string            `json:"id"`
	KilledTasks       int64             `json:"killed_tasks"`
	Leader            string            `json:"leader"`
	LogDir            string            `json:"log_dir"`
	LostTasks         int64             `json:"lost_tasks"`
	// OrphanTasks            []Something       `json:"orphan_tasks"`
	PID          string  `json:"pid"`
	Slaves       []Slave `json:"slaves"`
	StagedTasks  int64   `json:"staged_tasks"`
	StartTime    float64 `json:"start_time"`
	StartedTasks int64   `json:"started_tasks"`
	// UnregisteredFrameworks []Something       `json:"unregistered_frameworks"`
	Version string `json:"version"`
}

type Slave struct {
	Active         bool              `json:"active"`
	Attributes     map[string]string `json:"attributes"`
	Hostname       string            `json:"hostname"`
	ID             string            `json:"id"`
	PID            string            `json:"pid"`
	RegisteredTime float64           `json:"registered_time"`
	Resources      Resources         `json:"resources"`
}

type Resources struct {
	CPUs  float64 `json:"cpus"`
	Disk  int64   `json:"disk"`
	Mem   int64   `json:"mem"`
	Ports string  `json:"ports"`
}

type Framework struct {
	Active           bool        `json:"active"`
	Checkpoint       bool        `json:"checkpoint"`
	CompletedTasks   []Task      `json:"checkpoint"`
	FailoverTimeout  int64       `json:"failover_timeout"`
	Hostname         string      `json:"hostname"`
	ID               string      `json:"id"`
	Name             string      `json:"name"`
	OfferedResources Resources   `json:"offered_resources"`
	Offers           []Resources `json:"offers"`
	RegisteredTime   float64     `json:"registered_time"`
	ReregisteredTime float64     `json:"reregistered_time"`
	Resources        Resources   `json:"resources"`
	Role             string      `json:"role"`
	Tasks            []Task      `json:"tasks"`
	UnregisteredTime float64     `json:"unregistered_time"`
	UsedResources    Resources   `json:"used_resources"`
	User             string      `json:"user"`
	WebUIURL         string      `json:"webui_url"`
}

type Task struct {
	ExecutorID  string    `json:"executor_id"`
	FrameworkID string    `json:"framework_id"`
	ID          string    `json:"id"`
	Labels      []Label   `json:"labels"`
	Name        string    `json:"name"`
	Resources   Resources `json:"resources"`
	SlaveID     string    `json:"slave_id"`
	State       string    `json:"state"`
	Statuses    []Status  `json:"statuses"`
}

type Label struct {
	Key   string `json:"key"`
	Value string `json:"value"`
}

type Status struct {
	State     string  `json:"state"`
	Timestamp float64 `json:"timestamp"`
}

type Client struct {
	client *http.Client
	host   string
}

func NewClient(mesosHost string) *Client {
	return &Client{
		client: &http.Client{},
		host:   mesosHost,
	}
}

func (c *Client) GetState() (State, error) {
	stateURL := url.URL{Scheme: "http", Host: c.host, Path: "/state.json"}
	var decodedResponse State
	err := c.do(stateURL, &decodedResponse)
	if err != nil {
		return decodedResponse, err
	}
	// strip "master@" from beginning of leader
	leader := decodedResponse.Leader[7:]

	leaderURL := url.URL{Scheme: "http", Host: leader, Path: "/state.json"}
	err = c.do(leaderURL, &decodedResponse)

	return decodedResponse, err
}

func (c *Client) do(endpoint url.URL, response interface{}) error {
	req, err := http.NewRequest("GET", endpoint.String(), nil)
	if err != nil {
		return err
	}

	resp, err := c.client.Do(req)
	if err != nil {
		return err
	}
	defer resp.Body.Close()

	jsonDecoder := json.NewDecoder(resp.Body)
	if err := jsonDecoder.Decode(&response); err != nil {
		return err
	}
	return nil
}
<commit_msg>use TrimPrefix instead of slicing<commit_after>package mesos

import (
	"encoding/json"
	"net/http"
	"net/url"
	"strings"
)

type State struct {
	ActivatedSlaves int64   `json:"activated_slaves"`
	BuildDate       string  `json:"build_date"`
	BuildTime       float64 `json:"build_time"`
	BuildUser       string  `json:"build_user"`
	// CompletedFrameworks    []Something       `json:"completed_frameworks"`
	DeactivatedSlaves int64             `json:"deactivated_slaves"`
	ElectedTime       float64           `json:"elected_time"`
	FailedTasks       int64             `json:"failed_tasks"`
	FinishedTasks     int64             `json:"finished_tasks"`
	Flags             map[string]string `json:"flags"`
	Frameworks        []Framework       `json:"frameworks"`
	GitSHA            string            `json:"git_sha"`
	GitTag            string            `json:"git_tag"`
	Hostname          string            `json:"hostname"`
	ID                string            `json:"id"`
	KilledTasks       int64             `json:"killed_tasks"`
	Leader            string            `json:"leader"`
	LogDir            string            `json:"log_dir"`
	LostTasks         int64             `json:"lost_tasks"`
	// OrphanTasks            []Something       `json:"orphan_tasks"`
	PID          string  `json:"pid"`
	Slaves       []Slave `json:"slaves"`
	StagedTasks  int64   `json:"staged_tasks"`
	StartTime    float64 `json:"start_time"`
	StartedTasks int64   `json:"started_tasks"`
	// UnregisteredFrameworks []Something       `json:"unregistered_frameworks"`
	Version string `json:"version"`
}

type Slave struct {
	Active         bool              `json:"active"`
	Attributes     map[string]string `json:"attributes"`
	Hostname       string            `json:"hostname"`
	ID             string            `json:"id"`
	PID            string            `json:"pid"`
	RegisteredTime float64           `json:"registered_time"`
	Resources      Resources         `json:"resources"`
}

type Resources struct {
	CPUs  float64 `json:"cpus"`
	Disk  int64   `json:"disk"`
	Mem   int64   `json:"mem"`
	Ports string  `json:"ports"`
}

type Framework struct {
	Active           bool        `json:"active"`
	Checkpoint       bool        `json:"checkpoint"`
	CompletedTasks   []Task      `json:"checkpoint"`
	FailoverTimeout  int64       `json:"failover_timeout"`
	Hostname         string      `json:"hostname"`
	ID               string      `json:"id"`
	Name             string      `json:"name"`
	OfferedResources Resources   `json:"offered_resources"`
	Offers           []Resources `json:"offers"`
	RegisteredTime   float64     `json:"registered_time"`
	ReregisteredTime float64     `json:"reregistered_time"`
	Resources        Resources   `json:"resources"`
	Role             string      `json:"role"`
	Tasks            []Task      `json:"tasks"`
	UnregisteredTime float64     `json:"unregistered_time"`
	UsedResources    Resources   `json:"used_resources"`
	User             string      `json:"user"`
	WebUIURL         string      `json:"webui_url"`
}

type Task struct {
	ExecutorID  string    `json:"executor_id"`
	FrameworkID string    `json:"framework_id"`
	ID          string    `json:"id"`
	Labels      []Label   `json:"labels"`
	Name        string    `json:"name"`
	Resources   Resources `json:"resources"`
	SlaveID     string    `json:"slave_id"`
	State       string    `json:"state"`
	Statuses    []Status  `json:"statuses"`
}

type Label struct {
	Key   string `json:"key"`
	Value string `json:"value"`
}

type Status struct {
	State     string  `json:"state"`
	Timestamp float64 `json:"timestamp"`
}

type Client struct {
	client *http.Client
	host   string
}

func NewClient(mesosHost string) *Client {
	return &Client{
		client: &http.Client{},
		host:   mesosHost,
	}
}

func (c *Client) GetState() (State, error) {
	stateURL := url.URL{Scheme: "http", Host: c.host, Path: "/state.json"}
	var decodedResponse State
	err := c.do(stateURL, &decodedResponse)
	if err != nil {
		return decodedResponse, err
	}
	// strip "master@" from beginning of leader
	leader := strings.TrimPrefix(decodedResponse.Leader, "master@")

	leaderURL := url.URL{Scheme: "http", Host: leader, Path: "/state.json"}
	err = c.do(leaderURL, &decodedResponse)

	return decodedResponse, err
}

func (c *Client) do(endpoint url.URL, response interface{}) error {
	req, err := http.NewRequest("GET", endpoint.String(), nil)
	if err != nil {
		return err
	}

	resp, err := c.client.Do(req)
	if err != nil {
		return err
	}
	defer resp.Body.Close()

	jsonDecoder := json.NewDecoder(resp.Body)
	if err := jsonDecoder.Decode(&response); err != nil {
		return err
	}
	return nil
}
<|endoftext|>


<commit_before>package command

import (
	"os"
	"path/filepath"
	"time"

	"github.com/mackerelio/mackerel-agent/config"
	"github.com/mackerelio/mackerel-agent/metadata"
)

func metadataGenerators(conf *config.Config) []*metadata.Generator {
	generators := make([]*metadata.Generator, 0, len(conf.MetadataPlugins))

	workdir := os.Getenv("MACKEREL_PLUGIN_WORKDIR")
	if workdir == "" {
		workdir = os.TempDir()
	}

	for name, pluginConfig := range conf.MetadataPlugins {
		generator := &metadata.Generator{
			Name:     name,
			Config:   pluginConfig,
			Tempfile: filepath.Join(workdir, "mackerel-metadata-"+name),
		}
		logger.Debugf("Metadata plugin generator created: %#v %#v", generator, generator.Config)
		generators = append(generators, generator)
	}

	return generators
}

type metadataResult struct {
	namespace string
	metadata  interface{}
}

func runMetadataLoop(c *Context, termMetadataCh <-chan struct{}, quit <-chan struct{}) {
	resultCh := make(chan *metadataResult)
	for _, g := range c.Agent.MetadataGenerators {
		go runEachMetadataLoop(g, resultCh, quit)
	}

	exit := false
	for !exit {
		select {
		case <-time.After(1 * time.Minute):
		case <-termMetadataCh:
			logger.Debugf("received 'term' chan for metadata loop")
			exit = true
		}

		results := []*metadataResult{}
	ConsumeResults:
		for {
			select {
			case result := <-resultCh:
				results = append(results, result)
			default:
				break ConsumeResults
			}
		}

		for _, result := range results {
			err := c.API.PutMetadata(c.Host.ID, result.namespace, result.metadata)
			if err != nil {
				logger.Errorf("put metadata %q failed: %s", result.namespace, err.Error())
				// do not retry because huge metadata will be always rejected
				clearMetadataCache(c.Agent.MetadataGenerators, result.namespace)
				continue
			}
		}
	}
}

func clearMetadataCache(generators []*metadata.Generator, namespace string) {
	for _, g := range generators {
		if g.Name == namespace {
			g.Clear()
			return
		}
	}
}

func runEachMetadataLoop(g *metadata.Generator, resultCh chan<- *metadataResult, quit <-chan struct{}) {
	interval := g.Interval()
	nextInterval := 10 * time.Second
	nextTime := time.Now()

	for {
		select {
		case <-time.After(nextInterval):
			metadata, err := g.Fetch()

			// case for laptop sleep mode (now >> nextTime + interval)
			now := time.Now()
			nextInterval = interval - (now.Sub(nextTime) % interval)
			nextTime = now.Add(nextInterval)

			if err != nil {
				logger.Warningf("metadata plugin %q: %s", g.Name, err.Error())
				continue
			}

			if !g.IsChanged(metadata) {
				logger.Debugf("skipping metadata %q, metadata does not change", g.Name)
				continue
			}

			if err := g.Save(metadata); err != nil {
				logger.Warningf("metadata plugin %q: %s", g.Name, err.Error())
				continue
			}

			logger.Debugf("generated metadata %q (saved cache to file: %s)", g.Name, g.Tempfile)
			resultCh <- &metadataResult{
				namespace: g.Name,
				metadata:  metadata,
			}

		case <-quit:
			return
		}
	}
}
<commit_msg>change the temporary file for metadata<commit_after>package command

import (
	"os"
	"path/filepath"
	"time"

	"github.com/mackerelio/mackerel-agent/config"
	"github.com/mackerelio/mackerel-agent/metadata"
)

func metadataGenerators(conf *config.Config) []*metadata.Generator {
	generators := make([]*metadata.Generator, 0, len(conf.MetadataPlugins))

	workdir := os.Getenv("MACKEREL_PLUGIN_WORKDIR")
	if workdir == "" {
		workdir = os.TempDir()
	}

	for name, pluginConfig := range conf.MetadataPlugins {
		generator := &metadata.Generator{
			Name:     name,
			Config:   pluginConfig,
			Tempfile: filepath.Join(workdir, "mackerel-metadata", name),
		}
		logger.Debugf("Metadata plugin generator created: %#v %#v", generator, generator.Config)
		generators = append(generators, generator)
	}

	return generators
}

type metadataResult struct {
	namespace string
	metadata  interface{}
}

func runMetadataLoop(c *Context, termMetadataCh <-chan struct{}, quit <-chan struct{}) {
	resultCh := make(chan *metadataResult)
	for _, g := range c.Agent.MetadataGenerators {
		go runEachMetadataLoop(g, resultCh, quit)
	}

	exit := false
	for !exit {
		select {
		case <-time.After(1 * time.Minute):
		case <-termMetadataCh:
			logger.Debugf("received 'term' chan for metadata loop")
			exit = true
		}

		results := []*metadataResult{}
	ConsumeResults:
		for {
			select {
			case result := <-resultCh:
				results = append(results, result)
			default:
				break ConsumeResults
			}
		}

		for _, result := range results {
			err := c.API.PutMetadata(c.Host.ID, result.namespace, result.metadata)
			if err != nil {
				logger.Errorf("put metadata %q failed: %s", result.namespace, err.Error())
				// do not retry because huge metadata will be always rejected
				clearMetadataCache(c.Agent.MetadataGenerators, result.namespace)
				continue
			}
		}
	}
}

func clearMetadataCache(generators []*metadata.Generator, namespace string) {
	for _, g := range generators {
		if g.Name == namespace {
			g.Clear()
			return
		}
	}
}

func runEachMetadataLoop(g *metadata.Generator, resultCh chan<- *metadataResult, quit <-chan struct{}) {
	interval := g.Interval()
	nextInterval := 10 * time.Second
	nextTime := time.Now()

	for {
		select {
		case <-time.After(nextInterval):
			metadata, err := g.Fetch()

			// case for laptop sleep mode (now >> nextTime + interval)
			now := time.Now()
			nextInterval = interval - (now.Sub(nextTime) % interval)
			nextTime = now.Add(nextInterval)

			if err != nil {
				logger.Warningf("metadata plugin %q: %s", g.Name, err.Error())
				continue
			}

			if !g.IsChanged(metadata) {
				logger.Debugf("skipping metadata %q, metadata does not change", g.Name)
				continue
			}

			if err := g.Save(metadata); err != nil {
				logger.Warningf("metadata plugin %q: %s", g.Name, err.Error())
				continue
			}

			logger.Debugf("generated metadata %q (saved cache to file: %s)", g.Name, g.Tempfile)
			resultCh <- &metadataResult{
				namespace: g.Name,
				metadata:  metadata,
			}

		case <-quit:
			return
		}
	}
}
<|endoftext|>