<commit_before>package main

import (
	"bufio"
	"encoding/json"
	"io/ioutil"
	"log"
	"net"
	"os"
	"os/exec"
	"path/filepath"
	"strings"
	"sync"
	"time"

	"github.com/kelseyhightower/terminus/facts"
	"golang.org/x/sys/unix"
)

// SystemFacts holds the system facts.
type SystemFacts struct {
	Architecture string
	BootID       string
	Date         Date
	Domainname   string
	Hostname     string
	Network      Network
	Kernel       Kernel
	MachineID    string
	Memory       Memory
	OSRelease    OSRelease
	Swap         Swap
	Uptime       int64
	LoadAverage  LoadAverage

	mu sync.Mutex
}

// Holds the load average facts.
type LoadAverage struct {
	One  uint64
	Five uint64
	Ten  uint64
}

// Date holds the date facts.
type Date struct {
	Unix int64
	UTC  string
}

// Swap holds the swap facts.
type Swap struct {
	Total uint64
	Free  uint64
}

// OSRelease holds the OS release facts.
type OSRelease struct {
	Name       string
	ID         string
	PrettyName string
	Version    string
	VersionID  string
}

// Kernel holds the kernel facts.
type Kernel struct {
	Name    string
	Release string
	Version string
}

// Memory holds the memory facts.
type Memory struct {
	Total    uint64
	Free     uint64
	Shared   uint64
	Buffered uint64
}

// Network holds the network facts.
type Network struct {
	Interfaces Interfaces
}

// Interfaces holds the interface facts.
type Interfaces map[string]Interface

// Interface holds facts for a single interface.
type Interface struct {
	Name         string
	Index        int
	HardwareAddr string
	IpAddresses  []string
}

func getFacts() *facts.Facts {
	f := facts.New()
	systemFacts := getSystemFacts()
	f.Add("System", systemFacts)
	processExternalFacts(externalFactsDir, f)
	return f
}

func getSystemFacts() *SystemFacts {
	facts := new(SystemFacts)
	var wg sync.WaitGroup

	wg.Add(7)
	go facts.getOSRelease(&wg)
	go facts.getInterfaces(&wg)
	go facts.getBootID(&wg)
	go facts.getMachineID(&wg)
	go facts.getUname(&wg)
	go facts.getSysInfo(&wg)
	go facts.getDate(&wg)

	wg.Wait()
	return facts
}

func (f *SystemFacts) getDate(wg *sync.WaitGroup) {
	defer wg.Done()

	now := time.Now()
	f.Date.Unix = now.Unix()
	f.Date.UTC = now.UTC().String()

	return
}

func (f *SystemFacts) getSysInfo(wg *sync.WaitGroup) {
	defer wg.Done()

	var info unix.Sysinfo_t
	if err := unix.Sysinfo(&info); err != nil {
		log.Println(err.Error())
		return
	}

	f.mu.Lock()
	defer f.mu.Unlock()

	f.Memory.Total = info.Totalram
	f.Memory.Free = info.Freeram
	f.Memory.Shared = info.Sharedram
	f.Memory.Buffered = info.Bufferram

	f.Swap.Total = info.Totalswap
	f.Swap.Free = info.Freeswap

	f.Uptime = info.Uptime

	f.LoadAverage.One = info.Loads[0]
	f.LoadAverage.Five = info.Loads[1]
	f.LoadAverage.Ten = info.Loads[2]

	return
}

func (f *SystemFacts) getOSRelease(wg *sync.WaitGroup) {
	defer wg.Done()
	osReleaseFile, err := os.Open("/etc/os-release")
	if err != nil {
		log.Println(err.Error())
		return
	}
	defer osReleaseFile.Close()

	f.mu.Lock()
	defer f.mu.Unlock()
	scanner := bufio.NewScanner(osReleaseFile)
	for scanner.Scan() {
		columns := strings.Split(scanner.Text(), "=")
		key := columns[0]
		value := strings.Trim(strings.TrimSpace(columns[1]), `"`)
		switch key {
		case "NAME":
			f.OSRelease.Name = value
		case "ID":
			f.OSRelease.ID = value
		case "PRETTY_NAME":
			f.OSRelease.PrettyName = value
		case "VERSION":
			f.OSRelease.Version = value
		case "VERSION_ID":
			f.OSRelease.VersionID = value
		}
	}
	return
}

func (f *SystemFacts) getMachineID(wg *sync.WaitGroup) {
	defer wg.Done()
	machineIDFile, err := os.Open("/etc/machine-id")
	if err != nil {
		log.Println(err.Error())
		return
	}
	defer machineIDFile.Close()
	data, err := ioutil.ReadAll(machineIDFile)
	if err != nil {
		log.Println(err.Error())
		return
	}

	f.mu.Lock()
	defer f.mu.Unlock()

	f.MachineID = strings.TrimSpace(string(data))
	return
}

func (f *SystemFacts) getBootID(wg *sync.WaitGroup) {
	defer wg.Done()
	bootIDFile, err := os.Open("/proc/sys/kernel/random/boot_id")
	if err != nil {
		log.Println(err.Error())
		return
	}
	defer bootIDFile.Close()
	data, err := ioutil.ReadAll(bootIDFile)
	if err != nil {
		log.Println(err.Error())
		return
	}

	f.mu.Lock()
	defer f.mu.Unlock()

	f.BootID = strings.TrimSpace(string(data))
	return
}

func (f *SystemFacts) getInterfaces(wg *sync.WaitGroup) {
	defer wg.Done()
	ls, err := net.Interfaces()
	if err != nil {
		log.Println(err.Error())
		return
	}

	f.mu.Lock()
	defer f.mu.Unlock()

	m := make(Interfaces)
	for _, i := range ls {
		ipaddreses := make([]string, 0)
		addrs, err := i.Addrs()
		if err != nil {
			log.Println(err.Error())
			return
		}
		for _, ip := range addrs {
			ipaddreses = append(ipaddreses, ip.String())
		}
		m[i.Name] = Interface{
			Name:         i.Name,
			Index:        i.Index,
			HardwareAddr: i.HardwareAddr.String(),
			IpAddresses:  ipaddreses,
		}
	}
	f.Network.Interfaces = m
	return
}

func (f *SystemFacts) getUname(wg *sync.WaitGroup) {
	defer wg.Done()

	var buf unix.Utsname
	err := unix.Uname(&buf)
	if err != nil {
		log.Println(err.Error())
		return
	}

	f.mu.Lock()
	defer f.mu.Unlock()

	f.Domainname = charsToString(buf.Domainname)
	f.Architecture = charsToString(buf.Machine)
	f.Hostname = charsToString(buf.Nodename)
	f.Kernel.Name = charsToString(buf.Sysname)
	f.Kernel.Release = charsToString(buf.Release)
	f.Kernel.Version = charsToString(buf.Version)
	return
}

func processExternalFacts(dir string, f *facts.Facts) {
	d, err := os.Open(dir)
	if err != nil {
		log.Println(err)
		return
	}
	defer d.Close()

	files, err := d.Readdir(0)
	if err != nil {
		log.Println(err)
		return
	}

	executableFacts := make([]string, 0)
	staticFacts := make([]string, 0)

	for _, fi := range files {
		name := filepath.Join(dir, fi.Name())
		if isExecutable(fi) {
			executableFacts = append(executableFacts, name)
			continue
		}
		if strings.HasSuffix(name, ".json") {
			staticFacts = append(staticFacts, name)
		}
	}

	var wg sync.WaitGroup
	for _, p := range staticFacts {
		p := p
		wg.Add(1)
		go factsFromFile(p, f, &wg)
	}
	for _, p := range executableFacts {
		p := p
		wg.Add(1)
		go factsFromExec(p, f, &wg)
	}
	wg.Wait()
}

func factsFromFile(path string, f *facts.Facts, wg *sync.WaitGroup) {
	defer wg.Done()
	data, err := ioutil.ReadFile(path)
	if err != nil {
		log.Println(err)
		return
	}
	var result interface{}
	err = json.Unmarshal(data, &result)
	if err != nil {
		log.Println(err)
		return
	}
	f.Add(strings.TrimSuffix(filepath.Base(path), ".json"), result)
}

func factsFromExec(path string, f *facts.Facts, wg *sync.WaitGroup) {
	defer wg.Done()
	out, err := exec.Command(path).Output()
	if err != nil {
		log.Println(err)
		return
	}
	var result interface{}
	err = json.Unmarshal(out, &result)
	if err != nil {
		log.Println(err)
		return
	}
	f.Add(filepath.Base(path), result)
}

func isExecutable(fi os.FileInfo) bool {
	if m := fi.Mode(); !m.IsDir() && m&0111 != 0 {
		return true
	}
	return false
}

func charsToString(ca [65]int8) string {
	s := make([]byte, len(ca))
	var lens int
	for ; lens < len(ca); lens++ {
		if ca[lens] == 0 {
			break
		}
		s[lens] = uint8(ca[lens])
	}
	return string(s[0:lens])
}
<commit_msg>Adds FileSystem facts from /etc/mtab<commit_after>package main

import (
	"bufio"
	"bytes"
	"encoding/json"
	"io/ioutil"
	"log"
	"net"
	"os"
	"os/exec"
	"path/filepath"
	"strconv"
	"strings"
	"sync"
	"time"

	"github.com/kelseyhightower/terminus/facts"
	"golang.org/x/sys/unix"
)

// SystemFacts holds the system facts.
type SystemFacts struct {
	Architecture string
	BootID       string
	Date         Date
	Domainname   string
	Hostname     string
	Network      Network
	Kernel       Kernel
	MachineID    string
	Memory       Memory
	OSRelease    OSRelease
	Swap         Swap
	Uptime       int64
	LoadAverage  LoadAverage
	FileSystems  FileSystems

	mu sync.Mutex
}

// Holds the load average facts.
type LoadAverage struct {
	One  uint64
	Five uint64
	Ten  uint64
}

// Date holds the date facts.
type Date struct {
	Unix int64
	UTC  string
}

// Swap holds the swap facts.
type Swap struct {
	Total uint64
	Free  uint64
}

// OSRelease holds the OS release facts.
type OSRelease struct {
	Name       string
	ID         string
	PrettyName string
	Version    string
	VersionID  string
}

// Kernel holds the kernel facts.
type Kernel struct {
	Name    string
	Release string
	Version string
}

// Memory holds the memory facts.
type Memory struct {
	Total    uint64
	Free     uint64
	Shared   uint64
	Buffered uint64
}

// Network holds the network facts.
type Network struct {
	Interfaces Interfaces
}

// Interfaces holds the interface facts.
type Interfaces map[string]Interface

// Interface holds facts for a single interface.
type Interface struct {
	Name         string
	Index        int
	HardwareAddr string
	IpAddresses  []string
}

// FileSystems holds the Filesystem facts.
type FileSystems map[string]FileSystem

// FileSystem holds facts for a filesystem (man fstab).
type FileSystem struct {
	Device     string
	MountPoint string
	Type       string
	Options    []string
	DumpFreq   uint64
	PassNo     uint64
}

func getFacts() *facts.Facts {
	f := facts.New()
	systemFacts := getSystemFacts()
	f.Add("System", systemFacts)
	processExternalFacts(externalFactsDir, f)
	return f
}

func getSystemFacts() *SystemFacts {
	facts := new(SystemFacts)
	var wg sync.WaitGroup

	wg.Add(8)
	go facts.getOSRelease(&wg)
	go facts.getInterfaces(&wg)
	go facts.getBootID(&wg)
	go facts.getMachineID(&wg)
	go facts.getUname(&wg)
	go facts.getSysInfo(&wg)
	go facts.getDate(&wg)
	go facts.getFileSystems(&wg)

	wg.Wait()
	return facts
}

func (f *SystemFacts) getDate(wg *sync.WaitGroup) {
	defer wg.Done()

	now := time.Now()
	f.Date.Unix = now.Unix()
	f.Date.UTC = now.UTC().String()

	return
}

func (f *SystemFacts) getSysInfo(wg *sync.WaitGroup) {
	defer wg.Done()

	var info unix.Sysinfo_t
	if err := unix.Sysinfo(&info); err != nil {
		log.Println(err.Error())
		return
	}

	f.mu.Lock()
	defer f.mu.Unlock()

	f.Memory.Total = info.Totalram
	f.Memory.Free = info.Freeram
	f.Memory.Shared = info.Sharedram
	f.Memory.Buffered = info.Bufferram

	f.Swap.Total = info.Totalswap
	f.Swap.Free = info.Freeswap

	f.Uptime = info.Uptime

	f.LoadAverage.One = info.Loads[0]
	f.LoadAverage.Five = info.Loads[1]
	f.LoadAverage.Ten = info.Loads[2]

	return
}

func (f *SystemFacts) getOSRelease(wg *sync.WaitGroup) {
	defer wg.Done()
	osReleaseFile, err := os.Open("/etc/os-release")
	if err != nil {
		log.Println(err.Error())
		return
	}
	defer osReleaseFile.Close()

	f.mu.Lock()
	defer f.mu.Unlock()
	scanner := bufio.NewScanner(osReleaseFile)
	for scanner.Scan() {
		columns := strings.Split(scanner.Text(), "=")
		key := columns[0]
		value := strings.Trim(strings.TrimSpace(columns[1]), `"`)
		switch key {
		case "NAME":
			f.OSRelease.Name = value
		case "ID":
			f.OSRelease.ID = value
		case "PRETTY_NAME":
			f.OSRelease.PrettyName = value
		case "VERSION":
			f.OSRelease.Version = value
		case "VERSION_ID":
			f.OSRelease.VersionID = value
		}
	}
	return
}

func (f *SystemFacts) getMachineID(wg *sync.WaitGroup) {
	defer wg.Done()
	machineIDFile, err := os.Open("/etc/machine-id")
	if err != nil {
		log.Println(err.Error())
		return
	}
	defer machineIDFile.Close()
	data, err := ioutil.ReadAll(machineIDFile)
	if err != nil {
		log.Println(err.Error())
		return
	}

	f.mu.Lock()
	defer f.mu.Unlock()

	f.MachineID = strings.TrimSpace(string(data))
	return
}

func (f *SystemFacts) getBootID(wg *sync.WaitGroup) {
	defer wg.Done()
	bootIDFile, err := os.Open("/proc/sys/kernel/random/boot_id")
	if err != nil {
		log.Println(err.Error())
		return
	}
	defer bootIDFile.Close()
	data, err := ioutil.ReadAll(bootIDFile)
	if err != nil {
		log.Println(err.Error())
		return
	}

	f.mu.Lock()
	defer f.mu.Unlock()

	f.BootID = strings.TrimSpace(string(data))
	return
}

func (f *SystemFacts) getInterfaces(wg *sync.WaitGroup) {
	defer wg.Done()
	ls, err := net.Interfaces()
	if err != nil {
		log.Println(err.Error())
		return
	}

	f.mu.Lock()
	defer f.mu.Unlock()

	m := make(Interfaces)
	for _, i := range ls {
		ipaddreses := make([]string, 0)
		addrs, err := i.Addrs()
		if err != nil {
			log.Println(err.Error())
			return
		}
		for _, ip := range addrs {
			ipaddreses = append(ipaddreses, ip.String())
		}
		m[i.Name] = Interface{
			Name:         i.Name,
			Index:        i.Index,
			HardwareAddr: i.HardwareAddr.String(),
			IpAddresses:  ipaddreses,
		}
	}
	f.Network.Interfaces = m
	return
}

func (f *SystemFacts) getUname(wg *sync.WaitGroup) {
	defer wg.Done()

	var buf unix.Utsname
	err := unix.Uname(&buf)
	if err != nil {
		log.Println(err.Error())
		return
	}

	f.mu.Lock()
	defer f.mu.Unlock()

	f.Domainname = charsToString(buf.Domainname)
	f.Architecture = charsToString(buf.Machine)
	f.Hostname = charsToString(buf.Nodename)
	f.Kernel.Name = charsToString(buf.Sysname)
	f.Kernel.Release = charsToString(buf.Release)
	f.Kernel.Version = charsToString(buf.Version)
	return
}

func (f *SystemFacts) getFileSystems(wg *sync.WaitGroup) {
	defer wg.Done()

	mtab, err := ioutil.ReadFile("/etc/mtab")
	if err != nil {
		log.Println(err.Error())
		return
	}

	fsMap := make(FileSystems)

	f.mu.Lock()
	defer f.mu.Unlock()

	s := bufio.NewScanner(bytes.NewBuffer(mtab))
	for s.Scan() {
		line := s.Text()
		if string(line[0]) == "#" {
			continue
		}
		fields := strings.Fields(s.Text())
		fs := FileSystem{}
		fs.Device = fields[0]
		fs.MountPoint = fields[1]
		fs.Type = fields[2]
		fs.Options = strings.Split(fields[3], ",")
		dumpFreq, err := strconv.ParseUint(fields[4], 10, 64)
		if err != nil {
			log.Println(err.Error())
			return
		}
		fs.DumpFreq = dumpFreq

		passNo, err := strconv.ParseUint(fields[4], 10, 64)
		if err != nil {
			log.Println(err.Error())
			return
		}
		fs.PassNo = passNo

		fsMap[fs.Device] = fs
	}

	f.FileSystems = fsMap
	return
}

func processExternalFacts(dir string, f *facts.Facts) {
	d, err := os.Open(dir)
	if err != nil {
		log.Println(err)
		return
	}
	defer d.Close()

	files, err := d.Readdir(0)
	if err != nil {
		log.Println(err)
		return
	}

	executableFacts := make([]string, 0)
	staticFacts := make([]string, 0)

	for _, fi := range files {
		name := filepath.Join(dir, fi.Name())
		if isExecutable(fi) {
			executableFacts = append(executableFacts, name)
			continue
		}
		if strings.HasSuffix(name, ".json") {
			staticFacts = append(staticFacts, name)
		}
	}

	var wg sync.WaitGroup
	for _, p := range staticFacts {
		p := p
		wg.Add(1)
		go factsFromFile(p, f, &wg)
	}
	for _, p := range executableFacts {
		p := p
		wg.Add(1)
		go factsFromExec(p, f, &wg)
	}
	wg.Wait()
}

func factsFromFile(path string, f *facts.Facts, wg *sync.WaitGroup) {
	defer wg.Done()
	data, err := ioutil.ReadFile(path)
	if err != nil {
		log.Println(err)
		return
	}
	var result interface{}
	err = json.Unmarshal(data, &result)
	if err != nil {
		log.Println(err)
		return
	}
	f.Add(strings.TrimSuffix(filepath.Base(path), ".json"), result)
}

func factsFromExec(path string, f *facts.Facts, wg *sync.WaitGroup) {
	defer wg.Done()
	out, err := exec.Command(path).Output()
	if err != nil {
		log.Println(err)
		return
	}
	var result interface{}
	err = json.Unmarshal(out, &result)
	if err != nil {
		log.Println(err)
		return
	}
	f.Add(filepath.Base(path), result)
}

func isExecutable(fi os.FileInfo) bool {
	if m := fi.Mode(); !m.IsDir() && m&0111 != 0 {
		return true
	}
	return false
}

func charsToString(ca [65]int8) string {
	s := make([]byte, len(ca))
	var lens int
	for ; lens < len(ca); lens++ {
		if ca[lens] == 0 {
			break
		}
		s[lens] = uint8(ca[lens])
	}
	return string(s[0:lens])
}
<|endoftext|>


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

package thrift

import (
	"bytes"
	"io"
	"io/ioutil"
	"net/http"
	"net/url"
	"strconv"
	"sync"
)

type THttpClient struct {
	// the most recent response, if any
	responseBuffer     *bytes.Buffer
	url                *url.URL
	requestBuffer      *bytes.Buffer
	header             http.Header
	nsecConnectTimeout int64
	nsecReadTimeout    int64
	httpClient         *http.Client
}

type THttpClientTransportFactory struct {
	url    string
	isPost bool
}

func (p *THttpClientTransportFactory) GetTransport(trans TTransport) TTransport {
	if trans != nil {
		t, ok := trans.(*THttpClient)
		if ok && t.url != nil {
			if t.requestBuffer != nil {
				t2, _ := NewTHttpPostClient(t.url.String())
				return t2
			}
			t2, _ := NewTHttpClient(t.url.String())
			return t2
		}
	}
	if p.isPost {
		s, _ := NewTHttpPostClient(p.url)
		return s
	}
	s, _ := NewTHttpClient(p.url)
	return s
}

func NewTHttpClientTransportFactory(url string) *THttpClientTransportFactory {
	return &THttpClientTransportFactory{url: url, isPost: false}
}

func NewTHttpPostClientTransportFactory(url string) *THttpClientTransportFactory {
	return &THttpClientTransportFactory{url: url, isPost: true}
}

func newHttpClient() *http.Client {
	return &http.Client{
		// Setting a non-nil Transport means that each client will get its own
		// idle connection pool.  This is way more stable than sharing a single
		// pool for this entire server.
		Transport: &http.Transport{}, // TODO set dial and TLS timeouts?
	}
}

func NewTHttpClient(urlstr string) (TTransport, error) {
	parsedURL, err := url.Parse(urlstr)
	if err != nil {
		return nil, err
	}
	response, err := http.Get(urlstr)
	if response != nil && response.Body != nil {
		defer response.Body.Close()
	}
	if err != nil {
		return nil, err
	}
	bs, err := ioutil.ReadAll(response.Body)
	if err != nil {
		return nil, err
	}
	buf := bytes.NewBuffer(bs)

	return &THttpClient{
		responseBuffer: buf,
		url:            parsedURL,
		httpClient:     newHttpClient(),
	}, nil
}

// bufPool is a set of buffers used for the request bodies. Buffers should
// be Reset() before they are added to the pool.
var bufPool = sync.Pool{
	New: func() interface{} {
		return bytes.NewBuffer(make([]byte, 0, 1024))
	},
}

func getBuffer() *bytes.Buffer {
	buf := bufPool.Get().(*bytes.Buffer)
	if buf == nil {
		buf = bytes.NewBuffer(make([]byte, 0, 1024))
	}
	return buf
}

func NewTHttpPostClient(urlstr string) (TTransport, error) {
	parsedURL, err := url.Parse(urlstr)
	if err != nil {
		return nil, err
	}
	return &THttpClient{
		url:           parsedURL,
		requestBuffer: getBuffer(),
		header:        http.Header{},
		httpClient:    newHttpClient(),
	}, nil
}

// Set the HTTP Header for this specific Thrift Transport
// It is important that you first assert the TTransport as a THttpClient type
// like so:
//
// httpTrans := trans.(THttpClient)
// httpTrans.SetHeader("User-Agent","Thrift Client 1.0")
func (p *THttpClient) SetHeader(key string, value string) {
	p.header.Set(key, value)
}

// Get the HTTP Header represented by the supplied Header Key for this specific Thrift Transport
// It is important that you first assert the TTransport as a THttpClient type
// like so:
//
// httpTrans := trans.(THttpClient)
// hdrValue := httpTrans.GetHeader("User-Agent")
func (p *THttpClient) GetHeader(key string) string {
	return p.header.Get(key)
}

// Deletes the HTTP Header given a Header Key for this specific Thrift Transport
// It is important that you first assert the TTransport as a THttpClient type
// like so:
//
// httpTrans := trans.(THttpClient)
// httpTrans.DelHeader("User-Agent")
func (p *THttpClient) DelHeader(key string) {
	p.header.Del(key)
}

func (p *THttpClient) Open() error {
	// do nothing
	return nil
}

func (p *THttpClient) IsOpen() bool {
	return p.responseBuffer != nil || p.requestBuffer != nil
}

func (p *THttpClient) Peek() bool {
	return p.IsOpen()
}

func (p *THttpClient) Close() error {
	if p.responseBuffer != nil {
		p.responseBuffer = nil
	}
	if p.requestBuffer != nil {
		p.requestBuffer.Reset()
		bufPool.Put(p.requestBuffer)
		p.requestBuffer = nil
	}
	return nil
}

func (p *THttpClient) Read(buf []byte) (int, error) {
	if p.responseBuffer == nil {
		return 0, NewTTransportException(NOT_OPEN, "Response buffer is empty, no request.")
	}
	n, err := p.responseBuffer.Read(buf)
	if n > 0 && (err == nil || err == io.EOF) {
		return n, nil
	}
	return n, NewTTransportExceptionFromError(err)
}

func (p *THttpClient) ReadByte() (c byte, err error) {
	return readByte(p.responseBuffer)
}

func (p *THttpClient) Write(buf []byte) (int, error) {
	n, err := p.requestBuffer.Write(buf)
	return n, err
}

func (p *THttpClient) WriteByte(c byte) error {
	return p.requestBuffer.WriteByte(c)
}

func (p *THttpClient) WriteString(s string) (n int, err error) {
	return p.requestBuffer.WriteString(s)
}

// bufferCloser makes a bytes.Buffer into an io.Closer and returns unused
// buffers to the pool.
type bufferCloser struct {
	*bytes.Buffer
}

// THttpClient reuses its buffer from one request to the next. If that
// buffer is not reset, well, it'll just keep getting extended from
// one request to the next. Make the bytes.Buffer into an io.Closer
// and rely on the http transport to call Close() (and therefore
// to call Reset() and return it to the pool).
func (b *bufferCloser) Close() error {
	b.Reset()
	bufPool.Put(b.Buffer)
	return nil
}

func (p *THttpClient) Flush() error {
	req, err := http.NewRequest("POST", p.url.String(), &bufferCloser{p.requestBuffer})
	if err != nil {
		return NewTTransportExceptionFromError(err)
	}
	// Reset the request/response state
	p.requestBuffer = getBuffer()
	p.responseBuffer = nil
	// Do the call
	p.header.Set("Content-Type", "application/x-thrift")
	req.Header = p.header
	response, err := p.httpClient.Do(req)
	if response != nil && response.Body != nil {
		defer response.Body.Close()
	}

	if err != nil {
		return NewTTransportExceptionFromError(err)
	}
	if response.StatusCode != http.StatusOK {
		// TODO(pomack) log bad response
		return NewTTransportException(UNKNOWN_TRANSPORT_EXCEPTION, "HTTP Response code: "+strconv.Itoa(response.StatusCode))
	}
	bs, err := ioutil.ReadAll(response.Body)
	if err != nil {
		return err
	}
	p.responseBuffer = bytes.NewBuffer(bs)
	return nil
}
<commit_msg>don’t reuse headers across HTTP calls<commit_after>/*
 * Licensed to the Apache Software Foundation (ASF) under one
 * or more contributor license agreements. See the NOTICE file
 * distributed with this work for additional information
 * regarding copyright ownership. The ASF licenses this file
 * to you under the Apache License, Version 2.0 (the
 * "License"); you may not use this file except in compliance
 * with the License. You may obtain a copy of the License at
 *
 *   http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing,
 * software distributed under the License is distributed on an
 * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
 * KIND, either express or implied. See the License for the
 * specific language governing permissions and limitations
 * under the License.
 */

package thrift

import (
	"bytes"
	"io"
	"io/ioutil"
	"net/http"
	"net/url"
	"strconv"
	"sync"
)

type THttpClient struct {
	// the most recent response, if any
	responseBuffer     *bytes.Buffer
	url                *url.URL
	requestBuffer      *bytes.Buffer
	header             http.Header
	nsecConnectTimeout int64
	nsecReadTimeout    int64
	httpClient         *http.Client
}

type THttpClientTransportFactory struct {
	url    string
	isPost bool
}

func (p *THttpClientTransportFactory) GetTransport(trans TTransport) TTransport {
	if trans != nil {
		t, ok := trans.(*THttpClient)
		if ok && t.url != nil {
			if t.requestBuffer != nil {
				t2, _ := NewTHttpPostClient(t.url.String())
				return t2
			}
			t2, _ := NewTHttpClient(t.url.String())
			return t2
		}
	}
	if p.isPost {
		s, _ := NewTHttpPostClient(p.url)
		return s
	}
	s, _ := NewTHttpClient(p.url)
	return s
}

func NewTHttpClientTransportFactory(url string) *THttpClientTransportFactory {
	return &THttpClientTransportFactory{url: url, isPost: false}
}

func NewTHttpPostClientTransportFactory(url string) *THttpClientTransportFactory {
	return &THttpClientTransportFactory{url: url, isPost: true}
}

func newHttpClient() *http.Client {
	return &http.Client{
		// Setting a non-nil Transport means that each client will get its own
		// idle connection pool.  This is way more stable than sharing a single
		// pool for this entire server.
		Transport: &http.Transport{}, // TODO set dial and TLS timeouts?
	}
}

func NewTHttpClient(urlstr string) (TTransport, error) {
	parsedURL, err := url.Parse(urlstr)
	if err != nil {
		return nil, err
	}
	response, err := http.Get(urlstr)
	if response != nil && response.Body != nil {
		defer response.Body.Close()
	}
	if err != nil {
		return nil, err
	}
	bs, err := ioutil.ReadAll(response.Body)
	if err != nil {
		return nil, err
	}
	buf := bytes.NewBuffer(bs)

	return &THttpClient{
		responseBuffer: buf,
		url:            parsedURL,
		httpClient:     newHttpClient(),
	}, nil
}

// bufPool is a set of buffers used for the request bodies. Buffers should
// be Reset() before they are added to the pool.
var bufPool = sync.Pool{
	New: func() interface{} {
		return bytes.NewBuffer(make([]byte, 0, 1024))
	},
}

func getBuffer() *bytes.Buffer {
	buf := bufPool.Get().(*bytes.Buffer)
	if buf == nil {
		buf = bytes.NewBuffer(make([]byte, 0, 1024))
	}
	return buf
}

func NewTHttpPostClient(urlstr string) (TTransport, error) {
	parsedURL, err := url.Parse(urlstr)
	if err != nil {
		return nil, err
	}
	return &THttpClient{
		url:           parsedURL,
		requestBuffer: getBuffer(),
		header:        http.Header{},
		httpClient:    newHttpClient(),
	}, nil
}

// Set the HTTP Header for this specific Thrift Transport
// It is important that you first assert the TTransport as a THttpClient type
// like so:
//
// httpTrans := trans.(THttpClient)
// httpTrans.SetHeader("User-Agent","Thrift Client 1.0")
func (p *THttpClient) SetHeader(key string, value string) {
	p.header.Set(key, value)
}

// Get the HTTP Header represented by the supplied Header Key for this specific Thrift Transport
// It is important that you first assert the TTransport as a THttpClient type
// like so:
//
// httpTrans := trans.(THttpClient)
// hdrValue := httpTrans.GetHeader("User-Agent")
func (p *THttpClient) GetHeader(key string) string {
	return p.header.Get(key)
}

// Deletes the HTTP Header given a Header Key for this specific Thrift Transport
// It is important that you first assert the TTransport as a THttpClient type
// like so:
//
// httpTrans := trans.(THttpClient)
// httpTrans.DelHeader("User-Agent")
func (p *THttpClient) DelHeader(key string) {
	p.header.Del(key)
}

func (p *THttpClient) Open() error {
	// do nothing
	return nil
}

func (p *THttpClient) IsOpen() bool {
	return p.responseBuffer != nil || p.requestBuffer != nil
}

func (p *THttpClient) Peek() bool {
	return p.IsOpen()
}

func (p *THttpClient) Close() error {
	if p.responseBuffer != nil {
		p.responseBuffer = nil
	}
	if p.requestBuffer != nil {
		p.requestBuffer.Reset()
		bufPool.Put(p.requestBuffer)
		p.requestBuffer = nil
	}
	return nil
}

func (p *THttpClient) Read(buf []byte) (int, error) {
	if p.responseBuffer == nil {
		return 0, NewTTransportException(NOT_OPEN, "Response buffer is empty, no request.")
	}
	n, err := p.responseBuffer.Read(buf)
	if n > 0 && (err == nil || err == io.EOF) {
		return n, nil
	}
	return n, NewTTransportExceptionFromError(err)
}

func (p *THttpClient) ReadByte() (c byte, err error) {
	return readByte(p.responseBuffer)
}

func (p *THttpClient) Write(buf []byte) (int, error) {
	n, err := p.requestBuffer.Write(buf)
	return n, err
}

func (p *THttpClient) WriteByte(c byte) error {
	return p.requestBuffer.WriteByte(c)
}

func (p *THttpClient) WriteString(s string) (n int, err error) {
	return p.requestBuffer.WriteString(s)
}

// bufferCloser makes a bytes.Buffer into an io.Closer and returns unused
// buffers to the pool.
type bufferCloser struct {
	*bytes.Buffer
}

// THttpClient reuses its buffer from one request to the next. If that
// buffer is not reset, well, it'll just keep getting extended from
// one request to the next. Make the bytes.Buffer into an io.Closer
// and rely on the http transport to call Close() (and therefore
// to call Reset() and return it to the pool).
func (b *bufferCloser) Close() error {
	b.Reset()
	bufPool.Put(b.Buffer)
	return nil
}

func (p *THttpClient) Flush() error {
	req, err := http.NewRequest("POST", p.url.String(), &bufferCloser{p.requestBuffer})
	if err != nil {
		return NewTTransportExceptionFromError(err)
	}
	// Reset the request/response state
	p.requestBuffer = getBuffer()
	p.responseBuffer = nil
	// Do the call
	req.Header = p.header
	// http2 reads the header outside of the call to Do(), so make a new
	// header for next time here.
	p.header = http.Header{}
	// Also, I'm not sure why the headers are reused from request to request,
	// so make a copy of them in case anyone is expecting that.
	for k, v := range req.Header {
		p.header[k] = v
	}
	req.Header.Set("Content-Type", "application/x-thrift")
	response, err := p.httpClient.Do(req)
	if response != nil && response.Body != nil {
		defer response.Body.Close()
	}

	if err != nil {
		return NewTTransportExceptionFromError(err)
	}
	if response.StatusCode != http.StatusOK {
		// TODO(pomack) log bad response
		return NewTTransportException(UNKNOWN_TRANSPORT_EXCEPTION, "HTTP Response code: "+strconv.Itoa(response.StatusCode))
	}
	bs, err := ioutil.ReadAll(response.Body)
	if err != nil {
		return err
	}
	p.responseBuffer = bytes.NewBuffer(bs)
	return nil
}
<|endoftext|>


<commit_before>package mdata

import (
	"bytes"
	"encoding/binary"
	"errors"
	"fmt"
	"sort"
	"strings"
	"time"

	"github.com/dgryski/go-tsz"
	"github.com/gocql/gocql"
	"github.com/raintank/met"
	"github.com/raintank/metrictank/iter"
	"github.com/raintank/metrictank/mdata/chunk"
	"github.com/raintank/worldping-api/pkg/log"
)

// write aggregated data to cassandra.

const Month_sec = 60 * 60 * 24 * 28

const keyspace_schema = `CREATE KEYSPACE IF NOT EXISTS raintank WITH replication = {'class': 'SimpleStrategy', 'replication_factor': 1}  AND durable_writes = true`
const table_schema = `CREATE TABLE IF NOT EXISTS raintank.metric (
    key ascii,
    ts int,
    data blob,
    PRIMARY KEY (key, ts)
) WITH COMPACT STORAGE
    AND CLUSTERING ORDER BY (ts DESC)
    AND compaction = {'class': 'org.apache.cassandra.db.compaction.DateTieredCompactionStrategy'}
    AND compression = {'sstable_compression': 'org.apache.cassandra.io.compress.LZ4Compressor'}
    AND read_repair_chance = 0.0
    AND dclocal_read_repair_chance = 0`

var (
	errChunkTooSmall      = errors.New("unpossibly small chunk in cassandra")
	errUnknownChunkFormat = errors.New("unrecognized chunk format in cassandra")
	errStartBeforeEnd     = errors.New("start must be before end.")

	cassBlockDuration met.Timer
	cassGetDuration   met.Timer
	cassPutDuration   met.Timer

	cassChunksPerRow    met.Meter
	cassRowsPerResponse met.Meter
	cassToIterDuration  met.Timer

	chunkSaveOk   met.Count
	chunkSaveFail met.Count
	// it's pretty expensive/impossible to do chunk size in mem vs in cassandra etc, but we can more easily measure chunk sizes when we operate on them
	chunkSizeAtSave met.Meter
	chunkSizeAtLoad met.Meter
)

/*
https://godoc.org/github.com/gocql/gocql#Session
Session is the interface used by users to interact with the database.
It's safe for concurrent use by multiple goroutines and a typical usage scenario is to have one global session
object to interact with the whole Cassandra cluster.
*/

type cassandraStore struct {
	session          *gocql.Session
	writeQueues      []chan *ChunkWriteRequest
	readQueue        chan *ChunkReadRequest
	writeQueueMeters []met.Meter
}

func NewCassandraStore(stats met.Backend, addrs, consistency string, timeout, readers, writers, readqsize, writeqsize int) (*cassandraStore, error) {
	cluster := gocql.NewCluster(strings.Split(addrs, ",")...)
	cluster.Consistency = gocql.ParseConsistency(consistency)
	cluster.Timeout = time.Duration(timeout) * time.Millisecond
	cluster.NumConns = writers
	var err error
	tmpSession, err := cluster.CreateSession()
	if err != nil {
		return nil, err
	}
	// ensure the keyspace and table exist.
	err = tmpSession.Query(keyspace_schema).Exec()
	if err != nil {
		return nil, err
	}
	err = tmpSession.Query(table_schema).Exec()
	if err != nil {
		return nil, err
	}
	tmpSession.Close()
	cluster.Keyspace = "raintank"
	session, err := cluster.CreateSession()
	if err != nil {
		return nil, err
	}
	c := &cassandraStore{
		session:          session,
		writeQueues:      make([]chan *ChunkWriteRequest, writers),
		readQueue:        make(chan *ChunkReadRequest, readqsize),
		writeQueueMeters: make([]met.Meter, writers),
	}

	for i := 0; i < writers; i++ {
		c.writeQueues[i] = make(chan *ChunkWriteRequest, writeqsize)
		c.writeQueueMeters[i] = stats.NewMeter(fmt.Sprintf("cassandra.write_queue.%d.items", i+1), 0)
		go c.processWriteQueue(c.writeQueues[i], c.writeQueueMeters[i])
	}

	for i := 0; i < readers; i++ {
		go c.processReadQueue()
	}

	return c, err
}

func (c *cassandraStore) InitMetrics(stats met.Backend) {
	cassBlockDuration = stats.NewTimer("cassandra.block_duration", 0)
	cassGetDuration = stats.NewTimer("cassandra.get_duration", 0)
	cassPutDuration = stats.NewTimer("cassandra.put_duration", 0)

	cassChunksPerRow = stats.NewMeter("cassandra.chunks_per_row", 0)
	cassRowsPerResponse = stats.NewMeter("cassandra.rows_per_response", 0)
	cassToIterDuration = stats.NewTimer("cassandra.to_iter_duration", 0)

	chunkSaveOk = stats.NewCount("chunks.save_ok")
	chunkSaveFail = stats.NewCount("chunks.save_fail")
	chunkSizeAtSave = stats.NewMeter("chunk_size.at_save", 0)
	chunkSizeAtLoad = stats.NewMeter("chunk_size.at_load", 0)
}

func (c *cassandraStore) Add(cwr *ChunkWriteRequest) {
	sum := 0
	for _, char := range cwr.key {
		sum += int(char)
	}
	which := sum % len(c.writeQueues)
	c.writeQueueMeters[which].Value(int64(len(c.writeQueues[which])))
	c.writeQueues[which] <- cwr
	c.writeQueueMeters[which].Value(int64(len(c.writeQueues[which])))
}

/* process writeQueue.
 */
func (c *cassandraStore) processWriteQueue(queue chan *ChunkWriteRequest, meter met.Meter) {
	tick := time.Tick(time.Duration(1) * time.Second)
	for {
		select {
		case <-tick:
			meter.Value(int64(len(queue)))
		case cwr := <-queue:
			meter.Value(int64(len(queue)))
			log.Debug("CS: starting to save %s:%d %v", cwr.key, cwr.chunk.T0, cwr.chunk)
			//log how long the chunk waited in the queue before we attempted to save to cassandra
			cassBlockDuration.Value(time.Now().Sub(cwr.timestamp))

			data := cwr.chunk.Series.Bytes()
			chunkSizeAtSave.Value(int64(len(data)))
			version := chunk.FormatStandardGoTsz
			buf := new(bytes.Buffer)
			err := binary.Write(buf, binary.LittleEndian, uint8(version))
			if err != nil {
				// TODO
			}
			_, err = buf.Write(data)
			if err != nil {
				// TODO
			}
			success := false
			attempts := 0
			for !success {
				err := c.insertChunk(cwr.key, cwr.chunk.T0, buf.Bytes(), int(cwr.ttl))
				if err == nil {
					success = true
					cwr.chunk.Saved = true
					SendPersistMessage(cwr.key, cwr.chunk.T0)
					log.Debug("CS: save complete. %s:%d %v", cwr.key, cwr.chunk.T0, cwr.chunk)
					chunkSaveOk.Inc(1)
				} else {
					if (attempts % 20) == 0 {
						log.Warn("CS: failed to save chunk to cassandra after %d attempts. %v, %s", attempts+1, cwr.chunk, err)
					}
					chunkSaveFail.Inc(1)
					sleepTime := 100 * attempts
					if sleepTime > 2000 {
						sleepTime = 2000
					}
					time.Sleep(time.Duration(sleepTime) * time.Millisecond)
					attempts++
				}
			}
		}
	}
}

// Insert Chunks into Cassandra.
//
// key: is the metric_id
// ts: is the start of the aggregated time range.
// data: is the payload as bytes.
func (c *cassandraStore) insertChunk(key string, t0 uint32, data []byte, ttl int) error {
	// for unit tests
	if c.session == nil {
		return nil
	}
	query := fmt.Sprintf("INSERT INTO metric (key, ts, data) values(?,?,?) USING TTL %d", ttl)
	row_key := fmt.Sprintf("%s_%d", key, t0/Month_sec) // "month number" based on unix timestamp (rounded down)
	pre := time.Now()
	ret := c.session.Query(query, row_key, t0, data).Exec()
	cassPutDuration.Value(time.Now().Sub(pre))
	return ret
}

type outcome struct {
	month   uint32
	sortKey uint32
	i       *gocql.Iter
}
type asc []outcome

func (o asc) Len() int           { return len(o) }
func (o asc) Swap(i, j int)      { o[i], o[j] = o[j], o[i] }
func (o asc) Less(i, j int) bool { return o[i].sortKey < o[j].sortKey }

func (c *cassandraStore) processReadQueue() {
	for crr := range c.readQueue {
		crr.out <- outcome{crr.month, crr.sortKey, c.session.Query(crr.q, crr.p...).Iter()}
	}
}

// Basic search of cassandra.
// start inclusive, end exclusive
func (c *cassandraStore) Search(key string, start, end uint32) ([]iter.Iter, error) {
	iters := make([]iter.Iter, 0)
	if start > end {
		return iters, errStartBeforeEnd
	}

	crrs := make([]*ChunkReadRequest, 0)

	query := func(month, sortKey uint32, q string, p ...interface{}) {
		crrs = append(crrs, &ChunkReadRequest{month, sortKey, q, p, nil})
	}

	start_month := start - (start % Month_sec)       // starting row has to be at, or before, requested start
	end_month := (end - 1) - ((end - 1) % Month_sec) // ending row has to be include the last point we might need

	pre := time.Now()

	// unfortunately in the database we only have the t0's of all chunks.
	// this means we can easily make sure to include the correct last chunk (just query for a t0 < end, the last chunk will contain the last needed data)
	// but it becomes hard to find which should be the first chunk to include. we can't just query for start <= t0 because than will miss some data at the beginning
	// we can't assume we know the chunkSpan so we can't just calculate the t0 >= start - <some-predefined-number> because chunkSpans may change over time.
	// we effectively need all chunks with a t0 > start, as well as the last chunk with a t0 <= start.
	// since we make sure that you can only use chunkSpans so that Month_sec % chunkSpan == 0, we know that this previous chunk will always be in the same row
	// as the one that has start_month.

	row_key := fmt.Sprintf("%s_%d", key, start_month/Month_sec)

	query(start_month, start_month, "SELECT ts, data FROM metric WHERE key=? AND ts <= ? Limit 1", row_key, start)

	if start_month == end_month {
		// we need a selection of the row between startTs and endTs
		row_key = fmt.Sprintf("%s_%d", key, start_month/Month_sec)
		query(start_month, start_month+1, "SELECT ts, data FROM metric WHERE key = ? AND ts > ? AND ts < ? ORDER BY ts ASC", row_key, start, end)
	} else {
		// get row_keys for each row we need to query.
		for month := start_month; month <= end_month; month += Month_sec {
			row_key = fmt.Sprintf("%s_%d", key, month/Month_sec)
			if month == start_month {
				// we want from startTs to the end of the row.
				query(month, month+1, "SELECT ts, data FROM metric WHERE key = ? AND ts >= ? ORDER BY ts ASC", row_key, start+1)
			} else if month == end_month {
				// we want from start of the row till the endTs.
				query(month, month, "SELECT ts, data FROM metric WHERE key = ? AND ts <= ? ORDER BY ts ASC", row_key, end-1)
			} else {
				// we want all columns
				query(month, month, "SELECT ts, data FROM metric WHERE key = ? ORDER BY ts ASC", row_key)
			}
		}
	}
	numQueries := len(crrs)
	results := make(chan outcome, numQueries)
	for i := range crrs {
		crrs[i].out = results
		c.readQueue <- crrs[i]
	}
	outcomes := make([]outcome, 0, numQueries)

	seen := 0
	for o := range results {
		seen += 1
		outcomes = append(outcomes, o)
		if seen == numQueries {
			close(results)
			cassGetDuration.Value(time.Now().Sub(pre))
			break
		}
	}
	pre = time.Now()
	// we have all of the results, but they could have arrived in any order.
	sort.Sort(asc(outcomes))

	var b []byte
	var ts int
	for _, outcome := range outcomes {
		chunks := int64(0)
		for outcome.i.Scan(&ts, &b) {
			chunks += 1
			chunkSizeAtLoad.Value(int64(len(b)))
			if len(b) < 2 {
				log.Error(3, errChunkTooSmall.Error())
				return iters, errChunkTooSmall
			}
			if chunk.Format(b[0]) != chunk.FormatStandardGoTsz {
				log.Error(3, errUnknownChunkFormat.Error())
				return iters, errUnknownChunkFormat
			}
			it, err := tsz.NewIterator(b[1:])
			if err != nil {
				log.Error(3, "failed to unpack cassandra payload. %s", err)
				return iters, err
			}
			iters = append(iters, iter.New(it, true))
		}
		cassChunksPerRow.Value(chunks)
		err := outcome.i.Close()
		if err != nil {
			log.Error(3, "cassandra query error. %s", err)
		}
	}
	cassToIterDuration.Value(time.Now().Sub(pre))
	cassRowsPerResponse.Value(int64(len(outcomes)))
	log.Debug("CS: searchCassandra(): %d outcomes (queries), %d total iters", len(outcomes), len(iters))
	return iters, nil
}

func (c *cassandraStore) Stop() {
	c.session.Close()
}
<commit_msg>cleanup useless todo's. no need to check err return<commit_after>package mdata

import (
	"bytes"
	"encoding/binary"
	"errors"
	"fmt"
	"sort"
	"strings"
	"time"

	"github.com/dgryski/go-tsz"
	"github.com/gocql/gocql"
	"github.com/raintank/met"
	"github.com/raintank/metrictank/iter"
	"github.com/raintank/metrictank/mdata/chunk"
	"github.com/raintank/worldping-api/pkg/log"
)

// write aggregated data to cassandra.

const Month_sec = 60 * 60 * 24 * 28

const keyspace_schema = `CREATE KEYSPACE IF NOT EXISTS raintank WITH replication = {'class': 'SimpleStrategy', 'replication_factor': 1}  AND durable_writes = true`
const table_schema = `CREATE TABLE IF NOT EXISTS raintank.metric (
    key ascii,
    ts int,
    data blob,
    PRIMARY KEY (key, ts)
) WITH COMPACT STORAGE
    AND CLUSTERING ORDER BY (ts DESC)
    AND compaction = {'class': 'org.apache.cassandra.db.compaction.DateTieredCompactionStrategy'}
    AND compression = {'sstable_compression': 'org.apache.cassandra.io.compress.LZ4Compressor'}
    AND read_repair_chance = 0.0
    AND dclocal_read_repair_chance = 0`

var (
	errChunkTooSmall      = errors.New("unpossibly small chunk in cassandra")
	errUnknownChunkFormat = errors.New("unrecognized chunk format in cassandra")
	errStartBeforeEnd     = errors.New("start must be before end.")

	cassBlockDuration met.Timer
	cassGetDuration   met.Timer
	cassPutDuration   met.Timer

	cassChunksPerRow    met.Meter
	cassRowsPerResponse met.Meter
	cassToIterDuration  met.Timer

	chunkSaveOk   met.Count
	chunkSaveFail met.Count
	// it's pretty expensive/impossible to do chunk size in mem vs in cassandra etc, but we can more easily measure chunk sizes when we operate on them
	chunkSizeAtSave met.Meter
	chunkSizeAtLoad met.Meter
)

/*
https://godoc.org/github.com/gocql/gocql#Session
Session is the interface used by users to interact with the database.
It's safe for concurrent use by multiple goroutines and a typical usage scenario is to have one global session
object to interact with the whole Cassandra cluster.
*/

type cassandraStore struct {
	session          *gocql.Session
	writeQueues      []chan *ChunkWriteRequest
	readQueue        chan *ChunkReadRequest
	writeQueueMeters []met.Meter
}

func NewCassandraStore(stats met.Backend, addrs, consistency string, timeout, readers, writers, readqsize, writeqsize int) (*cassandraStore, error) {
	cluster := gocql.NewCluster(strings.Split(addrs, ",")...)
	cluster.Consistency = gocql.ParseConsistency(consistency)
	cluster.Timeout = time.Duration(timeout) * time.Millisecond
	cluster.NumConns = writers
	var err error
	tmpSession, err := cluster.CreateSession()
	if err != nil {
		return nil, err
	}
	// ensure the keyspace and table exist.
	err = tmpSession.Query(keyspace_schema).Exec()
	if err != nil {
		return nil, err
	}
	err = tmpSession.Query(table_schema).Exec()
	if err != nil {
		return nil, err
	}
	tmpSession.Close()
	cluster.Keyspace = "raintank"
	session, err := cluster.CreateSession()
	if err != nil {
		return nil, err
	}
	c := &cassandraStore{
		session:          session,
		writeQueues:      make([]chan *ChunkWriteRequest, writers),
		readQueue:        make(chan *ChunkReadRequest, readqsize),
		writeQueueMeters: make([]met.Meter, writers),
	}

	for i := 0; i < writers; i++ {
		c.writeQueues[i] = make(chan *ChunkWriteRequest, writeqsize)
		c.writeQueueMeters[i] = stats.NewMeter(fmt.Sprintf("cassandra.write_queue.%d.items", i+1), 0)
		go c.processWriteQueue(c.writeQueues[i], c.writeQueueMeters[i])
	}

	for i := 0; i < readers; i++ {
		go c.processReadQueue()
	}

	return c, err
}

func (c *cassandraStore) InitMetrics(stats met.Backend) {
	cassBlockDuration = stats.NewTimer("cassandra.block_duration", 0)
	cassGetDuration = stats.NewTimer("cassandra.get_duration", 0)
	cassPutDuration = stats.NewTimer("cassandra.put_duration", 0)

	cassChunksPerRow = stats.NewMeter("cassandra.chunks_per_row", 0)
	cassRowsPerResponse = stats.NewMeter("cassandra.rows_per_response", 0)
	cassToIterDuration = stats.NewTimer("cassandra.to_iter_duration", 0)

	chunkSaveOk = stats.NewCount("chunks.save_ok")
	chunkSaveFail = stats.NewCount("chunks.save_fail")
	chunkSizeAtSave = stats.NewMeter("chunk_size.at_save", 0)
	chunkSizeAtLoad = stats.NewMeter("chunk_size.at_load", 0)
}

func (c *cassandraStore) Add(cwr *ChunkWriteRequest) {
	sum := 0
	for _, char := range cwr.key {
		sum += int(char)
	}
	which := sum % len(c.writeQueues)
	c.writeQueueMeters[which].Value(int64(len(c.writeQueues[which])))
	c.writeQueues[which] <- cwr
	c.writeQueueMeters[which].Value(int64(len(c.writeQueues[which])))
}

/* process writeQueue.
 */
func (c *cassandraStore) processWriteQueue(queue chan *ChunkWriteRequest, meter met.Meter) {
	tick := time.Tick(time.Duration(1) * time.Second)
	for {
		select {
		case <-tick:
			meter.Value(int64(len(queue)))
		case cwr := <-queue:
			meter.Value(int64(len(queue)))
			log.Debug("CS: starting to save %s:%d %v", cwr.key, cwr.chunk.T0, cwr.chunk)
			//log how long the chunk waited in the queue before we attempted to save to cassandra
			cassBlockDuration.Value(time.Now().Sub(cwr.timestamp))

			data := cwr.chunk.Series.Bytes()
			chunkSizeAtSave.Value(int64(len(data)))
			version := chunk.FormatStandardGoTsz
			buf := new(bytes.Buffer)
			binary.Write(buf, binary.LittleEndian, uint8(version))
			buf.Write(data)
			success := false
			attempts := 0
			for !success {
				err := c.insertChunk(cwr.key, cwr.chunk.T0, buf.Bytes(), int(cwr.ttl))
				if err == nil {
					success = true
					cwr.chunk.Saved = true
					SendPersistMessage(cwr.key, cwr.chunk.T0)
					log.Debug("CS: save complete. %s:%d %v", cwr.key, cwr.chunk.T0, cwr.chunk)
					chunkSaveOk.Inc(1)
				} else {
					if (attempts % 20) == 0 {
						log.Warn("CS: failed to save chunk to cassandra after %d attempts. %v, %s", attempts+1, cwr.chunk, err)
					}
					chunkSaveFail.Inc(1)
					sleepTime := 100 * attempts
					if sleepTime > 2000 {
						sleepTime = 2000
					}
					time.Sleep(time.Duration(sleepTime) * time.Millisecond)
					attempts++
				}
			}
		}
	}
}

// Insert Chunks into Cassandra.
//
// key: is the metric_id
// ts: is the start of the aggregated time range.
// data: is the payload as bytes.
func (c *cassandraStore) insertChunk(key string, t0 uint32, data []byte, ttl int) error {
	// for unit tests
	if c.session == nil {
		return nil
	}
	query := fmt.Sprintf("INSERT INTO metric (key, ts, data) values(?,?,?) USING TTL %d", ttl)
	row_key := fmt.Sprintf("%s_%d", key, t0/Month_sec) // "month number" based on unix timestamp (rounded down)
	pre := time.Now()
	ret := c.session.Query(query, row_key, t0, data).Exec()
	cassPutDuration.Value(time.Now().Sub(pre))
	return ret
}

type outcome struct {
	month   uint32
	sortKey uint32
	i       *gocql.Iter
}
type asc []outcome

func (o asc) Len() int           { return len(o) }
func (o asc) Swap(i, j int)      { o[i], o[j] = o[j], o[i] }
func (o asc) Less(i, j int) bool { return o[i].sortKey < o[j].sortKey }

func (c *cassandraStore) processReadQueue() {
	for crr := range c.readQueue {
		crr.out <- outcome{crr.month, crr.sortKey, c.session.Query(crr.q, crr.p...).Iter()}
	}
}

// Basic search of cassandra.
// start inclusive, end exclusive
func (c *cassandraStore) Search(key string, start, end uint32) ([]iter.Iter, error) {
	iters := make([]iter.Iter, 0)
	if start > end {
		return iters, errStartBeforeEnd
	}

	crrs := make([]*ChunkReadRequest, 0)

	query := func(month, sortKey uint32, q string, p ...interface{}) {
		crrs = append(crrs, &ChunkReadRequest{month, sortKey, q, p, nil})
	}

	start_month := start - (start % Month_sec)       // starting row has to be at, or before, requested start
	end_month := (end - 1) - ((end - 1) % Month_sec) // ending row has to be include the last point we might need

	pre := time.Now()

	// unfortunately in the database we only have the t0's of all chunks.
	// this means we can easily make sure to include the correct last chunk (just query for a t0 < end, the last chunk will contain the last needed data)
	// but it becomes hard to find which should be the first chunk to include. we can't just query for start <= t0 because than will miss some data at the beginning
	// we can't assume we know the chunkSpan so we can't just calculate the t0 >= start - <some-predefined-number> because chunkSpans may change over time.
	// we effectively need all chunks with a t0 > start, as well as the last chunk with a t0 <= start.
	// since we make sure that you can only use chunkSpans so that Month_sec % chunkSpan == 0, we know that this previous chunk will always be in the same row
	// as the one that has start_month.

	row_key := fmt.Sprintf("%s_%d", key, start_month/Month_sec)

	query(start_month, start_month, "SELECT ts, data FROM metric WHERE key=? AND ts <= ? Limit 1", row_key, start)

	if start_month == end_month {
		// we need a selection of the row between startTs and endTs
		row_key = fmt.Sprintf("%s_%d", key, start_month/Month_sec)
		query(start_month, start_month+1, "SELECT ts, data FROM metric WHERE key = ? AND ts > ? AND ts < ? ORDER BY ts ASC", row_key, start, end)
	} else {
		// get row_keys for each row we need to query.
		for month := start_month; month <= end_month; month += Month_sec {
			row_key = fmt.Sprintf("%s_%d", key, month/Month_sec)
			if month == start_month {
				// we want from startTs to the end of the row.
				query(month, month+1, "SELECT ts, data FROM metric WHERE key = ? AND ts >= ? ORDER BY ts ASC", row_key, start+1)
			} else if month == end_month {
				// we want from start of the row till the endTs.
				query(month, month, "SELECT ts, data FROM metric WHERE key = ? AND ts <= ? ORDER BY ts ASC", row_key, end-1)
			} else {
				// we want all columns
				query(month, month, "SELECT ts, data FROM metric WHERE key = ? ORDER BY ts ASC", row_key)
			}
		}
	}
	numQueries := len(crrs)
	results := make(chan outcome, numQueries)
	for i := range crrs {
		crrs[i].out = results
		c.readQueue <- crrs[i]
	}
	outcomes := make([]outcome, 0, numQueries)

	seen := 0
	for o := range results {
		seen += 1
		outcomes = append(outcomes, o)
		if seen == numQueries {
			close(results)
			cassGetDuration.Value(time.Now().Sub(pre))
			break
		}
	}
	pre = time.Now()
	// we have all of the results, but they could have arrived in any order.
	sort.Sort(asc(outcomes))

	var b []byte
	var ts int
	for _, outcome := range outcomes {
		chunks := int64(0)
		for outcome.i.Scan(&ts, &b) {
			chunks += 1
			chunkSizeAtLoad.Value(int64(len(b)))
			if len(b) < 2 {
				log.Error(3, errChunkTooSmall.Error())
				return iters, errChunkTooSmall
			}
			if chunk.Format(b[0]) != chunk.FormatStandardGoTsz {
				log.Error(3, errUnknownChunkFormat.Error())
				return iters, errUnknownChunkFormat
			}
			it, err := tsz.NewIterator(b[1:])
			if err != nil {
				log.Error(3, "failed to unpack cassandra payload. %s", err)
				return iters, err
			}
			iters = append(iters, iter.New(it, true))
		}
		cassChunksPerRow.Value(chunks)
		err := outcome.i.Close()
		if err != nil {
			log.Error(3, "cassandra query error. %s", err)
		}
	}
	cassToIterDuration.Value(time.Now().Sub(pre))
	cassRowsPerResponse.Value(int64(len(outcomes)))
	log.Debug("CS: searchCassandra(): %d outcomes (queries), %d total iters", len(outcomes), len(iters))
	return iters, nil
}

func (c *cassandraStore) Stop() {
	c.session.Close()
}
<|endoftext|>


<commit_before>package container

import (
	"encoding/json"
	"fmt"
	"io"
	"io/ioutil"
	"net/http"
	"net/url"

	"errors"
	"strings"

	"github.com/cloudfoundry-incubator/garden/api"
	"github.com/pivotal-cf-experimental/garden-dot-net/process"

	"code.google.com/p/go.net/websocket"
	"github.com/pivotal-golang/lager"
)

type container struct {
	containerizerURL url.URL
	handle           string
	logger           lager.Logger
}

type netInResponse struct {
	HostPort    uint32 `json:"hostPort"`
	ErrorString string `json:"error"`
}

type ProcessStreamEvent struct {
	MessageType    string          `json:"type"`
	ApiProcessSpec api.ProcessSpec `json:"pspec"`
	Data           string          `json:"data"`
}

func NewContainer(containerizerURL url.URL, handle string, logger lager.Logger) *container {
	return &container{
		containerizerURL: containerizerURL,
		handle:           handle,
		logger:           logger,
	}
}

func (container *container) Handle() string {
	return container.handle
}

func (container *container) Stop(kill bool) error {
	return nil
}

func (container *container) Info() (api.ContainerInfo, error) {
	url := container.containerizerURL.String() + "/api/containers/" + container.Handle() + "/properties"
	container.logger.Info("GETTING INFO", lager.Data{
		"url": url,
	})
	response, err := http.Get(url)
	if err != nil {
		container.logger.Info("ERROR GETTING PROPERTIES", lager.Data{
			"error": err,
		})
		return api.ContainerInfo{}, err
	}
	defer response.Body.Close()
	rawJSON, err := ioutil.ReadAll(response.Body)
	if err != nil {
		return api.ContainerInfo{}, err
	}
	properties := api.Properties{}
	err = json.Unmarshal(rawJSON, &properties)
	if err != nil {
		container.logger.Info("ERROR UNMARSHALING PROPERTIES", lager.Data{
			"error": err,
		})
		return api.ContainerInfo{}, nil
	}
	containerInfo := api.ContainerInfo{Properties: properties}
	container.logger.Info("CONTAINER INFO", lager.Data{
		"containerInfo": containerInfo,
	})

	return containerInfo, nil
}

func (container *container) StreamIn(dstPath string, tarStream io.Reader) error {
	url := container.containerizerURL.String() + "/api/containers/" + container.Handle() + "/files?destination=" + dstPath

	req, err := http.NewRequest("PUT", url, tarStream)
	if err != nil {
		return err
	}
	req.Header.Set("Content-Type", "application/octet-stream")
	resp, err := http.DefaultClient.Do(req)
	if err != nil {
		return err
	}

	resp.Body.Close()
	return nil
}

func (container *container) StreamOut(srcPath string) (io.ReadCloser, error) {
	url := container.containerizerURL.String() + "/api/containers/" + container.Handle() + "/files?source=" + srcPath
	resp, err := http.Get(url)
	if err != nil {
		return nil, err
	}
	return resp.Body, nil
}

func (container *container) LimitBandwidth(limits api.BandwidthLimits) error {
	return nil
}
func (container *container) CurrentBandwidthLimits() (api.BandwidthLimits, error) {
	return api.BandwidthLimits{}, nil
}

func (container *container) LimitCPU(limits api.CPULimits) error {
	return nil
}

func (container *container) CurrentCPULimits() (api.CPULimits, error) {
	return api.CPULimits{}, nil
}

func (container *container) LimitDisk(limits api.DiskLimits) error {
	return nil
}
func (container *container) CurrentDiskLimits() (api.DiskLimits, error) {
	return api.DiskLimits{}, nil
}

func (container *container) LimitMemory(limits api.MemoryLimits) error {
	return nil
}
func (container *container) CurrentMemoryLimits() (api.MemoryLimits, error) {
	return api.MemoryLimits{}, nil
}

func (container *container) NetIn(hostPort, containerPort uint32) (uint32, uint32, error) {
	url := container.containerizerURL.String() + "/api/containers/" + container.Handle() + "/net/in"
	response, err := http.Post(url, "application/json", strings.NewReader(fmt.Sprintf(`{"hostPort": %v}`, hostPort)))
	if err != nil {
		return 0, 0, err
	}
	responseBody, err := ioutil.ReadAll(response.Body)
	if err != nil {
		return 0, 0, err
	}

	var responseJSON netInResponse
	err = json.Unmarshal(responseBody, &responseJSON)
	if err != nil {
		return 0, 0, err
	}

	if responseJSON.ErrorString != "" {
		return 0, 0, errors.New(responseJSON.ErrorString)
	}

	return responseJSON.HostPort, containerPort, err
}

func (container *container) NetOut(network string, port uint32) error {
	return nil
}

func (container *container) containerizerWS() string {
	u2 := container.containerizerURL
	u2.Scheme = "ws"
	return u2.String()
}

func (container *container) Run(processSpec api.ProcessSpec, processIO api.ProcessIO) (api.Process, error) {
	origin := "http://localhost/"
	wsUri := container.containerizerWS() + "/api/containers/" + container.handle + "/run"
	ws, err := websocket.Dial(wsUri, "", origin)
	if err != nil {
		return nil, err
	}
	websocket.JSON.Send(ws, ProcessStreamEvent{
		MessageType:    "run",
		ApiProcessSpec: processSpec,
	})

	proc := process.NewDotNetProcess()

	streamWebsocketIOToContainerizer(ws, processIO)
	go func() {
		err := streamWebsocketIOFromContainerizer(ws, processIO)
		if err != nil {
			proc.StreamOpen <- err.Error()
		}
		close(proc.StreamOpen)
	}()

	return proc, nil
}

func (container *container) Attach(uint32, api.ProcessIO) (api.Process, error) {
	return process.NewDotNetProcess(), nil
}

func (container *container) GetProperty(name string) (string, error) {
	url := container.containerizerURL.String() + "/api/containers/" + container.Handle() + "/properties/" + name
	response, err := http.Get(url)
	defer response.Body.Close()
	if err != nil {
		return "", err
	}
	property, err := ioutil.ReadAll(response.Body)
	if err != nil {
		return "", err
	}
	return string(property), nil
}

func (container *container) SetProperty(name string, value string) error {
	url := container.containerizerURL.String() + "/api/containers/" + container.Handle() + "/properties/" + name

	container.logger.Info("SETTING PROPERTY", lager.Data{
		"url": url,
	})
	request, err := http.NewRequest("PUT", url, strings.NewReader("[\""+value+"\"]"))
	if err != nil {
		return err
	}
	request.Header.Set("Content-Type", "application/json")
	response, err := http.DefaultClient.Do(request)
	if err != nil {
		return err
	}

	response.Body.Close()
	return nil
}

func (container *container) RemoveProperty(name string) error {
	url := container.containerizerURL.String() + "/api/containers/" + container.Handle() + "/properties/" + name

	request, err := http.NewRequest("DELETE", url, strings.NewReader(""))
	if err != nil {
		return err
	}
	response, err := http.DefaultClient.Do(request)
	if err != nil {
		return err
	}

	response.Body.Close()
	return nil
}

func streamWebsocketIOToContainerizer(ws *websocket.Conn, processIO api.ProcessIO) {
	if processIO.Stdin != nil {
		fiw := faninWriter{
			hasSink: make(chan struct{}),
		}
		fiw.AddSink(ws)
		fiw.AddSource(processIO.Stdin)
	}
}

func streamWebsocketIOFromContainerizer(ws *websocket.Conn, processIO api.ProcessIO) error {
	receiveStream := ProcessStreamEvent{}
	for {
		err := websocket.JSON.Receive(ws, &receiveStream)
		if err != nil {
			return err
		}

		if receiveStream.MessageType == "stdout" && processIO.Stdout != nil {
			io.WriteString(processIO.Stdout, receiveStream.Data)
		}
		if receiveStream.MessageType == "stderr" && processIO.Stderr != nil {
			io.WriteString(processIO.Stderr, receiveStream.Data)
		}

		if receiveStream.MessageType == "error" {
			return errors.New(receiveStream.Data)
		}
		if receiveStream.MessageType == "close" {
			return nil
		}
	}
}
<commit_msg>Handle urls properly [#79974132]<commit_after>package container

import (
	"encoding/json"
	"fmt"
	"io"
	"io/ioutil"
	"net/http"
	"net/url"

	"errors"
	"strings"

	"github.com/cloudfoundry-incubator/garden/api"
	"github.com/pivotal-cf-experimental/garden-dot-net/process"

	"code.google.com/p/go.net/websocket"
	"github.com/pivotal-golang/lager"
)

type container struct {
	containerizerURL url.URL
	handle           string
	logger           lager.Logger
}

type netInResponse struct {
	HostPort    uint32 `json:"hostPort"`
	ErrorString string `json:"error"`
}

type ProcessStreamEvent struct {
	MessageType    string          `json:"type"`
	ApiProcessSpec api.ProcessSpec `json:"pspec"`
	Data           string          `json:"data"`
}

func NewContainer(containerizerURL url.URL, handle string, logger lager.Logger) *container {
	return &container{
		containerizerURL: containerizerURL,
		handle:           handle,
		logger:           logger,
	}
}

func (container *container) Handle() string {
	return container.handle
}

func (container *container) Stop(kill bool) error {
	return nil
}

func (container *container) Info() (api.ContainerInfo, error) {
	url := container.containerizerURL.String() + "/api/containers/" + container.Handle() + "/properties"
	container.logger.Info("GETTING INFO", lager.Data{
		"url": url,
	})
	response, err := http.Get(url)
	if err != nil {
		container.logger.Info("ERROR GETTING PROPERTIES", lager.Data{
			"error": err,
		})
		return api.ContainerInfo{}, err
	}
	defer response.Body.Close()
	rawJSON, err := ioutil.ReadAll(response.Body)
	if err != nil {
		return api.ContainerInfo{}, err
	}
	properties := api.Properties{}
	err = json.Unmarshal(rawJSON, &properties)
	if err != nil {
		container.logger.Info("ERROR UNMARSHALING PROPERTIES", lager.Data{
			"error": err,
		})
		return api.ContainerInfo{}, nil
	}
	containerInfo := api.ContainerInfo{Properties: properties}
	container.logger.Info("CONTAINER INFO", lager.Data{
		"containerInfo": containerInfo,
	})

	return containerInfo, nil
}

func (container *container) StreamIn(dstPath string, tarStream io.Reader) error {
	url := container.containerizerURL.String() + "/api/containers/" + container.Handle() + "/files?destination=" + dstPath

	req, err := http.NewRequest("PUT", url, tarStream)
	if err != nil {
		return err
	}
	req.Header.Set("Content-Type", "application/octet-stream")
	resp, err := http.DefaultClient.Do(req)
	if err != nil {
		return err
	}

	resp.Body.Close()
	return nil
}

func (container *container) StreamOut(srcPath string) (io.ReadCloser, error) {
	url := container.containerizerURL.String() + "/api/containers/" + container.Handle() + "/files?source=" + srcPath
	resp, err := http.Get(url)
	if err != nil {
		return nil, err
	}
	return resp.Body, nil
}

func (container *container) LimitBandwidth(limits api.BandwidthLimits) error {
	return nil
}
func (container *container) CurrentBandwidthLimits() (api.BandwidthLimits, error) {
	return api.BandwidthLimits{}, nil
}

func (container *container) LimitCPU(limits api.CPULimits) error {
	return nil
}

func (container *container) CurrentCPULimits() (api.CPULimits, error) {
	return api.CPULimits{}, nil
}

func (container *container) LimitDisk(limits api.DiskLimits) error {
	return nil
}
func (container *container) CurrentDiskLimits() (api.DiskLimits, error) {
	return api.DiskLimits{}, nil
}

func (container *container) LimitMemory(limits api.MemoryLimits) error {
	return nil
}
func (container *container) CurrentMemoryLimits() (api.MemoryLimits, error) {
	return api.MemoryLimits{}, nil
}

func (container *container) NetIn(hostPort, containerPort uint32) (uint32, uint32, error) {
	url := container.containerizerURL.String() + "/api/containers/" + container.Handle() + "/net/in"
	response, err := http.Post(url, "application/json", strings.NewReader(fmt.Sprintf(`{"hostPort": %v}`, hostPort)))
	if err != nil {
		return 0, 0, err
	}
	responseBody, err := ioutil.ReadAll(response.Body)
	if err != nil {
		return 0, 0, err
	}

	var responseJSON netInResponse
	err = json.Unmarshal(responseBody, &responseJSON)
	if err != nil {
		return 0, 0, err
	}

	if responseJSON.ErrorString != "" {
		return 0, 0, errors.New(responseJSON.ErrorString)
	}

	return responseJSON.HostPort, containerPort, err
}

func (container *container) NetOut(network string, port uint32) error {
	return nil
}

func (container *container) containerizerWS() string {
	u2 := container.containerizerURL
	u2.Scheme = "ws"
	return u2.String()
}

func (container *container) Run(processSpec api.ProcessSpec, processIO api.ProcessIO) (api.Process, error) {
	origin := "http://localhost/"
	wsUri := container.containerizerWS() + "/api/containers/" + container.handle + "/run"
	ws, err := websocket.Dial(wsUri, "", origin)
	if err != nil {
		return nil, err
	}
	websocket.JSON.Send(ws, ProcessStreamEvent{
		MessageType:    "run",
		ApiProcessSpec: processSpec,
	})

	proc := process.NewDotNetProcess()

	streamWebsocketIOToContainerizer(ws, processIO)
	go func() {
		err := streamWebsocketIOFromContainerizer(ws, processIO)
		if err != nil {
			proc.StreamOpen <- err.Error()
		}
		close(proc.StreamOpen)
	}()

	return proc, nil
}

func (container *container) Attach(uint32, api.ProcessIO) (api.Process, error) {
	return process.NewDotNetProcess(), nil
}

func (container *container) GetProperty(name string) (string, error) {
	name = strings.Replace(name, ":", "♥", -1)
	requestUrl := container.containerizerURL
	requestUrl.Path += "/api/containers/" + container.Handle() + "/properties/" + name
	container.logger.Info("GET PROPERTY", lager.Data{
		"property": name,
		"url":      requestUrl,
	})
	response, err := http.Get(requestUrl.String())
	defer response.Body.Close()

	if err != nil {
		container.logger.Info("NO GET FOR YOU!", lager.Data{
			"error": err,
		})
		return "", err
	}
	property, err := ioutil.ReadAll(response.Body)
	if err != nil {
		return "", err
	}
	return string(property), nil
}

func (container *container) SetProperty(name string, value string) error {
	name = strings.Replace(name, ":", "♥", -1)
	requestUrl := container.containerizerURL
	requestUrl.Path += "/api/containers/" + container.Handle() + "/properties/" + name
	container.logger.Info("SET PROPERTY", lager.Data{
		"property": name,
		"url":      requestUrl,
	})
	request, err := http.NewRequest("PUT", requestUrl.String(), strings.NewReader("[\""+value+"\"]"))
	if err != nil {
		return err
	}
	request.Header.Set("Content-Type", "application/json")
	response, err := http.DefaultClient.Do(request)
	if err != nil {
		return err
	}

	response.Body.Close()
	return nil
}

func (container *container) RemoveProperty(name string) error {
	name = strings.Replace(name, ":", "♥", -1)
	requestUrl := container.containerizerURL
	requestUrl.Path += "/api/containers/" + container.Handle() + "/properties/" + name
	container.logger.Info("REMOVING PROPERTY", lager.Data{
		"property": name,
		"url":      requestUrl,
	})
	request, err := http.NewRequest("DELETE", requestUrl.String(), strings.NewReader(""))
	if err != nil {
		return err
	}
	response, err := http.DefaultClient.Do(request)
	if err != nil {
		return err
	}

	response.Body.Close()
	return nil
}

func streamWebsocketIOToContainerizer(ws *websocket.Conn, processIO api.ProcessIO) {
	if processIO.Stdin != nil {
		fiw := faninWriter{
			hasSink: make(chan struct{}),
		}
		fiw.AddSink(ws)
		fiw.AddSource(processIO.Stdin)
	}
}

func streamWebsocketIOFromContainerizer(ws *websocket.Conn, processIO api.ProcessIO) error {
	receiveStream := ProcessStreamEvent{}
	for {
		err := websocket.JSON.Receive(ws, &receiveStream)
		if err != nil {
			return err
		}

		if receiveStream.MessageType == "stdout" && processIO.Stdout != nil {
			io.WriteString(processIO.Stdout, receiveStream.Data)
		}
		if receiveStream.MessageType == "stderr" && processIO.Stderr != nil {
			io.WriteString(processIO.Stderr, receiveStream.Data)
		}

		if receiveStream.MessageType == "error" {
			return errors.New(receiveStream.Data)
		}
		if receiveStream.MessageType == "close" {
			return nil
		}
	}
}
<|endoftext|>


<commit_before>package container

import (
	"errors"
	"fmt"
	. "github.com/byrnedo/capitan/consts"
	"github.com/byrnedo/capitan/helpers"
	"github.com/byrnedo/capitan/logger"
	. "github.com/byrnedo/capitan/logger"
	"github.com/codeskyblue/go-sh"
	"io/ioutil"
	"math/rand"
	"os"
	"strings"
	"sync"
	"time"
)

var (
	colorList = []string{
		"white",
		"red",
		"green",
		"yellow",
		"blue",
		"magenta",
		"cyan",
	}

	nextColorIndex = rand.Intn(len(colorList) - 1)
)

// Get the next color to be used in log output
func nextColor() string {
	defer func() {
		nextColorIndex++
		if nextColorIndex >= len(colorList) {
			nextColorIndex = 0
		}
	}()
	return colorList[nextColorIndex]
}

type Link struct {
	Container string
	Alias     string
}

type AppliedAction string

const (
	Run     AppliedAction = "run"
	Start   AppliedAction = "start"
	Stop    AppliedAction = "stop"
	Kill    AppliedAction = "kill"
	Restart AppliedAction = "restart"
	Remove  AppliedAction = "remove"
)

type Hook struct {
	Scripts []string
	Ses     *ShellSession
}

type Hooks map[string]*Hook

// Runs a hook command if it exists for a specific container
func (h Hooks) Run(hookName string, ctr *Container) error {
	var (
		hook  *Hook
		found bool
		err   error
	)

	if hook, found = h[hookName]; !found {
		return nil
	}

	for _, script := range hook.Scripts {
		hook.Ses = NewContainerShellSession(ctr)
		hook.Ses.SetEnv("CAPITAN_HOOK_NAME", hookName)

		hook.Ses.Command("bash", "-c", script)

		hook.Ses.Stdout = os.Stdout
		hook.Ses.Stderr = os.Stderr
		hook.Ses.Stdin = os.Stdin

		if err = hook.Ses.Run(); err != nil {
			return err
		}
	}
	return nil
}

type Container struct {
	// Container name
	Name string
	// name of service (not including number)
	ServiceName string
	// non unique service id, eg the first col in config, "mongo" or "php"
	ServiceType string
	// the order defined in the config output
	Placement int
	// arguments to container
	ContainerArgs []string
	// image to use
	Image string
	// if supplied will do docker build on this path
	Build string
	// The arguments for the build command
	BuildArgs []string
	// command for container
	Command []string
	// links
	Links []Link
	// volumes from list
	VolumesFrom []string
	// hooks map for this definition
	Hooks Hooks
	// used in commands
	Action AppliedAction
	// the total number of containers to scale to.
	Scale int
	// the arguments for docker run / create
	RunArguments []interface{}
	// the project name
	ProjectName string
	// the project name separator, usually "_"
	ProjectNameSeparator string
	// the number of this container, relates to scale
	InstanceNumber int
	// Rm command given, therefore dont run as daemon
	Remove bool
}

// Builds an image for a container
func (set *Container) BuildImage() error {
	if err := set.Hooks.Run("before.build", set); err != nil {
		return err
	}

	args := append([]interface{}{
		"build",
	}, helpers.ToInterfaceSlice(set.BuildArgs)...)

	args = append(args, "--tag", set.Image, set.Build)

	if _, err := helpers.RunCmd(args...); err != nil {
		return err
	}
	if err := set.Hooks.Run("after.build", set); err != nil {
		return err
	}
	return nil
}

func (set *Container) runInForeground(cmd []interface{}, wg *sync.WaitGroup) error {

	var (
		ses *ShellSession
		err error
	)

	beforeStart := time.Now()

	initialArgs := []interface{}{
		"run",
		"-a", "stdout",
		"-a", "stderr",
		"-a", "stdin",
		"--sig-proxy=false",
	}

	if set.Remove {
		initialArgs = append(initialArgs, "--rm")
	}

	cmd = append(initialArgs, cmd...)
	if ses, err = set.startLoggedCommand(cmd); err != nil {
		return err
	}
	wg.Add(1)
	go func(name string) {
		ses.Wait()
		wg.Done()
	}(set.Name)

	if !helpers.WasContainerStartedAfterOrRetry(set.Name, beforeStart, 10, 200*time.Millisecond) {
		return errors.New(set.Name + " failed to start")
	}

	Debug.Println("Container deemed to have started after", beforeStart)

	if !helpers.ContainerIsRunning(set.Name) {
		exitCode := helpers.ContainerExitCode(set.Name)
		if exitCode != "0" {
			return errors.New(set.Name + " exited with non-zero exit code " + exitCode)
		}
	}

	return nil

}

func (set *Container) RecreateAndRun(attach bool, dryRun bool, wg *sync.WaitGroup) error {
	if !dryRun {
		set.Rm([]string{"-f"})
	}

	if err := set.Run(attach, dryRun, wg); err != nil {
		return err
	}
	return nil
}

func createCapitanContainerLabels(ctr *Container, args []interface{}) []interface{} {
	return []interface{}{
		"--label",
		UniqueLabelName + "=" + fmt.Sprintf("'%s'", args),
		"--label",
		ServiceLabelName + "=" + ctr.ServiceName,
		"--label",
		ProjectLabelName + "=" + ctr.ProjectName,
	}
}

// Run a container
func (set *Container) Create(dryRun bool) error {
	set.Action = Run

	Info.Println("Creating " + set.Name)
	if dryRun {
		return nil
	}
	if err := set.Hooks.Run("before.create", set); err != nil {
		return err
	}

	cmd := set.RunArguments
	labels := createCapitanContainerLabels(set, cmd)
	cmd = append(labels, cmd...)

	cmd = append([]interface{}{"create"}, cmd...)
	if err := set.launchDaemonCommand(cmd); err != nil {
		return err
	}

	return set.Hooks.Run("after.create", set)
}

// Run a container
func (set *Container) Run(attach bool, dryRun bool, wg *sync.WaitGroup) error {
	set.Action = Run

	Info.Println("Running " + set.Name)
	if dryRun {
		return nil
	}
	if err := set.Hooks.Run("before.run", set); err != nil {
		return err
	}

	cmd := set.RunArguments
	labels := createCapitanContainerLabels(set, cmd)
	cmd = append(labels, cmd...)

	if attach || set.Remove {

		if err := set.runInForeground(cmd, wg); err != nil {
			return err
		}

	} else {
		cmd = append([]interface{}{"run", "-d"}, cmd...)
		if err := set.launchDaemonCommand(cmd); err != nil {
			return err
		}
	}

	return set.Hooks.Run("after.run", set)
}

func (set *Container) launchDaemonCommand(cmd []interface{}) error {
	var (
		ses *ShellSession
		err error
	)
	ses = NewContainerShellSession(set)

	concStr := "docker "
	for _, arg := range cmd {
		concStr += fmt.Sprintf("%s", arg) + " "
	}
	concStr = strings.Trim(concStr, " ")

	err = ses.Command("bash", "-c", concStr).Run()
	return err
}

func (set *Container) startLoggedCommand(cmd []interface{}) (*ShellSession, error) {
	ses := NewContainerShellSession(set)

	color := nextColor()
	ses.Stdout = NewContainerLogWriter(os.Stdout, set.Name, color)
	ses.Stderr = NewContainerLogWriter(os.Stderr, set.Name, color)

	concStr := "docker "
	for _, arg := range cmd {
		concStr += fmt.Sprintf("%s", arg) + " "
	}
	concStr = strings.Trim(concStr, " ")

	err := ses.Command("bash", "-c", concStr).Start()

	return ses, err
}

// Create docker arg slice from container options
func (set *Container) GetRunArguments() []interface{} {
	imageName := set.Name
	if len(set.Image) > 0 {
		imageName = set.Image
	}

	var linkArgs = make([]interface{}, 0, len(set.Links)*2)
	for _, link := range set.Links {
		linkStr := link.Container
		if link.Alias != "" {
			linkStr += ":" + link.Alias
		}
		linkArgs = append(linkArgs, "--link", linkStr)
	}

	var volumesFromArgs = make([]interface{}, 0, len(set.VolumesFrom)*2)
	for _, vol := range set.VolumesFrom {
		volumesFromArgs = append(volumesFromArgs, "--volumes-from", vol)
	}

	cmd := append([]interface{}{"--name", set.Name}, helpers.ToInterfaceSlice(set.ContainerArgs)...)
	cmd = append(cmd, linkArgs...)
	cmd = append(cmd, volumesFromArgs...)
	cmd = append(cmd, imageName)
	cmd = append(cmd, helpers.ToInterfaceSlice(set.Command)...)
	return cmd
}

func (set *Container) Attach(wg *sync.WaitGroup) error {
	var (
		err error
		ses *ShellSession
	)
	if ses, err = set.startLoggedCommand(append([]interface{}{"attach", "--sig-proxy=false"}, set.Name)); err != nil {
		return err
	}
	wg.Add(1)

	go func(name string) {
		ses.Wait()
		wg.Done()
	}(set.Name)
	return nil
}

// Start a given container
// TODO needs to respect scale
func (set *Container) Start(attach bool, wg *sync.WaitGroup) error {
	var (
		err error
	)
	set.Action = Start
	if helpers.ContainerIsRunning(set.Name) {
		Info.Println("Already running", set.Name)
		if attach {
			if err = set.Attach(wg); err != nil {
				return err
			}
		}
		return nil
	}

	if err = set.Hooks.Run("before.start", set); err != nil {
		return err
	}

	if err = set.launchDaemonCommand(append([]interface{}{"start"}, set.Name)); err != nil {
		return err
	}
	if attach {
		if err = set.Attach(wg); err != nil {
			return err
		}
	}

	if err := set.Hooks.Run("after.start", set); err != nil {
		return err
	}
	return nil
}

// Restart the container
// TODO needs to respect scale
func (set *Container) Restart(args []string) error {
	set.Action = Restart
	if err := set.Hooks.Run("before.start", set); err != nil {
		return err
	}
	args = append(args, set.Name)
	if _, err := helpers.RunCmd(append([]interface{}{"restart"}, helpers.ToInterfaceSlice(args)...)...); err != nil {
		return err
	}
	if err := set.Hooks.Run("after.start", set); err != nil {
		return err
	}
	return nil
}

// Returns a containers IP
// TODO needs to respect scale
func (set *Container) IP() string {
	ses := sh.NewSession()
	ses.Stderr = ioutil.Discard
	out, err := ses.Command("docker", "inspect", "--type", "container", "--format", "{{.NetworkSettings.IPAddress}}", set.Name).Output()
	if err != nil {
		return ""
	}
	ip := strings.Trim(string(out), " \n")
	return ip
}

// Start streaming a container's logs
// TODO needs to respect scale
func (set *Container) Logs() (*sh.Session, error) {
	color := nextColor()
	ses := sh.NewSession()

	if logger.GetLevel() == DebugLevel {
		ses.ShowCMD = true
	}
	ses.Command("docker", "logs", "--tail", "10", "-f", set.Name)

	ses.Stdout = NewContainerLogWriter(os.Stdout, set.Name, color)
	ses.Stderr = NewContainerLogWriter(os.Stderr, set.Name, color)

	err := ses.Start()
	return ses, err
}

// Kills the container
// TODO needs to respect scale
func (set *Container) Kill(args []string) error {
	set.Action = Kill
	if err := set.Hooks.Run("before.kill", set); err != nil {
		return err
	}
	args = append(args, set.Name)
	if _, err := helpers.RunCmd(append([]interface{}{"kill"}, helpers.ToInterfaceSlice(args)...)...); err != nil {
		return err
	}
	if err := set.Hooks.Run("after.kill", set); err != nil {
		return err
	}
	return nil

}

// Stops the container
// TODO needs to respect scale
func (set *Container) Stop(args []string) error {
	set.Action = Stop
	if err := set.Hooks.Run("before.stop", set); err != nil {
		return err
	}
	args = append(args, set.Name)
	if _, err := helpers.RunCmd(append([]interface{}{"stop"}, helpers.ToInterfaceSlice(args)...)...); err != nil {
		return err
	}
	if err := set.Hooks.Run("after.stop", set); err != nil {
		return err
	}
	return nil
}

// Removes the container
// TODO needs to respect scale
func (set *Container) Rm(args []string) error {

	set.Action = Remove
	if err := set.Hooks.Run("before.rm", set); err != nil {
		return err
	}
	args = append(args, set.Name)
	if _, err := helpers.RunCmd(append([]interface{}{"rm"}, helpers.ToInterfaceSlice(args)...)...); err != nil {
		return err
	}
	if err := set.Hooks.Run("after.rm", set); err != nil {
		return err
	}
	return nil
}
<commit_msg>Better handling of rm command<commit_after>package container

import (
	"errors"
	"fmt"
	. "github.com/byrnedo/capitan/consts"
	"github.com/byrnedo/capitan/helpers"
	"github.com/byrnedo/capitan/logger"
	. "github.com/byrnedo/capitan/logger"
	"github.com/codeskyblue/go-sh"
	"io/ioutil"
	"math/rand"
	"os"
	"strings"
	"sync"
	"time"
)

var (
	colorList = []string{
		"white",
		"red",
		"green",
		"yellow",
		"blue",
		"magenta",
		"cyan",
	}

	nextColorIndex = rand.Intn(len(colorList) - 1)
)

// Get the next color to be used in log output
func nextColor() string {
	defer func() {
		nextColorIndex++
		if nextColorIndex >= len(colorList) {
			nextColorIndex = 0
		}
	}()
	return colorList[nextColorIndex]
}

type Link struct {
	Container string
	Alias     string
}

type AppliedAction string

const (
	Run     AppliedAction = "run"
	Start   AppliedAction = "start"
	Stop    AppliedAction = "stop"
	Kill    AppliedAction = "kill"
	Restart AppliedAction = "restart"
	Remove  AppliedAction = "remove"
)

type Hook struct {
	Scripts []string
	Ses     *ShellSession
}

type Hooks map[string]*Hook

// Runs a hook command if it exists for a specific container
func (h Hooks) Run(hookName string, ctr *Container) error {
	var (
		hook  *Hook
		found bool
		err   error
	)

	if hook, found = h[hookName]; !found {
		return nil
	}

	for _, script := range hook.Scripts {
		hook.Ses = NewContainerShellSession(ctr)
		hook.Ses.SetEnv("CAPITAN_HOOK_NAME", hookName)

		hook.Ses.Command("bash", "-c", script)

		hook.Ses.Stdout = os.Stdout
		hook.Ses.Stderr = os.Stderr
		hook.Ses.Stdin = os.Stdin

		if err = hook.Ses.Run(); err != nil {
			return err
		}
	}
	return nil
}

type Container struct {
	// Container name
	Name string
	// name of service (not including number)
	ServiceName string
	// non unique service id, eg the first col in config, "mongo" or "php"
	ServiceType string
	// the order defined in the config output
	Placement int
	// arguments to container
	ContainerArgs []string
	// image to use
	Image string
	// if supplied will do docker build on this path
	Build string
	// The arguments for the build command
	BuildArgs []string
	// command for container
	Command []string
	// links
	Links []Link
	// volumes from list
	VolumesFrom []string
	// hooks map for this definition
	Hooks Hooks
	// used in commands
	Action AppliedAction
	// the total number of containers to scale to.
	Scale int
	// the arguments for docker run / create
	RunArguments []interface{}
	// the project name
	ProjectName string
	// the project name separator, usually "_"
	ProjectNameSeparator string
	// the number of this container, relates to scale
	InstanceNumber int
	// Rm command given, therefore dont run as daemon
	Remove bool
}

// Builds an image for a container
func (set *Container) BuildImage() error {
	if err := set.Hooks.Run("before.build", set); err != nil {
		return err
	}

	args := append([]interface{}{
		"build",
	}, helpers.ToInterfaceSlice(set.BuildArgs)...)

	args = append(args, "--tag", set.Image, set.Build)

	if _, err := helpers.RunCmd(args...); err != nil {
		return err
	}
	if err := set.Hooks.Run("after.build", set); err != nil {
		return err
	}
	return nil
}

func (set *Container) launchWithRmInForeground(cmd []interface{}) error {
	var (
		ses *ShellSession
		err error
	)

	initialArgs := []interface{}{
		"run",
		"-a", "stdout",
		"-a", "stderr",
		"-a", "stdin",
		"--sig-proxy=false",
		"--rm",
	}

	cmd = append(initialArgs, cmd...)
	if ses, err = set.startLoggedCommand(cmd); err != nil {
		return err
	}

	err = ses.Wait()
	if err != nil {
		return errors.New(set.Name + " exited with error: " + err.Error())
	}

	return nil
}

func (set *Container) launchInForeground(cmd []interface{}, wg *sync.WaitGroup) error {

	var (
		ses *ShellSession
		err error
	)

	beforeStart := time.Now()

	initialArgs := []interface{}{
		"run",
		"-a", "stdout",
		"-a", "stderr",
		"-a", "stdin",
		"--sig-proxy=false",
	}

	if set.Remove {
		initialArgs = append(initialArgs, "--rm")
	}

	cmd = append(initialArgs, cmd...)
	if ses, err = set.startLoggedCommand(cmd); err != nil {
		return err
	}
	wg.Add(1)
	go func(name string) {
		ses.Wait()
		wg.Done()
	}(set.Name)

	if !helpers.WasContainerStartedAfterOrRetry(set.Name, beforeStart, 10, 200 * time.Millisecond) {
		return errors.New(set.Name + " failed to start")
	}

	Debug.Println("Container deemed to have started after", beforeStart)
	if !helpers.ContainerIsRunning(set.Name) {
		exitCode := helpers.ContainerExitCode(set.Name)
		if exitCode != "0" {
			return errors.New(set.Name + " exited with non-zero exit code " + exitCode)
		}
	}

	return nil

}

func (set *Container) RecreateAndRun(attach bool, dryRun bool, wg *sync.WaitGroup) error {
	if !dryRun {
		set.Rm([]string{"-f"})
	}

	if err := set.Run(attach, dryRun, wg); err != nil {
		return err
	}
	return nil
}

func createCapitanContainerLabels(ctr *Container, args []interface{}) []interface{} {
	return []interface{}{
		"--label",
		UniqueLabelName + "=" + fmt.Sprintf("'%s'", args),
		"--label",
		ServiceLabelName + "=" + ctr.ServiceName,
		"--label",
		ProjectLabelName + "=" + ctr.ProjectName,
	}
}

// Run a container
func (set *Container) Create(dryRun bool) error {
	set.Action = Run

	Info.Println("Creating " + set.Name)
	if dryRun {
		return nil
	}
	if err := set.Hooks.Run("before.create", set); err != nil {
		return err
	}

	cmd := set.RunArguments
	labels := createCapitanContainerLabels(set, cmd)
	cmd = append(labels, cmd...)

	cmd = append([]interface{}{"create"}, cmd...)
	if err := set.launchDaemonCommand(cmd); err != nil {
		return err
	}

	return set.Hooks.Run("after.create", set)
}

// Run a container
func (set *Container) Run(attach bool, dryRun bool, wg *sync.WaitGroup) error {
	set.Action = Run

	Info.Println("Running " + set.Name)
	if dryRun {
		return nil
	}
	if err := set.Hooks.Run("before.run", set); err != nil {
		return err
	}

	cmd := set.RunArguments
	labels := createCapitanContainerLabels(set, cmd)
	cmd = append(labels, cmd...)

	if set.Remove {
		if err := set.launchWithRmInForeground(cmd); err != nil {
			return err
		}
	} else if attach {
		if err := set.launchInForeground(cmd, wg); err != nil {
			return err
		}
	} else {
		cmd = append([]interface{}{"run", "-d"}, cmd...)
		if err := set.launchDaemonCommand(cmd); err != nil {
			return err
		}
	}

	return set.Hooks.Run("after.run", set)
}

func (set *Container) launchDaemonCommand(cmd []interface{}) error {
	var (
		ses *ShellSession
		err error
	)
	ses = NewContainerShellSession(set)

	concStr := "docker "
	for _, arg := range cmd {
		concStr += fmt.Sprintf("%s", arg) + " "
	}
	concStr = strings.Trim(concStr, " ")

	err = ses.Command("bash", "-c", concStr).Run()
	return err
}

func (set *Container) startLoggedCommand(cmd []interface{}) (*ShellSession, error) {
	ses := NewContainerShellSession(set)

	color := nextColor()
	ses.Stdout = NewContainerLogWriter(os.Stdout, set.Name, color)
	ses.Stderr = NewContainerLogWriter(os.Stderr, set.Name, color)

	concStr := "docker "
	for _, arg := range cmd {
		concStr += fmt.Sprintf("%s", arg) + " "
	}
	concStr = strings.Trim(concStr, " ")

	err := ses.Command("bash", "-c", concStr).Start()

	return ses, err
}

// Create docker arg slice from container options
func (set *Container) GetRunArguments() []interface{} {
	imageName := set.Name
	if len(set.Image) > 0 {
		imageName = set.Image
	}

	var linkArgs = make([]interface{}, 0, len(set.Links)*2)
	for _, link := range set.Links {
		linkStr := link.Container
		if link.Alias != "" {
			linkStr += ":" + link.Alias
		}
		linkArgs = append(linkArgs, "--link", linkStr)
	}

	var volumesFromArgs = make([]interface{}, 0, len(set.VolumesFrom)*2)
	for _, vol := range set.VolumesFrom {
		volumesFromArgs = append(volumesFromArgs, "--volumes-from", vol)
	}

	cmd := append([]interface{}{"--name", set.Name}, helpers.ToInterfaceSlice(set.ContainerArgs)...)
	cmd = append(cmd, linkArgs...)
	cmd = append(cmd, volumesFromArgs...)
	cmd = append(cmd, imageName)
	cmd = append(cmd, helpers.ToInterfaceSlice(set.Command)...)
	return cmd
}

func (set *Container) Attach(wg *sync.WaitGroup) error {
	var (
		err error
		ses *ShellSession
	)
	if ses, err = set.startLoggedCommand(append([]interface{}{"attach", "--sig-proxy=false"}, set.Name)); err != nil {
		return err
	}
	wg.Add(1)

	go func(name string) {
		ses.Wait()
		wg.Done()
	}(set.Name)
	return nil
}

// Start a given container
// TODO needs to respect scale
func (set *Container) Start(attach bool, wg *sync.WaitGroup) error {
	var (
		err error
	)
	set.Action = Start
	if helpers.ContainerIsRunning(set.Name) {
		Info.Println("Already running", set.Name)
		if attach {
			if err = set.Attach(wg); err != nil {
				return err
			}
		}
		return nil
	}

	if err = set.Hooks.Run("before.start", set); err != nil {
		return err
	}

	if err = set.launchDaemonCommand(append([]interface{}{"start"}, set.Name)); err != nil {
		return err
	}
	if attach {
		if err = set.Attach(wg); err != nil {
			return err
		}
	}

	if err := set.Hooks.Run("after.start", set); err != nil {
		return err
	}
	return nil
}

// Restart the container
// TODO needs to respect scale
func (set *Container) Restart(args []string) error {
	set.Action = Restart
	if err := set.Hooks.Run("before.start", set); err != nil {
		return err
	}
	args = append(args, set.Name)
	if _, err := helpers.RunCmd(append([]interface{}{"restart"}, helpers.ToInterfaceSlice(args)...)...); err != nil {
		return err
	}
	if err := set.Hooks.Run("after.start", set); err != nil {
		return err
	}
	return nil
}

// Returns a containers IP
// TODO needs to respect scale
func (set *Container) IP() string {
	ses := sh.NewSession()
	ses.Stderr = ioutil.Discard
	out, err := ses.Command("docker", "inspect", "--type", "container", "--format", "{{.NetworkSettings.IPAddress}}", set.Name).Output()
	if err != nil {
		return ""
	}
	ip := strings.Trim(string(out), " \n")
	return ip
}

// Start streaming a container's logs
// TODO needs to respect scale
func (set *Container) Logs() (*sh.Session, error) {
	color := nextColor()
	ses := sh.NewSession()

	if logger.GetLevel() == DebugLevel {
		ses.ShowCMD = true
	}
	ses.Command("docker", "logs", "--tail", "10", "-f", set.Name)

	ses.Stdout = NewContainerLogWriter(os.Stdout, set.Name, color)
	ses.Stderr = NewContainerLogWriter(os.Stderr, set.Name, color)

	err := ses.Start()
	return ses, err
}

// Kills the container
// TODO needs to respect scale
func (set *Container) Kill(args []string) error {
	set.Action = Kill
	if err := set.Hooks.Run("before.kill", set); err != nil {
		return err
	}
	args = append(args, set.Name)
	if _, err := helpers.RunCmd(append([]interface{}{"kill"}, helpers.ToInterfaceSlice(args)...)...); err != nil {
		return err
	}
	if err := set.Hooks.Run("after.kill", set); err != nil {
		return err
	}
	return nil

}

// Stops the container
// TODO needs to respect scale
func (set *Container) Stop(args []string) error {
	set.Action = Stop
	if err := set.Hooks.Run("before.stop", set); err != nil {
		return err
	}
	args = append(args, set.Name)
	if _, err := helpers.RunCmd(append([]interface{}{"stop"}, helpers.ToInterfaceSlice(args)...)...); err != nil {
		return err
	}
	if err := set.Hooks.Run("after.stop", set); err != nil {
		return err
	}
	return nil
}

// Removes the container
// TODO needs to respect scale
func (set *Container) Rm(args []string) error {

	set.Action = Remove
	if err := set.Hooks.Run("before.rm", set); err != nil {
		return err
	}
	args = append(args, set.Name)
	if _, err := helpers.RunCmd(append([]interface{}{"rm"}, helpers.ToInterfaceSlice(args)...)...); err != nil {
		return err
	}
	if err := set.Hooks.Run("after.rm", set); err != nil {
		return err
	}
	return nil
}
<|endoftext|>


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

package model

import (
	"context"

	"go.chromium.org/luci/deploy/api/modelpb"
	"go.chromium.org/luci/gae/service/datastore"
)

// shouldRecordHistory returns true if the new history entry should be recorded.
//
// It is skipped if it is not sufficiently interesting compared to the last
// committed entry.
func shouldRecordHistory(next, prev *modelpb.AssetHistory) bool {
	nextDecision := next.Decision.Decision
	prevDecision := prev.Decision.Decision
	switch {
	case nextDecision != prevDecision:
		return true // changes are always interesting
	case IsActuateDecision(nextDecision):
		return true // active actuations are also always interesting
	case nextDecision == modelpb.ActuationDecision_SKIP_UPTODATE:
		return false // repeating UPTODATE decisions are boring, it is steady state
	case nextDecision == modelpb.ActuationDecision_SKIP_DISABLED:
		return false // repeating DISABLED decisions are also boring
	case nextDecision == modelpb.ActuationDecision_SKIP_LOCKED:
		return !sameLocks(next.Decision.Locks, prev.Decision.Locks)
	case nextDecision == modelpb.ActuationDecision_SKIP_BROKEN:
		return true // errors are always interesting (for retries and alerts)
	default:
		panic("unreachable")
	}
}

func sameLocks(a, b []*modelpb.ActuationLock) bool {
	if len(a) != len(b) {
		return false
	}
	for i := range a {
		if a[i].Id != b[i].Id {
			return false
		}
	}
	return true
}

// commitHistory prepares the history entries for commit and emits TQ tasks.
//
// Must be called inside a transaction. Returns a list of entities to
// transactionally store.
func commitHistory(ctx context.Context, entries []*modelpb.AssetHistory) ([]interface{}, error) {
	// TODO: Emit TQ tasks.
	toPut := make([]interface{}, len(entries))
	for idx, entry := range entries {
		toPut[idx] = &AssetHistory{
			ID:      entry.HistoryId,
			Parent:  datastore.NewKey(ctx, "Asset", entry.AssetId, 0, nil),
			Entry:   entry,
			Created: asTime(entry.Actuation.Created),
		}
	}
	return toPut, nil
}
<commit_msg>[deploy] Don't record uninteresting UPTODATE decisions.<commit_after>// Copyright 2022 The LUCI Authors.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//      http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package model

import (
	"context"

	"go.chromium.org/luci/deploy/api/modelpb"
	"go.chromium.org/luci/gae/service/datastore"
)

// shouldRecordHistory returns true if the new history entry should be recorded.
//
// It is skipped if it is not sufficiently interesting compared to the last
// committed entry.
func shouldRecordHistory(next, prev *modelpb.AssetHistory) bool {
	nextDecision := next.Decision.Decision
	prevDecision := prev.Decision.Decision
	switch {
	case nextDecision == modelpb.ActuationDecision_SKIP_UPTODATE && IsActuateDecision(prevDecision):
		return false // this particular transition is very common and not interesting
	case nextDecision != prevDecision:
		return true // other changes are always interesting
	case IsActuateDecision(nextDecision):
		return true // active actuations are also always interesting
	case nextDecision == modelpb.ActuationDecision_SKIP_UPTODATE:
		return false // repeating UPTODATE decisions are boring, it is steady state
	case nextDecision == modelpb.ActuationDecision_SKIP_DISABLED:
		return false // repeating DISABLED decisions are also boring
	case nextDecision == modelpb.ActuationDecision_SKIP_LOCKED:
		return !sameLocks(next.Decision.Locks, prev.Decision.Locks)
	case nextDecision == modelpb.ActuationDecision_SKIP_BROKEN:
		return true // errors are always interesting (for retries and alerts)
	default:
		panic("unreachable")
	}
}

func sameLocks(a, b []*modelpb.ActuationLock) bool {
	if len(a) != len(b) {
		return false
	}
	for i := range a {
		if a[i].Id != b[i].Id {
			return false
		}
	}
	return true
}

// commitHistory prepares the history entries for commit and emits TQ tasks.
//
// Must be called inside a transaction. Returns a list of entities to
// transactionally store.
func commitHistory(ctx context.Context, entries []*modelpb.AssetHistory) ([]interface{}, error) {
	// TODO: Emit TQ tasks.
	toPut := make([]interface{}, len(entries))
	for idx, entry := range entries {
		toPut[idx] = &AssetHistory{
			ID:      entry.HistoryId,
			Parent:  datastore.NewKey(ctx, "Asset", entry.AssetId, 0, nil),
			Entry:   entry,
			Created: asTime(entry.Actuation.Created),
		}
	}
	return toPut, nil
}
<|endoftext|>