<commit_before>// Copyright (c) 2013 Erik St. Martin, Brian Ketelsen. All rights reserved.
// Use of this source code is governed by The MIT License (MIT) that can be
// found in the LICENSE file.

package server

import (
	"crypto/sha1"
	"github.com/miekg/dns"
	"log"
	"os"
	"strings"
	"sync"
	"time"
)

const origTTL uint32 = 60

var cache *sigCache = newCache()
var inflight *single = new(single)

// TODO(miek): allow RRSIG queries and NSEC queries?

// ParseKeyFile read a DNSSEC keyfile as generated by dnssec-keygen or other
// utilities. It add ".key" for the public key and ".private" for the private key.
func ParseKeyFile(file string) (*dns.DNSKEY, dns.PrivateKey, error) {
	f, e := os.Open(file + ".key")
	if e != nil {
		return nil, nil, e
	}
	k, e := dns.ReadRR(f, file+".key")
	if e != nil {
		return nil, nil, e
	}
	f, e = os.Open(file + ".private")
	if e != nil {
		return nil, nil, e
	}
	p, e := k.(*dns.DNSKEY).ReadPrivateKey(f, file+".private")
	if e != nil {
		return nil, nil, e
	}
	return k.(*dns.DNSKEY), p, nil
}

// nsec creates (if needed) NSEC records that are included in the reply.
func (s *Server) nsec(m *dns.Msg) {
	if m.Rcode == dns.RcodeNameError {
		m.Ns = append(m.Ns, s.newNSEC(m.Question[0].Name))
	}
	if m.Rcode == dns.RcodeSuccess && len(m.Ns) == 1 {
		if _, ok := m.Ns[0].(*dns.SOA); ok {
			m.Ns = append(m.Ns, s.newNSEC(m.Question[0].Name))
		}
	}
}

// sign signs a message m, it takes care of negative or nodata responses as
// well by synthesising NSEC records. It will also cache the signatures, using
// a hash of the signed data as a key.
// We also fake the origin TTL in the signature, because we don't want to
// throw away signatures when services decide to have longer TTL. So we just
// set the origTTL to 60.
func (s *Server) sign(m *dns.Msg, bufsize uint16) {
	now := time.Now().UTC()
	incep := uint32(now.Add(-2 * time.Hour).Unix()) // 2 hours, be sure to catch daylight saving time and such
	expir := uint32(now.Add(7 * 24 * time.Hour).Unix()) // sign for a week

	// TODO(miek): repeating this two times?
	for _, r := range rrSets(m.Answer) {
		key := cache.key(r)
		if s := cache.search(key); s != nil {
			if s.ValidityPeriod(now.Add(-24 * time.Hour)) {
				m.Answer = append(m.Answer, s)
				continue
			}
			cache.remove(key)
		}
		sig, err, shared := inflight.Do(key, func() (*dns.RRSIG, error) {
			sig1 := s.newRRSIG(incep, expir)
			e := sig1.Sign(s.Privkey, r)
			if e != nil {
				log.Printf("Failed to sign: %s\n", e.Error())
			}
			return sig1, e
		})
		if err != nil {
			continue
		}
		if !shared {
			// is it possible to miss this, due the the c.dups > 0 in Do()? TODO(miek)
			cache.insert(key, sig)
		}
		m.Answer = append(m.Answer, dns.Copy(sig).(*dns.RRSIG))
	}
	for _, r := range rrSets(m.Ns) {
		key := cache.key(r)
		if s := cache.search(key); s != nil {
			if s.ValidityPeriod(now.Add(-24 * time.Hour)) {
				m.Ns = append(m.Ns, s)
				continue
			}
			cache.remove(key)
		}
		sig, err, shared := inflight.Do(key, func() (*dns.RRSIG, error) {
			sig1 := s.newRRSIG(incep, expir)
			e := sig1.Sign(s.Privkey, r)
			if e != nil {
				log.Printf("Failed to sign: %s\n", e.Error())
			}
			return sig1, e
		})
		if err != nil {
			continue
		}
		if !shared {
			// is it possible to miss this, due the the c.dups > 0 in Do()? TODO(miek)
			cache.insert(key, sig)
		}
		m.Ns = append(m.Ns, dns.Copy(sig).(*dns.RRSIG))
	}
	// TODO(miek): Forget the additional section for now
	if bufsize >= 512 || bufsize <= 4096 {
		m.Truncated = m.Len() > int(bufsize)
	}
	o := new(dns.OPT)
	o.Hdr.Name = "."
	o.Hdr.Rrtype = dns.TypeOPT
	o.SetDo()
	o.SetUDPSize(4096)
	m.Extra = append(m.Extra, o)
	return
}

func (s *Server) newRRSIG(incep, expir uint32) *dns.RRSIG {
	sig := new(dns.RRSIG)
	sig.Hdr.Rrtype = dns.TypeRRSIG
	sig.Hdr.Ttl = origTTL
	sig.OrigTtl = origTTL
	sig.Algorithm = s.Dnskey.Algorithm
	sig.KeyTag = s.KeyTag
	sig.Inception = incep
	sig.Expiration = expir
	sig.SignerName = s.Dnskey.Hdr.Name
	return sig
}

// newNSEC returns the NSEC record need to denial qname, or gives back a NODATA NSEC.
// TODO(miek): wildcard respones. Do we actually use DNS wildcards in SkyDNS?
func (s *Server) newNSEC(qname string) *dns.NSEC {
	qlabels := dns.SplitDomainName(qname)
	if len(qlabels) < s.domainLabels {
		// TODO(miek): can not happen...?
	}
	// Strip the last s.domainLabels, return up to 4 before
	// that. Four labels is the maximum qname we can handle.
	ls := len(qlabels) - s.domainLabels
	ls4 := ls - 4
	if ls4 < 0 {
		ls4 = 0
	}
	key := qlabels[ls4:ls]
	prev, next := s.registry.GetNSEC(strings.Join(key, "."))
	nsec := &dns.NSEC{Hdr: dns.RR_Header{Name: prev + s.domain + ".", Rrtype: dns.TypeNSEC, Class: dns.ClassINET, Ttl: 60},
		NextDomain: next + s.domain + "."}
	if prev == "" {
		nsec.TypeBitMap = []uint16{dns.TypeA, dns.TypeSOA, dns.TypeNS, dns.TypeAAAA, dns.TypeRRSIG, dns.TypeNSEC, dns.TypeDNSKEY}
	} else {
		nsec.TypeBitMap = []uint16{dns.TypeA, dns.TypeAAAA, dns.TypeSRV, dns.TypeNSEC}
	}
	return nsec
}

type rrset struct {
	qname  string
	qclass uint16
	qtype  uint16
}

func rrSets(rrs []dns.RR) map[rrset][]dns.RR {
	m := make(map[rrset][]dns.RR)
	for _, r := range rrs {
		if s, ok := m[rrset{r.Header().Name, r.Header().Class, r.Header().Rrtype}]; ok {
			s = append(s, r)
		} else {
			s := make([]dns.RR, 1, 3)
			s[0] = r
			m[rrset{r.Header().Name, r.Header().Class, r.Header().Rrtype}] = s
		}
	}
	if len(m) > 0 {
		return m
	}
	return nil
}

type sigCache struct {
	sync.RWMutex
	m map[string]*dns.RRSIG
}

func newCache() *sigCache {
	c := new(sigCache)
	c.m = make(map[string]*dns.RRSIG)
	return c
}

func (c *sigCache) remove(s string) {
	delete(c.m, s)
}

func (c *sigCache) insert(s string, r *dns.RRSIG) {
	c.Lock()
	defer c.Unlock()
	if _, ok := c.m[s]; !ok {
		c.m[s] = r
	}
}

func (c *sigCache) search(s string) *dns.RRSIG {
	c.RLock()
	defer c.RUnlock()
	if s, ok := c.m[s]; ok {
		// we want to return a copy here, because if we didn't the RRSIG
		// could be removed by another goroutine before the packet containing
		// this signature is send out.
		log.Println("DNS Signature retrieved from cache")
		return dns.Copy(s).(*dns.RRSIG)
	}
	return nil
}

// key uses the name, type and rdata, which is serialized and then hashed as the
// key for the lookup
func (c *sigCache) key(rrs []dns.RR) string {
	h := sha1.New()
	i := []byte(rrs[0].Header().Name)
	i = append(i, packUint16(rrs[0].Header().Rrtype)...)
	for _, r := range rrs {
		switch t := r.(type) { // we only do a few type, serialize these manually
		case *dns.SOA:
			i = append(i, packUint32(t.Serial)...)
			// we only fiddle with the serial so store that
		case *dns.SRV:
			i = append(i, packUint16(t.Priority)...)
			i = append(i, packUint16(t.Weight)...)
			i = append(i, packUint16(t.Weight)...)
			i = append(i, []byte(t.Target)...)
		case *dns.A:
			i = append(i, []byte(t.A)...)
		case *dns.AAAA:
			i = append(i, []byte(t.AAAA)...)
		case *dns.DNSKEY:
			// Need nothing more, the rdata stays the same during a run
		case *dns.NSEC:
			// nextname?
		default:
			// not handled
		}
	}
	return string(h.Sum(i))
}

func packUint16(i uint16) []byte { return []byte{byte(i >> 8), byte(i)} }
func packUint32(i uint32) []byte { return []byte{byte(i >> 24), byte(i >> 16), byte(i >> 8), byte(i)} }

// Adapted from singleinflight.go from the original Go Code. Copyright 2013 The Go Authors.
type call struct {
	wg   sync.WaitGroup
	val  *dns.RRSIG
	err  error
	dups int
}

type single struct {
	sync.Mutex
	m map[string]*call
}

func (g *single) Do(key string, fn func() (*dns.RRSIG, error)) (*dns.RRSIG, error, bool) {
	g.Lock()
	if g.m == nil {
		g.m = make(map[string]*call)
	}
	if c, ok := g.m[key]; ok {
		c.dups++
		g.Unlock()
		c.wg.Wait()
		return c.val, c.err, true
	}
	c := new(call)
	c.wg.Add(1)
	g.m[key] = c
	g.Unlock()

	c.val, c.err = fn()
	c.wg.Done()

	g.Lock()
	delete(g.m, key)
	g.Unlock()

	return c.val, c.err, c.dups > 0
}
<commit_msg>Add RRSIG to in the type list<commit_after>// Copyright (c) 2013 Erik St. Martin, Brian Ketelsen. All rights reserved.
// Use of this source code is governed by The MIT License (MIT) that can be
// found in the LICENSE file.

package server

import (
	"crypto/sha1"
	"github.com/miekg/dns"
	"log"
	"os"
	"strings"
	"sync"
	"time"
)

const origTTL uint32 = 60

var cache *sigCache = newCache()
var inflight *single = new(single)

// TODO(miek): allow RRSIG queries and NSEC queries?

// ParseKeyFile read a DNSSEC keyfile as generated by dnssec-keygen or other
// utilities. It add ".key" for the public key and ".private" for the private key.
func ParseKeyFile(file string) (*dns.DNSKEY, dns.PrivateKey, error) {
	f, e := os.Open(file + ".key")
	if e != nil {
		return nil, nil, e
	}
	k, e := dns.ReadRR(f, file+".key")
	if e != nil {
		return nil, nil, e
	}
	f, e = os.Open(file + ".private")
	if e != nil {
		return nil, nil, e
	}
	p, e := k.(*dns.DNSKEY).ReadPrivateKey(f, file+".private")
	if e != nil {
		return nil, nil, e
	}
	return k.(*dns.DNSKEY), p, nil
}

// nsec creates (if needed) NSEC records that are included in the reply.
func (s *Server) nsec(m *dns.Msg) {
	if m.Rcode == dns.RcodeNameError {
		m.Ns = append(m.Ns, s.newNSEC(m.Question[0].Name))
	}
	if m.Rcode == dns.RcodeSuccess && len(m.Ns) == 1 {
		if _, ok := m.Ns[0].(*dns.SOA); ok {
			m.Ns = append(m.Ns, s.newNSEC(m.Question[0].Name))
		}
	}
}

// sign signs a message m, it takes care of negative or nodata responses as
// well by synthesising NSEC records. It will also cache the signatures, using
// a hash of the signed data as a key.
// We also fake the origin TTL in the signature, because we don't want to
// throw away signatures when services decide to have longer TTL. So we just
// set the origTTL to 60.
func (s *Server) sign(m *dns.Msg, bufsize uint16) {
	now := time.Now().UTC()
	incep := uint32(now.Add(-2 * time.Hour).Unix()) // 2 hours, be sure to catch daylight saving time and such
	expir := uint32(now.Add(7 * 24 * time.Hour).Unix()) // sign for a week

	// TODO(miek): repeating this two times?
	for _, r := range rrSets(m.Answer) {
		key := cache.key(r)
		if s := cache.search(key); s != nil {
			if s.ValidityPeriod(now.Add(-24 * time.Hour)) {
				m.Answer = append(m.Answer, s)
				continue
			}
			cache.remove(key)
		}
		sig, err, shared := inflight.Do(key, func() (*dns.RRSIG, error) {
			sig1 := s.newRRSIG(incep, expir)
			e := sig1.Sign(s.Privkey, r)
			if e != nil {
				log.Printf("Failed to sign: %s\n", e.Error())
			}
			return sig1, e
		})
		if err != nil {
			continue
		}
		if !shared {
			// is it possible to miss this, due the the c.dups > 0 in Do()? TODO(miek)
			cache.insert(key, sig)
		}
		m.Answer = append(m.Answer, dns.Copy(sig).(*dns.RRSIG))
	}
	for _, r := range rrSets(m.Ns) {
		key := cache.key(r)
		if s := cache.search(key); s != nil {
			if s.ValidityPeriod(now.Add(-24 * time.Hour)) {
				m.Ns = append(m.Ns, s)
				continue
			}
			cache.remove(key)
		}
		sig, err, shared := inflight.Do(key, func() (*dns.RRSIG, error) {
			sig1 := s.newRRSIG(incep, expir)
			e := sig1.Sign(s.Privkey, r)
			if e != nil {
				log.Printf("Failed to sign: %s\n", e.Error())
			}
			return sig1, e
		})
		if err != nil {
			continue
		}
		if !shared {
			// is it possible to miss this, due the the c.dups > 0 in Do()? TODO(miek)
			cache.insert(key, sig)
		}
		m.Ns = append(m.Ns, dns.Copy(sig).(*dns.RRSIG))
	}
	// TODO(miek): Forget the additional section for now
	if bufsize >= 512 || bufsize <= 4096 {
		m.Truncated = m.Len() > int(bufsize)
	}
	o := new(dns.OPT)
	o.Hdr.Name = "."
	o.Hdr.Rrtype = dns.TypeOPT
	o.SetDo()
	o.SetUDPSize(4096)
	m.Extra = append(m.Extra, o)
	return
}

func (s *Server) newRRSIG(incep, expir uint32) *dns.RRSIG {
	sig := new(dns.RRSIG)
	sig.Hdr.Rrtype = dns.TypeRRSIG
	sig.Hdr.Ttl = origTTL
	sig.OrigTtl = origTTL
	sig.Algorithm = s.Dnskey.Algorithm
	sig.KeyTag = s.KeyTag
	sig.Inception = incep
	sig.Expiration = expir
	sig.SignerName = s.Dnskey.Hdr.Name
	return sig
}

// newNSEC returns the NSEC record need to denial qname, or gives back a NODATA NSEC.
// TODO(miek): wildcard respones. Do we actually use DNS wildcards in SkyDNS?
func (s *Server) newNSEC(qname string) *dns.NSEC {
	qlabels := dns.SplitDomainName(qname)
	if len(qlabels) < s.domainLabels {
		// TODO(miek): can not happen...?
	}
	// Strip the last s.domainLabels, return up to 4 before
	// that. Four labels is the maximum qname we can handle.
	ls := len(qlabels) - s.domainLabels
	ls4 := ls - 4
	if ls4 < 0 {
		ls4 = 0
	}
	key := qlabels[ls4:ls]
	prev, next := s.registry.GetNSEC(strings.Join(key, "."))
	nsec := &dns.NSEC{Hdr: dns.RR_Header{Name: prev + s.domain + ".", Rrtype: dns.TypeNSEC, Class: dns.ClassINET, Ttl: 60},
		NextDomain: next + s.domain + "."}
	if prev == "" {
		nsec.TypeBitMap = []uint16{dns.TypeA, dns.TypeSOA, dns.TypeNS, dns.TypeAAAA, dns.TypeRRSIG, dns.TypeNSEC, dns.TypeDNSKEY}
	} else {
		nsec.TypeBitMap = []uint16{dns.TypeA, dns.TypeAAAA, dns.TypeSRV, dns.TypeRRSIG, dns.TypeNSEC}
	}
	return nsec
}

type rrset struct {
	qname  string
	qclass uint16
	qtype  uint16
}

func rrSets(rrs []dns.RR) map[rrset][]dns.RR {
	m := make(map[rrset][]dns.RR)
	for _, r := range rrs {
		if s, ok := m[rrset{r.Header().Name, r.Header().Class, r.Header().Rrtype}]; ok {
			s = append(s, r)
		} else {
			s := make([]dns.RR, 1, 3)
			s[0] = r
			m[rrset{r.Header().Name, r.Header().Class, r.Header().Rrtype}] = s
		}
	}
	if len(m) > 0 {
		return m
	}
	return nil
}

type sigCache struct {
	sync.RWMutex
	m map[string]*dns.RRSIG
}

func newCache() *sigCache {
	c := new(sigCache)
	c.m = make(map[string]*dns.RRSIG)
	return c
}

func (c *sigCache) remove(s string) {
	delete(c.m, s)
}

func (c *sigCache) insert(s string, r *dns.RRSIG) {
	c.Lock()
	defer c.Unlock()
	if _, ok := c.m[s]; !ok {
		c.m[s] = r
	}
}

func (c *sigCache) search(s string) *dns.RRSIG {
	c.RLock()
	defer c.RUnlock()
	if s, ok := c.m[s]; ok {
		// we want to return a copy here, because if we didn't the RRSIG
		// could be removed by another goroutine before the packet containing
		// this signature is send out.
		log.Println("DNS Signature retrieved from cache")
		return dns.Copy(s).(*dns.RRSIG)
	}
	return nil
}

// key uses the name, type and rdata, which is serialized and then hashed as the
// key for the lookup
func (c *sigCache) key(rrs []dns.RR) string {
	h := sha1.New()
	i := []byte(rrs[0].Header().Name)
	i = append(i, packUint16(rrs[0].Header().Rrtype)...)
	for _, r := range rrs {
		switch t := r.(type) { // we only do a few type, serialize these manually
		case *dns.SOA:
			i = append(i, packUint32(t.Serial)...)
			// we only fiddle with the serial so store that
		case *dns.SRV:
			i = append(i, packUint16(t.Priority)...)
			i = append(i, packUint16(t.Weight)...)
			i = append(i, packUint16(t.Weight)...)
			i = append(i, []byte(t.Target)...)
		case *dns.A:
			i = append(i, []byte(t.A)...)
		case *dns.AAAA:
			i = append(i, []byte(t.AAAA)...)
		case *dns.DNSKEY:
			// Need nothing more, the rdata stays the same during a run
		case *dns.NSEC:
			i = append(i, []byte(t.NextDomain)...)
			for _, v := range t.TypeBitMap {
				i = append(i, packUint16(v)...)
			}
		default:
			log.Printf("DNS Signature for unhandled type %T seen", t)
		}
	}
	return string(h.Sum(i))
}

func packUint16(i uint16) []byte { return []byte{byte(i >> 8), byte(i)} }
func packUint32(i uint32) []byte { return []byte{byte(i >> 24), byte(i >> 16), byte(i >> 8), byte(i)} }

// Adapted from singleinflight.go from the original Go Code. Copyright 2013 The Go Authors.
type call struct {
	wg   sync.WaitGroup
	val  *dns.RRSIG
	err  error
	dups int
}

type single struct {
	sync.Mutex
	m map[string]*call
}

func (g *single) Do(key string, fn func() (*dns.RRSIG, error)) (*dns.RRSIG, error, bool) {
	g.Lock()
	if g.m == nil {
		g.m = make(map[string]*call)
	}
	if c, ok := g.m[key]; ok {
		c.dups++
		g.Unlock()
		c.wg.Wait()
		return c.val, c.err, true
	}
	c := new(call)
	c.wg.Add(1)
	g.m[key] = c
	g.Unlock()

	c.val, c.err = fn()
	c.wg.Done()

	g.Lock()
	delete(g.m, key)
	g.Unlock()

	return c.val, c.err, c.dups > 0
}
<|endoftext|>


<commit_before>// Copyright 2013 The Prometheus 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 main

import (
	"bytes"
	"encoding/binary"
	"fmt"
	"hash/fnv"
	"net"
	"regexp"
	"strconv"
	"strings"

	"github.com/prometheus/client_golang/model"
	"github.com/prometheus/client_golang/prometheus"
	"github.com/prometheus/log"
)

const (
	defaultHelp = "Metric autogenerated by statsd_bridge."
	regErrF     = "A change of configuration created inconsistent metrics for " +
		"%q. You have to restart the statsd_bridge, and you should " +
		"consider the effects on your monitoring setup. Error: %s"
)

var (
	illegalCharsRE = regexp.MustCompile(`[^a-zA-Z0-9_]`)

	hash   = fnv.New64a()
	strBuf bytes.Buffer // Used for hashing.
	intBuf = make([]byte, 8)
)

// hashNameAndLabels returns a hash value of the provided name string and all
// the label names and values in the provided labels map.
//
// Not safe for concurrent use! (Uses a shared buffer and hasher to save on
// allocations.)
func hashNameAndLabels(name string, labels prometheus.Labels) uint64 {
	hash.Reset()
	strBuf.Reset()
	strBuf.WriteString(name)
	hash.Write(strBuf.Bytes())
	binary.BigEndian.PutUint64(intBuf, model.LabelsToSignature(labels))
	hash.Write(intBuf)
	return hash.Sum64()
}

type CounterContainer struct {
	Elements map[uint64]prometheus.Counter
}

func NewCounterContainer() *CounterContainer {
	return &CounterContainer{
		Elements: make(map[uint64]prometheus.Counter),
	}
}

func (c *CounterContainer) Get(metricName string, labels prometheus.Labels) prometheus.Counter {
	hash := hashNameAndLabels(metricName, labels)
	counter, ok := c.Elements[hash]
	if !ok {
		counter = prometheus.NewCounter(prometheus.CounterOpts{
			Name:        metricName,
			Help:        defaultHelp,
			ConstLabels: labels,
		})
		c.Elements[hash] = counter
		if err := prometheus.Register(counter); err != nil {
			log.Fatalf(regErrF, metricName, err)
		}
	}
	return counter
}

type GaugeContainer struct {
	Elements map[uint64]prometheus.Gauge
}

func NewGaugeContainer() *GaugeContainer {
	return &GaugeContainer{
		Elements: make(map[uint64]prometheus.Gauge),
	}
}

func (c *GaugeContainer) Get(metricName string, labels prometheus.Labels) prometheus.Gauge {
	hash := hashNameAndLabels(metricName, labels)
	gauge, ok := c.Elements[hash]
	if !ok {
		gauge = prometheus.NewGauge(prometheus.GaugeOpts{
			Name:        metricName,
			Help:        defaultHelp,
			ConstLabels: labels,
		})
		c.Elements[hash] = gauge
		if err := prometheus.Register(gauge); err != nil {
			log.Fatalf(regErrF, metricName, err)
		}
	}
	return gauge
}

type SummaryContainer struct {
	Elements map[uint64]prometheus.Summary
}

func NewSummaryContainer() *SummaryContainer {
	return &SummaryContainer{
		Elements: make(map[uint64]prometheus.Summary),
	}
}

func (c *SummaryContainer) Get(metricName string, labels prometheus.Labels) prometheus.Summary {
	hash := hashNameAndLabels(metricName, labels)
	summary, ok := c.Elements[hash]
	if !ok {
		summary = prometheus.NewSummary(
			prometheus.SummaryOpts{
				Name:        metricName,
				Help:        defaultHelp,
				ConstLabels: labels,
			})
		c.Elements[hash] = summary
		if err := prometheus.Register(summary); err != nil {
			log.Fatalf(regErrF, metricName, err)
		}
	}
	return summary
}

type Event interface {
	MetricName() string
	Value() float64
}

type CounterEvent struct {
	metricName string
	value      float64
}

func (c *CounterEvent) MetricName() string { return c.metricName }
func (c *CounterEvent) Value() float64     { return c.value }

type GaugeEvent struct {
	metricName string
	value      float64
}

func (g *GaugeEvent) MetricName() string { return g.metricName }
func (g *GaugeEvent) Value() float64     { return g.value }

type TimerEvent struct {
	metricName string
	value      float64
}

func (t *TimerEvent) MetricName() string { return t.metricName }
func (t *TimerEvent) Value() float64     { return t.value }

type Events []Event

type Bridge struct {
	Counters  *CounterContainer
	Gauges    *GaugeContainer
	Summaries *SummaryContainer
	mapper    *metricMapper
}

func escapeMetricName(metricName string) string {
	// If a metric starts with a digit, prepend an underscore.
	if metricName[0] >= '0' && metricName[0] <= '9' {
		metricName = "_" + metricName
	}

	// Replace all illegal metric chars with underscores.
	metricName = illegalCharsRE.ReplaceAllString(metricName, "_")
	return metricName
}

func (b *Bridge) Listen(e <-chan Events) {
	for {
		events := <-e
		for _, event := range events {
			metricName := ""
			prometheusLabels := prometheus.Labels{}

			labels, present := b.mapper.getMapping(event.MetricName())
			if present {
				metricName = labels["name"]
				for label, value := range labels {
					if label != "name" {
						prometheusLabels[label] = value
					}
				}
			} else {
				metricName = escapeMetricName(event.MetricName())
			}

			switch event.(type) {
			case *CounterEvent:
				counter := b.Counters.Get(
					metricName+"_counter",
					prometheusLabels,
				)
				counter.Add(event.Value())

				eventStats.WithLabelValues("counter").Inc()

			case *GaugeEvent:
				gauge := b.Gauges.Get(
					metricName+"_gauge",
					prometheusLabels,
				)
				gauge.Set(event.Value())

				eventStats.WithLabelValues("gauge").Inc()

			case *TimerEvent:
				summary := b.Summaries.Get(
					metricName+"_timer",
					prometheusLabels,
				)
				summary.Observe(event.Value())

				eventStats.WithLabelValues("timer").Inc()

			default:
				log.Println("Unsupported event type")
				eventStats.WithLabelValues("illegal").Inc()
			}
		}
	}
}

func NewBridge(mapper *metricMapper) *Bridge {
	return &Bridge{
		Counters:  NewCounterContainer(),
		Gauges:    NewGaugeContainer(),
		Summaries: NewSummaryContainer(),
		mapper:    mapper,
	}
}

type StatsDListener struct {
	conn *net.UDPConn
}

func buildEvent(statType, metric string, value float64) (Event, error) {
	switch statType {
	case "c":
		return &CounterEvent{
			metricName: metric,
			value:      float64(value),
		}, nil
	case "g":
		return &GaugeEvent{
			metricName: metric,
			value:      float64(value),
		}, nil
	case "ms":
		return &TimerEvent{
			metricName: metric,
			value:      float64(value),
		}, nil
	case "s":
		return nil, fmt.Errorf("No support for StatsD sets")
	default:
		return nil, fmt.Errorf("Bad stat type %s", statType)
	}
}

func (l *StatsDListener) Listen(e chan<- Events) {
	// TODO: evaluate proper size according to MTU
	var buf [512]byte
	for {
		n, _, err := l.conn.ReadFromUDP(buf[0:])
		if err != nil {
			log.Fatal(err)
		}
		l.handlePacket(buf[0:n], e)
	}
}

func (l *StatsDListener) handlePacket(packet []byte, e chan<- Events) {
	lines := strings.Split(string(packet), "\n")
	events := Events{}
	for _, line := range lines {
		if line == "" {
			continue
		}

		elements := strings.Split(line, ":")
		if len(elements) < 2 {
			networkStats.WithLabelValues("malformed_line").Inc()
			log.Println("Bad line from StatsD:", line)
			continue
		}
		metric := elements[0]
		samples := elements[1:]
		for _, sample := range samples {
			components := strings.Split(sample, "|")
			samplingFactor := 1.0
			if len(components) < 2 || len(components) > 3 {
				networkStats.WithLabelValues("malformed_component").Inc()
				log.Println("Bad component on line:", line)
				continue
			}
			valueStr, statType := components[0], components[1]
			value, err := strconv.ParseFloat(valueStr, 64)
			if err != nil {
				log.Printf("Bad value %s on line: %s", valueStr, line)
				networkStats.WithLabelValues("malformed_value").Inc()
				continue
			}

			if len(components) == 3 {
				if statType != "c" {
					log.Println("Illegal sampling factor for non-counter metric on line", line)
					networkStats.WithLabelValues("illegal_sample_factor").Inc()
				}
				samplingStr := components[2]
				if samplingStr[0] != '@' {
					log.Printf("Invalid sampling factor %s on line %s", samplingStr, line)
					networkStats.WithLabelValues("invalid_sample_factor").Inc()
					continue
				}
				samplingFactor, err = strconv.ParseFloat(samplingStr[1:], 64)
				if err != nil {
					log.Printf("Invalid sampling factor %s on line %s", samplingStr, line)
					networkStats.WithLabelValues("invalid_sample_factor").Inc()
					continue
				}
				if samplingFactor == 0 {
					// This should never happen, but avoid division by zero if it does.
					log.Printf("Invalid zero sampling factor %s on line %s, setting to 1", samplingStr, line)
					samplingFactor = 1
				}
				value /= samplingFactor
			}

			event, err := buildEvent(statType, metric, value)
			if err != nil {
				log.Printf("Error building event on line %s: %s", line, err)
				networkStats.WithLabelValues("illegal_event").Inc()
				continue
			}
			events = append(events, event)
			networkStats.WithLabelValues("legal").Inc()
		}
	}
	e <- events
}
<commit_msg>Move from client_golang/model to common/model<commit_after>// Copyright 2013 The Prometheus 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 main

import (
	"bytes"
	"encoding/binary"
	"fmt"
	"hash/fnv"
	"net"
	"regexp"
	"strconv"
	"strings"

	"github.com/prometheus/client_golang/prometheus"
	"github.com/prometheus/common/model"
	"github.com/prometheus/log"
)

const (
	defaultHelp = "Metric autogenerated by statsd_bridge."
	regErrF     = "A change of configuration created inconsistent metrics for " +
		"%q. You have to restart the statsd_bridge, and you should " +
		"consider the effects on your monitoring setup. Error: %s"
)

var (
	illegalCharsRE = regexp.MustCompile(`[^a-zA-Z0-9_]`)

	hash   = fnv.New64a()
	strBuf bytes.Buffer // Used for hashing.
	intBuf = make([]byte, 8)
)

// hashNameAndLabels returns a hash value of the provided name string and all
// the label names and values in the provided labels map.
//
// Not safe for concurrent use! (Uses a shared buffer and hasher to save on
// allocations.)
func hashNameAndLabels(name string, labels prometheus.Labels) uint64 {
	hash.Reset()
	strBuf.Reset()
	strBuf.WriteString(name)
	hash.Write(strBuf.Bytes())
	binary.BigEndian.PutUint64(intBuf, model.LabelsToSignature(labels))
	hash.Write(intBuf)
	return hash.Sum64()
}

type CounterContainer struct {
	Elements map[uint64]prometheus.Counter
}

func NewCounterContainer() *CounterContainer {
	return &CounterContainer{
		Elements: make(map[uint64]prometheus.Counter),
	}
}

func (c *CounterContainer) Get(metricName string, labels prometheus.Labels) prometheus.Counter {
	hash := hashNameAndLabels(metricName, labels)
	counter, ok := c.Elements[hash]
	if !ok {
		counter = prometheus.NewCounter(prometheus.CounterOpts{
			Name:        metricName,
			Help:        defaultHelp,
			ConstLabels: labels,
		})
		c.Elements[hash] = counter
		if err := prometheus.Register(counter); err != nil {
			log.Fatalf(regErrF, metricName, err)
		}
	}
	return counter
}

type GaugeContainer struct {
	Elements map[uint64]prometheus.Gauge
}

func NewGaugeContainer() *GaugeContainer {
	return &GaugeContainer{
		Elements: make(map[uint64]prometheus.Gauge),
	}
}

func (c *GaugeContainer) Get(metricName string, labels prometheus.Labels) prometheus.Gauge {
	hash := hashNameAndLabels(metricName, labels)
	gauge, ok := c.Elements[hash]
	if !ok {
		gauge = prometheus.NewGauge(prometheus.GaugeOpts{
			Name:        metricName,
			Help:        defaultHelp,
			ConstLabels: labels,
		})
		c.Elements[hash] = gauge
		if err := prometheus.Register(gauge); err != nil {
			log.Fatalf(regErrF, metricName, err)
		}
	}
	return gauge
}

type SummaryContainer struct {
	Elements map[uint64]prometheus.Summary
}

func NewSummaryContainer() *SummaryContainer {
	return &SummaryContainer{
		Elements: make(map[uint64]prometheus.Summary),
	}
}

func (c *SummaryContainer) Get(metricName string, labels prometheus.Labels) prometheus.Summary {
	hash := hashNameAndLabels(metricName, labels)
	summary, ok := c.Elements[hash]
	if !ok {
		summary = prometheus.NewSummary(
			prometheus.SummaryOpts{
				Name:        metricName,
				Help:        defaultHelp,
				ConstLabels: labels,
			})
		c.Elements[hash] = summary
		if err := prometheus.Register(summary); err != nil {
			log.Fatalf(regErrF, metricName, err)
		}
	}
	return summary
}

type Event interface {
	MetricName() string
	Value() float64
}

type CounterEvent struct {
	metricName string
	value      float64
}

func (c *CounterEvent) MetricName() string { return c.metricName }
func (c *CounterEvent) Value() float64     { return c.value }

type GaugeEvent struct {
	metricName string
	value      float64
}

func (g *GaugeEvent) MetricName() string { return g.metricName }
func (g *GaugeEvent) Value() float64     { return g.value }

type TimerEvent struct {
	metricName string
	value      float64
}

func (t *TimerEvent) MetricName() string { return t.metricName }
func (t *TimerEvent) Value() float64     { return t.value }

type Events []Event

type Bridge struct {
	Counters  *CounterContainer
	Gauges    *GaugeContainer
	Summaries *SummaryContainer
	mapper    *metricMapper
}

func escapeMetricName(metricName string) string {
	// If a metric starts with a digit, prepend an underscore.
	if metricName[0] >= '0' && metricName[0] <= '9' {
		metricName = "_" + metricName
	}

	// Replace all illegal metric chars with underscores.
	metricName = illegalCharsRE.ReplaceAllString(metricName, "_")
	return metricName
}

func (b *Bridge) Listen(e <-chan Events) {
	for {
		events := <-e
		for _, event := range events {
			metricName := ""
			prometheusLabels := prometheus.Labels{}

			labels, present := b.mapper.getMapping(event.MetricName())
			if present {
				metricName = labels["name"]
				for label, value := range labels {
					if label != "name" {
						prometheusLabels[label] = value
					}
				}
			} else {
				metricName = escapeMetricName(event.MetricName())
			}

			switch event.(type) {
			case *CounterEvent:
				counter := b.Counters.Get(
					metricName+"_counter",
					prometheusLabels,
				)
				counter.Add(event.Value())

				eventStats.WithLabelValues("counter").Inc()

			case *GaugeEvent:
				gauge := b.Gauges.Get(
					metricName+"_gauge",
					prometheusLabels,
				)
				gauge.Set(event.Value())

				eventStats.WithLabelValues("gauge").Inc()

			case *TimerEvent:
				summary := b.Summaries.Get(
					metricName+"_timer",
					prometheusLabels,
				)
				summary.Observe(event.Value())

				eventStats.WithLabelValues("timer").Inc()

			default:
				log.Println("Unsupported event type")
				eventStats.WithLabelValues("illegal").Inc()
			}
		}
	}
}

func NewBridge(mapper *metricMapper) *Bridge {
	return &Bridge{
		Counters:  NewCounterContainer(),
		Gauges:    NewGaugeContainer(),
		Summaries: NewSummaryContainer(),
		mapper:    mapper,
	}
}

type StatsDListener struct {
	conn *net.UDPConn
}

func buildEvent(statType, metric string, value float64) (Event, error) {
	switch statType {
	case "c":
		return &CounterEvent{
			metricName: metric,
			value:      float64(value),
		}, nil
	case "g":
		return &GaugeEvent{
			metricName: metric,
			value:      float64(value),
		}, nil
	case "ms":
		return &TimerEvent{
			metricName: metric,
			value:      float64(value),
		}, nil
	case "s":
		return nil, fmt.Errorf("No support for StatsD sets")
	default:
		return nil, fmt.Errorf("Bad stat type %s", statType)
	}
}

func (l *StatsDListener) Listen(e chan<- Events) {
	// TODO: evaluate proper size according to MTU
	var buf [512]byte
	for {
		n, _, err := l.conn.ReadFromUDP(buf[0:])
		if err != nil {
			log.Fatal(err)
		}
		l.handlePacket(buf[0:n], e)
	}
}

func (l *StatsDListener) handlePacket(packet []byte, e chan<- Events) {
	lines := strings.Split(string(packet), "\n")
	events := Events{}
	for _, line := range lines {
		if line == "" {
			continue
		}

		elements := strings.Split(line, ":")
		if len(elements) < 2 {
			networkStats.WithLabelValues("malformed_line").Inc()
			log.Println("Bad line from StatsD:", line)
			continue
		}
		metric := elements[0]
		samples := elements[1:]
		for _, sample := range samples {
			components := strings.Split(sample, "|")
			samplingFactor := 1.0
			if len(components) < 2 || len(components) > 3 {
				networkStats.WithLabelValues("malformed_component").Inc()
				log.Println("Bad component on line:", line)
				continue
			}
			valueStr, statType := components[0], components[1]
			value, err := strconv.ParseFloat(valueStr, 64)
			if err != nil {
				log.Printf("Bad value %s on line: %s", valueStr, line)
				networkStats.WithLabelValues("malformed_value").Inc()
				continue
			}

			if len(components) == 3 {
				if statType != "c" {
					log.Println("Illegal sampling factor for non-counter metric on line", line)
					networkStats.WithLabelValues("illegal_sample_factor").Inc()
				}
				samplingStr := components[2]
				if samplingStr[0] != '@' {
					log.Printf("Invalid sampling factor %s on line %s", samplingStr, line)
					networkStats.WithLabelValues("invalid_sample_factor").Inc()
					continue
				}
				samplingFactor, err = strconv.ParseFloat(samplingStr[1:], 64)
				if err != nil {
					log.Printf("Invalid sampling factor %s on line %s", samplingStr, line)
					networkStats.WithLabelValues("invalid_sample_factor").Inc()
					continue
				}
				if samplingFactor == 0 {
					// This should never happen, but avoid division by zero if it does.
					log.Printf("Invalid zero sampling factor %s on line %s, setting to 1", samplingStr, line)
					samplingFactor = 1
				}
				value /= samplingFactor
			}

			event, err := buildEvent(statType, metric, value)
			if err != nil {
				log.Printf("Error building event on line %s: %s", line, err)
				networkStats.WithLabelValues("illegal_event").Inc()
				continue
			}
			events = append(events, event)
			networkStats.WithLabelValues("legal").Inc()
		}
	}
	e <- events
}
<|endoftext|>


<commit_before>package peco

import (
	"bufio"
	"io"
	"sync"
	"time"

	"github.com/lestrrat/go-pdebug"
	"github.com/peco/peco/internal/util"
	"github.com/peco/peco/pipeline"
	"github.com/pkg/errors"
	"golang.org/x/net/context"
)

func NewFilteredBuffer(src Buffer, page, perPage int) *FilteredBuffer {
	fb := FilteredBuffer{
		src: src,
	}

	s := perPage * (page - 1)
	if s > src.Size() {
		return &fb
	}

	selection := make([]int, 0, src.Size())
	e := s + perPage
	if e >= src.Size() {
		e = src.Size()
	}

	for i := s; i < e; i++ {
		selection = append(selection, i)
	}
	fb.selection = selection

	return &fb
}

func (flb *FilteredBuffer) Append(l Line) (Line, error) {
	return l, nil
}

// LineAt returns the line at index `i`. Note that the i-th element
// in this filtered buffer may actually correspond to a totally
// different line number in the source buffer.
func (flb FilteredBuffer) LineAt(i int) (Line, error) {
	if i >= len(flb.selection) {
		return nil, errors.Errorf("specified index %d is out of range", len(flb.selection))
	}
	return flb.src.LineAt(flb.selection[i])
}

// Size returns the number of lines in the buffer
func (flb FilteredBuffer) Size() int {
	return len(flb.selection)
}

func NewMemoryBuffer() *MemoryBuffer {
	return &MemoryBuffer{
		done: make(chan struct{}),
	}
}

func (mb *MemoryBuffer) Append(l Line) {
	mb.mutex.Lock()
	defer mb.mutex.Unlock()

	mb.lines = append(mb.lines, l)
}

func (mb *MemoryBuffer) Size() int {
	mb.mutex.RLock()
	defer mb.mutex.RUnlock()
	return len(mb.lines)
}

func (mb *MemoryBuffer) Reset() {
	mb.mutex.Lock()
	defer mb.mutex.Unlock()
	mb.done = make(chan struct{})
	mb.lines = []Line(nil)
}

func (mb *MemoryBuffer) Done() <-chan struct{} {
	mb.mutex.RLock()
	defer mb.mutex.RUnlock()
	return mb.done
}

func (mb *MemoryBuffer) Accept(ctx context.Context, p pipeline.Producer) {
	mb.mutex.Lock()
	defer mb.mutex.Unlock()
	if pdebug.Enabled {
		g := pdebug.Marker("MemoryBuffer.Accept")
		defer g.End()
	}
	defer close(mb.done)

	for {
		select {
		case <-ctx.Done():
			if pdebug.Enabled {
				pdebug.Printf("MemoryBuffer received context done")
			}
			return
		case v := <-p.OutCh():
			switch v.(type) {
			case error:
				if pipeline.IsEndMark(v.(error)) {
					if pdebug.Enabled {
						pdebug.Printf("MemoryBuffer received end mark (read %d lines)", len(mb.lines))
					}
					return
				}
			case Line:
				mb.lines = append(mb.lines, v.(Line))
			}
		}
	}
}

func (mb *MemoryBuffer) LineAt(n int) (Line, error) {
	mb.mutex.RLock()
	defer mb.mutex.RUnlock()

	if s := len(mb.lines); s <= 0 || n >= s {
		return nil, errors.New("empty buffer")
	}

	return mb.lines[n], nil
}

// Creates a new Source. Does not start processing the input until you
// call Setup()
func NewSource(in io.Reader, idgen lineIDGenerator, enableSep bool) *Source {
	return &Source{
		in:            in, // Note that this may be closed, so do not rely on it
		enableSep:     enableSep,
		done:          make(chan struct{}),
		idgen:         idgen,
		ready:         make(chan struct{}),
		start:         make(chan struct{}, 1),
		setupOnce:     sync.Once{},
		OutputChannel: pipeline.OutputChannel(make(chan interface{})),
	}
}

// Setup reads from the input os.File.
func (s *Source) Setup(state *Peco) {
	s.setupOnce.Do(func() {
		done := make(chan struct{})
		refresh := make(chan struct{}, 1)
		defer close(done)
		defer close(refresh)

		draw := func(state *Peco) {
			// Not a great thing to do, allowing nil to be passed
			// as state, but for testing I couldn't come up with anything
			// better for the moment
			if state != nil {
				state.Hub().SendDraw(false)
			}
		}

		go func() {
			ticker := time.NewTicker(100 * time.Millisecond)
			defer ticker.Stop()

			for {
				select {
				case <-done:
					draw(state)
					return
				case <-ticker.C:
					draw(state)
				}
			}
		}()

		// This sync.Once var is used to receive the notification
		// that there was at least 1 line read from the source
		var notify sync.Once
		notifycb := func() {
			// close the ready channel so others can be notified
			// that there's at least 1 line in the buffer
			close(s.ready)
		}
		scanner := bufio.NewScanner(s.in)
		defer func() {
			if util.IsTty(s.in) {
				return
			}
			if closer, ok := s.in.(io.Closer); ok {
				closer.Close()
			}
		}()

		readCount := 0
		for scanner.Scan() {
			txt := scanner.Text()
			readCount++
			s.Append(NewRawLine(s.idgen.next(), txt, s.enableSep))
			notify.Do(notifycb)
		}

		// XXX Just in case scanner.Scan() did not return a single line...
		// Note: this will be a no-op if notify.Do has been called before
		notify.Do(notifycb)
		// And also, close the done channel so we can tell the consumers
		// we have finished reading everything
		close(s.done)

		if pdebug.Enabled {
			pdebug.Printf("Read all %d lines from source", readCount)
		}
	})
}

// Start starts
func (s *Source) Start(ctx context.Context) {
	// I should be the only one running this method until I bail out
	if pdebug.Enabled {
		g := pdebug.Marker("Source.Start")
		defer g.End()
		defer pdebug.Printf("Source sent %d lines", len(s.lines))
	}
	s.start<-struct{}{}
	defer func() { <-s.start }()

	defer s.OutputChannel.SendEndMark("end of input")
	defer close(s.done)

	for _, l := range s.lines {
		select {
		case <-ctx.Done():
			if pdebug.Enabled {
				pdebug.Printf("Source: context.Done detected")
			}
			return
		case s.OutputChannel <- l:
			// no op
		}
	}
}

// Reset resets the state of the source object so that it
// is ready to feed the filters
func (s *Source) Reset() {
	s.done = make(chan struct{})
	s.OutputChannel = pipeline.OutputChannel(make(chan interface{}))
}

// Ready returns the "input ready" channel. It will be closed as soon as
// the first line of input is processed via Setup()
func (s *Source) Ready() <-chan struct{} {
	return s.ready
}

// Done returns the "read all lines" channel. It will be closed as soon as
// the all input has been read
func (s *Source) Done() <-chan struct{} {
	return s.done
}
<commit_msg>Localize locking so that the one can Accept while doing other things<commit_after>package peco

import (
	"bufio"
	"io"
	"sync"
	"time"

	"github.com/lestrrat/go-pdebug"
	"github.com/peco/peco/internal/util"
	"github.com/peco/peco/pipeline"
	"github.com/pkg/errors"
	"golang.org/x/net/context"
)

func NewFilteredBuffer(src Buffer, page, perPage int) *FilteredBuffer {
	fb := FilteredBuffer{
		src: src,
	}

	s := perPage * (page - 1)
	if s > src.Size() {
		return &fb
	}

	selection := make([]int, 0, src.Size())
	e := s + perPage
	if e >= src.Size() {
		e = src.Size()
	}

	for i := s; i < e; i++ {
		selection = append(selection, i)
	}
	fb.selection = selection

	return &fb
}

func (flb *FilteredBuffer) Append(l Line) (Line, error) {
	return l, nil
}

// LineAt returns the line at index `i`. Note that the i-th element
// in this filtered buffer may actually correspond to a totally
// different line number in the source buffer.
func (flb FilteredBuffer) LineAt(i int) (Line, error) {
	if i >= len(flb.selection) {
		return nil, errors.Errorf("specified index %d is out of range", len(flb.selection))
	}
	return flb.src.LineAt(flb.selection[i])
}

// Size returns the number of lines in the buffer
func (flb FilteredBuffer) Size() int {
	return len(flb.selection)
}

func NewMemoryBuffer() *MemoryBuffer {
	return &MemoryBuffer{
		done: make(chan struct{}),
	}
}

func (mb *MemoryBuffer) Append(l Line) {
	mb.mutex.Lock()
	defer mb.mutex.Unlock()

	mb.lines = append(mb.lines, l)
}

func (mb *MemoryBuffer) Size() int {
	mb.mutex.RLock()
	defer mb.mutex.RUnlock()
	return len(mb.lines)
}

func (mb *MemoryBuffer) Reset() {
	mb.mutex.Lock()
	defer mb.mutex.Unlock()
	mb.done = make(chan struct{})
	mb.lines = []Line(nil)
}

func (mb *MemoryBuffer) Done() <-chan struct{} {
	mb.mutex.RLock()
	defer mb.mutex.RUnlock()
	return mb.done
}

func (mb *MemoryBuffer) Accept(ctx context.Context, p pipeline.Producer) {
	if pdebug.Enabled {
		g := pdebug.Marker("MemoryBuffer.Accept")
		defer g.End()
	}
	defer func() {
		mb.mutex.Lock()
		close(mb.done)
		mb.mutex.Unlock()
	}()

	for {
		select {
		case <-ctx.Done():
			if pdebug.Enabled {
				pdebug.Printf("MemoryBuffer received context done")
			}
			return
		case v := <-p.OutCh():
			switch v.(type) {
			case error:
				if pipeline.IsEndMark(v.(error)) {
					if pdebug.Enabled {
						pdebug.Printf("MemoryBuffer received end mark (read %d lines)", len(mb.lines))
					}
					return
				}
			case Line:
				mb.mutex.Lock()
				mb.lines = append(mb.lines, v.(Line))
				mb.mutex.Unlock()
			}
		}
	}
}

func (mb *MemoryBuffer) LineAt(n int) (Line, error) {
	mb.mutex.RLock()
	defer mb.mutex.RUnlock()

	if s := len(mb.lines); s <= 0 || n >= s {
		return nil, errors.New("empty buffer")
	}

	return mb.lines[n], nil
}

// Creates a new Source. Does not start processing the input until you
// call Setup()
func NewSource(in io.Reader, idgen lineIDGenerator, enableSep bool) *Source {
	return &Source{
		in:            in, // Note that this may be closed, so do not rely on it
		enableSep:     enableSep,
		done:          make(chan struct{}),
		idgen:         idgen,
		ready:         make(chan struct{}),
		start:         make(chan struct{}, 1),
		setupOnce:     sync.Once{},
		OutputChannel: pipeline.OutputChannel(make(chan interface{})),
	}
}

// Setup reads from the input os.File.
func (s *Source) Setup(state *Peco) {
	s.setupOnce.Do(func() {
		done := make(chan struct{})
		refresh := make(chan struct{}, 1)
		defer close(done)
		defer close(refresh)

		draw := func(state *Peco) {
			// Not a great thing to do, allowing nil to be passed
			// as state, but for testing I couldn't come up with anything
			// better for the moment
			if state != nil {
				state.Hub().SendDraw(false)
			}
		}

		go func() {
			ticker := time.NewTicker(100 * time.Millisecond)
			defer ticker.Stop()

			for {
				select {
				case <-done:
					draw(state)
					return
				case <-ticker.C:
					draw(state)
				}
			}
		}()

		// This sync.Once var is used to receive the notification
		// that there was at least 1 line read from the source
		var notify sync.Once
		notifycb := func() {
			// close the ready channel so others can be notified
			// that there's at least 1 line in the buffer
			close(s.ready)
		}
		scanner := bufio.NewScanner(s.in)
		defer func() {
			if util.IsTty(s.in) {
				return
			}
			if closer, ok := s.in.(io.Closer); ok {
				closer.Close()
			}
		}()

		readCount := 0
		for scanner.Scan() {
			txt := scanner.Text()
			readCount++
			s.Append(NewRawLine(s.idgen.next(), txt, s.enableSep))
			notify.Do(notifycb)
		}

		// XXX Just in case scanner.Scan() did not return a single line...
		// Note: this will be a no-op if notify.Do has been called before
		notify.Do(notifycb)
		// And also, close the done channel so we can tell the consumers
		// we have finished reading everything
		close(s.done)

		if pdebug.Enabled {
			pdebug.Printf("Read all %d lines from source", readCount)
		}
	})
}

// Start starts
func (s *Source) Start(ctx context.Context) {
	// I should be the only one running this method until I bail out
	if pdebug.Enabled {
		g := pdebug.Marker("Source.Start")
		defer g.End()
		defer pdebug.Printf("Source sent %d lines", len(s.lines))
	}
	s.start<-struct{}{}
	defer func() { <-s.start }()

	defer s.OutputChannel.SendEndMark("end of input")
	defer close(s.done)

	for _, l := range s.lines {
		select {
		case <-ctx.Done():
			if pdebug.Enabled {
				pdebug.Printf("Source: context.Done detected")
			}
			return
		case s.OutputChannel <- l:
			// no op
		}
	}
}

// Reset resets the state of the source object so that it
// is ready to feed the filters
func (s *Source) Reset() {
	s.done = make(chan struct{})
	s.OutputChannel = pipeline.OutputChannel(make(chan interface{}))
}

// Ready returns the "input ready" channel. It will be closed as soon as
// the first line of input is processed via Setup()
func (s *Source) Ready() <-chan struct{} {
	return s.ready
}

// Done returns the "read all lines" channel. It will be closed as soon as
// the all input has been read
func (s *Source) Done() <-chan struct{} {
	return s.done
}
<|endoftext|>


<commit_before>/*
 Copyright 2015 Stack Inc.
 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 cli

import (
	"log"
	"os"
)

func buildBsp(t *Target, pm *PkgMgr, incls *[]string,
	libs *[]string) (string, error) {

	if t.Bsp == "" {
		return "", NewStackError("Expected a BSP")
	}

	bspPackage, err := pm.ResolvePkgName(t.Bsp)
	if err != nil {
		return "", NewStackError("No BSP package for " + t.Bsp + " exists")
	}

	if err = pm.Build(t, t.Bsp, nil, libs); err != nil {
		return "", err
	}

	*incls = append(*incls, bspPackage.Includes...)

	// A BSP doesn't have to contain source; don't fail if no library was
	// built.
	if lib := pm.GetPackageLib(t, bspPackage); NodeExist(lib) {
		*libs = append(*libs, lib)
	}

	var linkerScript string
	if bspPackage.LinkerScript != "" {
		linkerScript = bspPackage.BasePath + "/" + bspPackage.LinkerScript
	} else {
		linkerScript = ""
	}

	return linkerScript, nil
}

// Creates the set of compiler flags that should be specified when building a
// particular target-package pair.
func CreateCFlags(c *Compiler, t *Target, p *Package) string {
	cflags := c.Cflags + " " + p.Cflags + " " + t.Cflags

	// The 'test' identity causes the TEST symbol to be defined.  This allows
	// package code to behave differently in test builds.
	if t.HasIdentity("test") {
		cflags += " -DTEST"
	}

	return cflags
}

func PkgIncludeDirs(pkg *Package, t *Target, srcDir string) []string {
	incls := pkg.Includes
	incls = append(incls, srcDir)
	incls = append(incls, srcDir+"/arch/"+t.Arch)

	if t.HasIdentity("test") {
		testSrcDir := srcDir + "/test"
		incls = append(incls, testSrcDir)
		incls = append(incls, testSrcDir+"/arch/"+t.Arch)
	}

	return incls
}

// Recursively compiles all the .c and .s files in the specified directory.
// Architecture-specific files are also compiled.
func BuildDir(srcDir string, c *Compiler, t *Target, ignDirs []string) error {

	var err error

	log.Printf("[DEBUG] compiling src in base directory: %s", srcDir)

	// First change into the package src directory, and build all the objects
	// there
	os.Chdir(srcDir)

	// Ignore architecture-specific source files for now.
	ignDirs = append(ignDirs, "arch")

	if err = c.RecursiveCompile("*.c", 0, ignDirs); err != nil {
		return err
	}

	log.Printf("[DEBUG] compiling architecture specific src packages")

	archDir := srcDir + "/arch/" + t.Arch + "/"
	if NodeExist(archDir) {
		if err := os.Chdir(archDir); err != nil {
			return NewStackError(err.Error())
		}
		if err := c.RecursiveCompile("*.c", 0, nil); err != nil {
			return err
		}

		// compile assembly sources in recursive compile as well
		if err = c.RecursiveCompile("*.s", 1, nil); err != nil {
			return err
		}
	}

	return nil
}
<commit_msg>Don't try to build source files in the obj and bin directories.<commit_after>/*
 Copyright 2015 Stack Inc.
 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 cli

import (
	"log"
	"os"
)

func buildBsp(t *Target, pm *PkgMgr, incls *[]string,
	libs *[]string) (string, error) {

	if t.Bsp == "" {
		return "", NewStackError("Expected a BSP")
	}

	bspPackage, err := pm.ResolvePkgName(t.Bsp)
	if err != nil {
		return "", NewStackError("No BSP package for " + t.Bsp + " exists")
	}

	if err = pm.Build(t, t.Bsp, nil, libs); err != nil {
		return "", err
	}

	*incls = append(*incls, bspPackage.Includes...)

	// A BSP doesn't have to contain source; don't fail if no library was
	// built.
	if lib := pm.GetPackageLib(t, bspPackage); NodeExist(lib) {
		*libs = append(*libs, lib)
	}

	var linkerScript string
	if bspPackage.LinkerScript != "" {
		linkerScript = bspPackage.BasePath + "/" + bspPackage.LinkerScript
	} else {
		linkerScript = ""
	}

	return linkerScript, nil
}

// Creates the set of compiler flags that should be specified when building a
// particular target-package pair.
func CreateCFlags(c *Compiler, t *Target, p *Package) string {
	cflags := c.Cflags + " " + p.Cflags + " " + t.Cflags

	// The 'test' identity causes the TEST symbol to be defined.  This allows
	// package code to behave differently in test builds.
	if t.HasIdentity("test") {
		cflags += " -DTEST"
	}

	return cflags
}

func PkgIncludeDirs(pkg *Package, t *Target, srcDir string) []string {
	incls := pkg.Includes
	incls = append(incls, srcDir)
	incls = append(incls, srcDir+"/arch/"+t.Arch)

	if t.HasIdentity("test") {
		testSrcDir := srcDir + "/test"
		incls = append(incls, testSrcDir)
		incls = append(incls, testSrcDir+"/arch/"+t.Arch)
	}

	return incls
}

// Recursively compiles all the .c and .s files in the specified directory.
// Architecture-specific files are also compiled.
func BuildDir(srcDir string, c *Compiler, t *Target, ignDirs []string) error {

	var err error

	log.Printf("[DEBUG] compiling src in base directory: %s", srcDir)

	// First change into the package src directory, and build all the objects
	// there
	os.Chdir(srcDir)

	// Don't recurse into destination directories.
	ignDirs = append(ignDirs, "obj")
	ignDirs = append(ignDirs, "bin")

	// Ignore architecture-specific source files for now.  Use a temporary
	// string array here so that the "arch" director is not ignored in the
	// subsequent architecture-specific compile phase.
	baseIgnDirs := append(ignDirs, "arch")

	if err = c.RecursiveCompile("*.c", 0, baseIgnDirs); err != nil {
		return err
	}

	log.Printf("[DEBUG] compiling architecture specific src packages")

	archDir := srcDir + "/arch/" + t.Arch + "/"
	if NodeExist(archDir) {
		if err := os.Chdir(archDir); err != nil {
			return NewStackError(err.Error())
		}
		if err := c.RecursiveCompile("*.c", 0, ignDirs); err != nil {
			return err
		}

		// compile assembly sources in recursive compile as well
		if err = c.RecursiveCompile("*.s", 1, ignDirs); err != nil {
			return err
		}
	}

	return nil
}
<|endoftext|>


<commit_before>package main

import (
	"bufio"
	"encoding/json"
	"fmt"
	"io"
	"net"
	"os"
	"strconv"
	"time"

	"golang.org/x/net/context"

	"github.com/Sirupsen/logrus"
	"github.com/docker/docker/api/types"
	"github.com/docker/docker/api/types/container"
	"github.com/docker/docker/api/types/network"
	dcli "github.com/docker/docker/client"
	"github.com/garyburd/redigo/redis"
)

// DockerClient for running code
var DockerClient *dcli.Client

// DockerAPIVersion is the API version connect to docker
const DockerAPIVersion = "1.24"

// Runner runs the code
type Runner struct {
	Lang          string `json:"lang"`
	Source        string `json:"source"`
	Version       string `json:"version"`
	Timeout       int    `json:"timeout"` // How long is the code going to run
	closeNotifier <-chan bool
	logger        *logrus.Logger
	containerID   string
}

// Runnerthrottle Limit the max throttle for runner
var Runnerthrottle chan struct{}

// FetchCode get the code from Redis Server according to the UUID
func FetchCode(uuid string, redisConn redis.Conn) (r *Runner, err error) {
	value, err := redis.Bytes(redisConn.Do("GET", uuid+"#run"))

	if err != nil {
		return
	}

	r = &Runner{}
	err = json.Unmarshal(value, r)
	return
}

// Run the code in the container
func (r *Runner) Run(output messages, conn redis.Conn, uuid string) {
	Runnerthrottle <- struct{}{}
	defer func() { <-Runnerthrottle }()

	err := r.createContainer(uuid)
	if err != nil {
		r.logger.Errorf("Container %s cannot be created - %v", uuid, err)
		return
	}

	stdoutReader, stdoutWriter := io.Pipe()
	stdinReader, stdinWriter := io.Pipe()

	if err != nil {
		r.logger.Error(err)
	}

	defer stdinWriter.Close()
	defer stdoutWriter.Close()

	go pipeStdin(conn, uuid, stdinWriter, r.logger)
	go pipeStdout(stdoutReader, output, r.logger)

	// Start running the container
	err = r.startContainer()
	if err != nil {
		r.logger.Errorf("Container %s cannot be started - %v", uuid, err)
		return
	}
	defer DockerClient.ContainerRemove(context.Background(), r.containerID, types.ContainerRemoveOptions{
		Force: true,
	})

	successChan := make(chan struct{})
	errorChan := make(chan error)

	go r.waitContainer(successChan, errorChan)

	go func(r *Runner, uuid string) {
		err = r.attachContainer(stdoutWriter, stdinReader)
		if err != nil {
			r.logger.Errorf("Container %s cannot be attached - %v", uuid, err)
		}
	}(r, uuid)

	select {
	case <-r.closeNotifier:
		DockerClient.ContainerStop(context.Background(), r.containerID, nil)
		r.logger.Infof("Container %s is stopped since the streamming has been halted", uuid)
	case <-successChan:
		r.logger.Infof("Container %s is executed successfully", uuid)
	case err := <-errorChan:
		r.logger.Errorf("Container %s failed caused by - %v", uuid, err)
	case <-time.After(time.Duration(r.Timeout) * time.Second):
		msg := fmt.Sprintf("Container %s is terminated caused by 15 sec timeout\n", uuid)
		r.logger.Error(msg)
		output <- msg
	}
}

func pipeStdin(conn redis.Conn, uuid string, stdin *io.PipeWriter, logger *logrus.Logger) {
	psc := redis.PubSubConn{Conn: conn}
	psc.Subscribe(uuid + "#stdin")

	defer func() {
		psc.Unsubscribe(uuid + "#stdin")
		psc.Close()
		conn.Close()
	}()

StdinSubscriptionLoop:
	for {
		switch n := psc.Receive().(type) {
		case redis.Message:
			stdinData := strconv.QuoteToASCII(string(n.Data))
			logger.Infof("Message: %s %s", n.Channel, stdinData)
			stdin.Write(n.Data)
		case error:
			break StdinSubscriptionLoop
		}
	}
	logger.Info("Stdin subscription closed")
}

func pipeStdout(stdout *io.PipeReader, output messages, logger *logrus.Logger) {
	buffer := make([]byte, 512)
	for {
		n, err := stdout.Read(buffer)
		if err != nil {
			if err != io.EOF {
				stdout.Close()
				logger.Error(err)
			}

			close(output)
			break
		}

		data := buffer[0:n]
		output <- string(data)

		// Clear the buffer
		for i := 0; i < n; i++ {
			buffer[i] = 0
		}
	}
}

// NewDockerClient create a new docker client
func NewDockerClient() (*dcli.Client, error) {
	os.Setenv("DOCKER_API_VERSION", DockerAPIVersion)
	return dcli.NewEnvClient()
}

var imageMapper = map[string]string{
	"swift":  "koderunr-swift",
	"ruby":   "koderunr-ruby",
	"python": "koderunr-python",
	"go":     "koderunr-go",
	"c":      "koderunr-c",
	"dotnet": "koderunr-dotnet",
	"fsharp": "koderunr-fsharp",
}

func (r *Runner) image() string {
	selectedVersion := r.Version
	availableVersions := (*appConfig.Languages)[r.Lang].Versions

	if selectedVersion == "" {
		if len(availableVersions) > 0 {
			selectedVersion = availableVersions[0]
		} else {
			selectedVersion = "latest"
		}
	}
	return fmt.Sprintf("%s:%s", imageMapper[r.Lang], selectedVersion)
}

func (r *Runner) createContainer(uuid string) error {
	ctx, cancel := context.WithTimeout(context.Background(), 1*time.Second)
	defer cancel()

	cmd := []string{r.Source, uuid}
	lang := (*appConfig.Languages)[r.Lang]

	ctr, err := DockerClient.ContainerCreate(ctx,
		&container.Config{
			Cmd:             cmd,
			Image:           r.image(),
			OpenStdin:       true,
			AttachStdin:     true,
			AttachStdout:    true,
			AttachStderr:    true,
			NetworkDisabled: true,
		},
		&container.HostConfig{
			Privileged: false,
			CapDrop:    []string{"all"},
			Resources: container.Resources{
				CPUQuota:   lang.GetCPUQuota(),
				MemorySwap: -1,
				Memory:     lang.GetMemory(),
				PidsLimit:  lang.GetPidsLimit(),
			},
		},
		&network.NetworkingConfig{},
		uuid,
	)

	if err != nil {
		return err
	}

	r.containerID = ctr.ID
	return nil
}

func (r *Runner) startContainer() error {
	ctx, cancel := context.WithTimeout(context.Background(), 2*time.Second)
	defer cancel()

	return DockerClient.ContainerStart(ctx, r.containerID, types.ContainerStartOptions{})

}

func (r *Runner) attachContainer(stdoutWriter *io.PipeWriter, stdinReader *io.PipeReader) error {
	hijackResp, err := DockerClient.ContainerAttach(context.Background(), r.containerID, types.ContainerAttachOptions{
		Stdin:  true,
		Stdout: true,
		Stderr: true,
		Stream: true,
	})

	if err != nil {
		return err
	}

	go func(reader *bufio.Reader) {
		for {
			io.Copy(stdoutWriter, reader)
			if err != nil {
				if err != io.EOF {
					r.logger.Error(err)
				}
				break
			}
		}
	}(hijackResp.Reader)

	go func(writer net.Conn) {
		scanner := bufio.NewScanner(stdinReader)
		for scanner.Scan() {
			fmt.Fprintf(writer, "%s\n", scanner.Text())
		}
	}(hijackResp.Conn)

	return nil
}

func (r *Runner) waitContainer(successChan chan<- struct{}, errorChan chan<- error) {
	ctx, cancel := context.WithTimeout(context.Background(), time.Duration(r.Timeout)*time.Second)
	defer cancel()

	_, err := DockerClient.ContainerWait(ctx, r.containerID)

	if err != nil {
		r.logger.Error(err)
		errorChan <- err
		return
	}
	successChan <- struct{}{}
}
<commit_msg>tidy up<commit_after>package main

import (
	"bufio"
	"encoding/json"
	"fmt"
	"io"
	"net"
	"os"
	"strconv"
	"time"

	"golang.org/x/net/context"

	"github.com/Sirupsen/logrus"
	"github.com/docker/docker/api/types"
	"github.com/docker/docker/api/types/container"
	"github.com/docker/docker/api/types/network"
	dcli "github.com/docker/docker/client"
	"github.com/garyburd/redigo/redis"
)

// DockerClient for running code
var DockerClient *dcli.Client

// DockerAPIVersion is the API version connect to docker
const DockerAPIVersion = "1.24"

// Runner runs the code
type Runner struct {
	Lang          string `json:"lang"`
	Source        string `json:"source"`
	Version       string `json:"version"`
	Timeout       int    `json:"timeout"` // How long is the code going to run
	closeNotifier <-chan bool
	logger        *logrus.Logger
	containerID   string
}

// Runnerthrottle Limit the max throttle for runner
var Runnerthrottle chan struct{}

// FetchCode get the code from Redis Server according to the UUID
func FetchCode(uuid string, redisConn redis.Conn) (r *Runner, err error) {
	value, err := redis.Bytes(redisConn.Do("GET", uuid+"#run"))

	if err != nil {
		return
	}

	r = &Runner{}
	err = json.Unmarshal(value, r)
	return
}

// Run the code in the container
func (r *Runner) Run(output messages, conn redis.Conn, uuid string) {
	Runnerthrottle <- struct{}{}
	defer func() { <-Runnerthrottle }()

	err := r.createContainer(uuid)
	if err != nil {
		r.logger.Errorf("Container %s cannot be created - %v", uuid, err)
		return
	}

	stdoutReader, stdoutWriter := io.Pipe()
	stdinReader, stdinWriter := io.Pipe()

	if err != nil {
		r.logger.Error(err)
	}

	defer stdinWriter.Close()
	defer stdoutWriter.Close()

	go pipeStdin(conn, uuid, stdinWriter, r.logger)
	go pipeStdout(stdoutReader, output, r.logger)

	// Start running the container
	err = r.startContainer()
	if err != nil {
		r.logger.Errorf("Container %s cannot be started - %v", uuid, err)
		return
	}
	defer func() {
		r.logger.Infof("Removing container %s", r.containerID)
		DockerClient.ContainerRemove(context.Background(), r.containerID, types.ContainerRemoveOptions{
			Force: true,
		})
		r.logger.Infof("Container %s removed successfully", r.containerID)
	}()

	go func(r *Runner, uuid string) {
		err = r.attachContainer(stdoutWriter, stdinReader)
		if err != nil {
			r.logger.Errorf("Container %s cannot be attached - %v", uuid, err)
		}
	}(r, uuid)

	ctx, cancel := context.WithTimeout(context.Background(), time.Duration(r.Timeout)*time.Second)
	defer cancel()

	if r.waitContainer(ctx) == nil {
		r.logger.Infof("Container %s is executed successfully", uuid)
		return
	}

	select {
	case <-r.closeNotifier:
		DockerClient.ContainerStop(context.Background(), r.containerID, nil)
		r.logger.Infof("Container %s is stopped since the streamming has been halted", uuid)
	case <-ctx.Done():
		if ctx.Err() != nil {
			msg := fmt.Sprintf("Container %s is terminated caused by 15 sec timeout\n", uuid)
			r.logger.Error(msg)
			output <- msg
		}
	}
}

func pipeStdin(conn redis.Conn, uuid string, stdin *io.PipeWriter, logger *logrus.Logger) {
	psc := redis.PubSubConn{Conn: conn}
	psc.Subscribe(uuid + "#stdin")

	defer func() {
		psc.Unsubscribe(uuid + "#stdin")
		psc.Close()
		conn.Close()
	}()

StdinSubscriptionLoop:
	for {
		switch n := psc.Receive().(type) {
		case redis.Message:
			stdinData := strconv.QuoteToASCII(string(n.Data))
			logger.Infof("Message: %s %s", n.Channel, stdinData)
			stdin.Write(n.Data)
		case error:
			break StdinSubscriptionLoop
		}
	}
	logger.Info("Stdin subscription closed")
}

func pipeStdout(stdout *io.PipeReader, output messages, logger *logrus.Logger) {
	buffer := make([]byte, 512)
	for {
		n, err := stdout.Read(buffer)
		if err != nil {
			if err != io.EOF {
				stdout.Close()
				logger.Error(err)
			}

			close(output)
			break
		}

		data := buffer[0:n]
		output <- string(data)

		// Clear the buffer
		for i := 0; i < n; i++ {
			buffer[i] = 0
		}
	}
}

// NewDockerClient create a new docker client
func NewDockerClient() (*dcli.Client, error) {
	os.Setenv("DOCKER_API_VERSION", DockerAPIVersion)
	return dcli.NewEnvClient()
}

var imageMapper = map[string]string{
	"swift":  "koderunr-swift",
	"ruby":   "koderunr-ruby",
	"python": "koderunr-python",
	"go":     "koderunr-go",
	"c":      "koderunr-c",
	"dotnet": "koderunr-dotnet",
	"fsharp": "koderunr-fsharp",
}

func (r *Runner) image() string {
	selectedVersion := r.Version
	availableVersions := (*appConfig.Languages)[r.Lang].Versions

	if selectedVersion == "" {
		if len(availableVersions) > 0 {
			selectedVersion = availableVersions[0]
		} else {
			selectedVersion = "latest"
		}
	}
	return fmt.Sprintf("%s:%s", imageMapper[r.Lang], selectedVersion)
}

func (r *Runner) createContainer(uuid string) error {
	ctx, cancel := context.WithTimeout(context.Background(), 1*time.Second)
	defer cancel()

	cmd := []string{r.Source, uuid}
	lang := (*appConfig.Languages)[r.Lang]

	ctr, err := DockerClient.ContainerCreate(ctx,
		&container.Config{
			Cmd:             cmd,
			Image:           r.image(),
			OpenStdin:       true,
			AttachStdin:     true,
			AttachStdout:    true,
			AttachStderr:    true,
			NetworkDisabled: true,
		},
		&container.HostConfig{
			Privileged: false,
			CapDrop:    []string{"all"},
			Resources: container.Resources{
				CPUQuota:   lang.GetCPUQuota(),
				MemorySwap: -1,
				Memory:     lang.GetMemory(),
				PidsLimit:  lang.GetPidsLimit(),
			},
		},
		&network.NetworkingConfig{},
		uuid,
	)

	if err != nil {
		return err
	}

	r.containerID = ctr.ID
	return nil
}

func (r *Runner) startContainer() error {
	ctx, cancel := context.WithTimeout(context.Background(), 2*time.Second)
	defer cancel()

	return DockerClient.ContainerStart(ctx, r.containerID, types.ContainerStartOptions{})

}

func (r *Runner) attachContainer(stdoutWriter *io.PipeWriter, stdinReader *io.PipeReader) error {
	hijackResp, err := DockerClient.ContainerAttach(context.Background(), r.containerID, types.ContainerAttachOptions{
		Stdin:  true,
		Stdout: true,
		Stderr: true,
		Stream: true,
	})

	if err != nil {
		return err
	}

	go func(reader *bufio.Reader) {
		for {
			io.Copy(stdoutWriter, reader)
			if err != nil {
				if err != io.EOF {
					r.logger.Error(err)
				}
				break
			}
		}
	}(hijackResp.Reader)

	go func(writer net.Conn) {
		scanner := bufio.NewScanner(stdinReader)
		for scanner.Scan() {
			fmt.Fprintf(writer, "%s\n", scanner.Text())
		}
	}(hijackResp.Conn)

	return nil
}

func (r *Runner) waitContainer(ctx context.Context) error {
	_, err := DockerClient.ContainerWait(ctx, r.containerID)

	return err
}
<|endoftext|>


<commit_before>package main

import (
	"fmt"
	"github.com/opennetworkinglab/onos-warden/warden"
	"golang.org/x/net/context"
	"google.golang.org/grpc"
	"google.golang.org/grpc/grpclog"
	"google.golang.org/grpc/peer"
	"io"
	"net"
	"sync"
	"time"
)

type cluster struct {
	ad    *warden.ClusterAdvertisement
	agent warden.ClusterAgentService_AgentClustersServer
}

type request struct {
	req    *warden.ClusterRequest
	client warden.ClusterClientService_ServerClustersServer
}

type key struct {
	cId   string
	cType string
}

type recvAd interface {
	Send(*warden.ClusterAdvertisement) error
	Context() context.Context
}

type recvReq interface {
	Send(*warden.ClusterRequest) error
	Context() context.Context
}

type wardenServer struct {
	lock sync.Mutex

	// mapping from key(ClusterId, ClusterType) to cluster resources
	clusters map[key]cluster

	// mapping from RequestId to key(ClusterId, ClusterType)
	requests map[string]key

	// registries of client and agent streams
	clients map[warden.ClusterClientService_ServerClustersServer]bool
	agents  map[warden.ClusterAgentService_AgentClustersServer]bool

	// registries of channels waiting for a cluster to be ready
	waiters map[key][]chan *warden.ClusterAdvertisement
}

func keyFromCluster(cl *cluster) key {
	return key{cl.ad.ClusterId, cl.ad.ClusterType}
}

func logClient(ctx context.Context, t string, o interface{}) {
	logConn(ctx, "Client", t, o)
}

func logAgent(ctx context.Context, t string, o interface{}) {
	logConn(ctx, "Agent", t, o)
}

func logConn(ctx context.Context, a, t string, o interface{}) {
	var pStr string
	p, ok := peer.FromContext(ctx)
	if !ok {
		pStr = "<unknown>"
	} else {
		pStr = fmt.Sprintf("%v (auth: %v)", p.Addr, p.AuthInfo)
	}
	if o != nil {
		fmt.Printf("%s: %s %s: %v\n", a, t, pStr, o)
	} else {
		fmt.Printf("%s: %s %s\n", a, t, pStr)
	}
}

func (s *wardenServer) Request(ctx context.Context, req *warden.ClusterRequest) (ad *warden.ClusterAdvertisement, err error) {
	logClient(ctx, "New request from", req)
	wait, err := s.processRequest(req)
	if err != nil {
		fmt.Printf("Error processing request %v\n%v\n", req, err)
		return nil, err
	}
	ad = <-wait
	logClient(ctx, "Sending ad to", ad)
	return ad, nil
}

func (s *wardenServer) sendSnapshot(c recvAd) error {
	// Note: callers to this method should hold s.lock
	for _, cluster := range s.clusters {
		err := c.Send(cluster.ad)
		if err != nil {
			return err
		} else {
			logClient(c.Context(), "Sending ad (snapshot) to", cluster.ad)
		}
	}
	return nil
}

func (s *wardenServer) List(_ *warden.Empty, stream warden.ClusterClientService_ListServer) error {
	logClient(stream.Context(), "List from", nil)
	s.lock.Lock()
	defer s.lock.Unlock()
	return s.sendSnapshot(stream)
}

func (s *wardenServer) ServerClusters(stream warden.ClusterClientService_ServerClustersServer) error {
	logClient(stream.Context(), "New stream from", nil)
	s.lock.Lock()
	// register the stream so that we can send it new information to all active client
	s.clients[stream] = true

	// we can use the defer mechanism to prune the stream when the stream is closed or encounters an error
	defer func() {
		s.lock.Lock()
		defer s.lock.Unlock()
		delete(s.clients, stream)
		//FIXME revoke all duration == -1 requests if client disconnects
	}()

	err := s.sendSnapshot(stream)
	s.lock.Unlock()
	if err != nil {
		return err
	}

	// Poll for requests from the client
	for {
		req, err := stream.Recv()
		if err == io.EOF {
			logClient(stream.Context(), "EOF (close) from", nil)
			return nil
		}
		if err != nil {
			logClient(stream.Context(), "Connection error from", err)
			return err
		}
		go s.processRequest(req)
	}
	return nil
}

func (s *wardenServer) sendUpdate(ad *warden.ClusterAdvertisement) {
	// Note: callers to this method should hold s.lock
	for c := range s.clients {
		err := c.Send(ad)
		if err != nil {
			logClient(c.Context(), "Error sending update to", err)
		} else {
			logClient(c.Context(), "Sent update to", ad)
		}
	}
}

func (s *wardenServer) updateCluster(cl *cluster) {
	// Note: callers must hold s.lock
	k := keyFromCluster(cl)
	existing, ok := s.clusters[k]
	if ok && cl.ad.RequestId != existing.ad.RequestId {
		// reservation is no longer assocated with the old request; delete the mapping
		delete(s.requests, existing.ad.RequestId)
	}
	s.clusters[k] = *cl
	if cl.ad.RequestId != "" {
		// update the request mapping (this is conservative, and likely won't change anything
		s.requests[cl.ad.RequestId] = k
	}

	fmt.Println("DEBUG: Got cluster", cl.ad)
	// If the cluster is ready, update all local waiters for this cluster
	if cl.ad.State == warden.ClusterAdvertisement_READY {
		w, ok := s.waiters[k]
		if ok {
			fmt.Printf("DEBUG: %d waiter(s) on %v\n", len(w), k)
			for _, ch := range w {
				ch <- cl.ad
			}
			// Remove the waiters, now that they have been updated
			delete(s.waiters, k)
		} else {
			fmt.Println("DEBUG: No waiters on", k)
		}
	}

	// Send update to all streaming clients
	s.sendUpdate(cl.ad)
}

func (s *wardenServer) deleteCluster(cl *cluster) {
	// Note: callers must hold s.lock
	k := keyFromCluster(cl)
	delete(s.clusters, k)
	if rId := cl.ad.RequestId; rId != "" {
		delete(s.requests, rId)
		//TODO need to send UNAVAILABLE
	}

	// Close all local waiters for this cluster
	w, ok := s.waiters[k]
	if ok {
		for _, ch := range w {
			close(ch)
		}
		// Remove the waiters, now that they have been closed
		delete(s.waiters, k)
	}

}

func (s *wardenServer) AgentClusters(stream warden.ClusterAgentService_AgentClustersServer) error {
	logAgent(stream.Context(), "New stream from", nil)

	s.lock.Lock()
	// register the stream into the inventory of active agent
	s.agents[stream] = true
	s.lock.Unlock()

	// defer mechanism to prune the inventory
	defer func() {
		s.lock.Lock()
		defer s.lock.Unlock()

		// remove cells from the warden map when agent disappears
		for _, cl := range s.clusters {
			//TODO maybe we should time these out instead? in case, the agent is coming right back
			if cl.agent == stream {
				s.deleteCluster(&cl)
			}
		}
		delete(s.agents, stream)
	}()

	// setup polling loop for receiving new cluster advertisements
	for {
		cl, err := stream.Recv()
		if err == io.EOF {
			logAgent(stream.Context(), "EOF (close) from", nil)
			return nil
		}
		if err != nil {
			logAgent(stream.Context(), "Connection error from", err)
			return err
		}
		logAgent(stream.Context(), "Update from", cl)
		s.lock.Lock()
		s.updateCluster(&cluster{cl, stream})
		s.lock.Unlock()
	}
	return nil
}

func (s *wardenServer) lookupRequest(req *warden.ClusterRequest) (*cluster, bool) {
	// Note: callers must hold s.lock
	rId := req.RequestId
	k, ok := s.requests[rId]
	if ok {
		ad, ok := s.clusters[k]
		if !ok {
			// this shouldn't happen, but we'll remove the stale request id to cluster id mapping
			delete(s.requests, rId)
			return nil, false
		}
		return &ad, true
	}
	return nil, false
}

func (s *wardenServer) assignRequest(req *warden.ClusterRequest) (*cluster, bool) {
	// Note: callers must hold s.lock
	for _, c := range s.clusters {
		if req.ClusterType != "" && req.ClusterType != c.ad.ClusterType {
			continue
		}
		if req.ClusterId != "" && req.ClusterId != c.ad.ClusterId {
			continue
		}
		// find the first one that is available
		if c.ad.State == warden.ClusterAdvertisement_AVAILABLE {
			k := key{c.ad.ClusterId, c.ad.ClusterType}
			// Update the request with cluster info
			req.ClusterType = c.ad.ClusterType
			req.ClusterId = c.ad.ClusterId

			// Mark the cluster as reserved internally so that it is not reassigned
			c.ad.State = warden.ClusterAdvertisement_RESERVED
			c.ad.RequestId = req.RequestId
			s.clusters[k] = c
			s.requests[req.RequestId] = k
			fmt.Println("Assigning cluster:", c.ad)
			return &c, true
		}
	}
	return nil, false
}

func (s *wardenServer) waitForReady(cl *cluster) (wait chan *warden.ClusterAdvertisement) {
	// Note: callers must hold s.lock
	// We allocate a buffered channel, so that we will not block if the cluster is already ready
	wait = make(chan *warden.ClusterAdvertisement, 1)
	if cl.ad.State == warden.ClusterAdvertisement_READY {
		// cluster is already ready, return immediately
		wait <- cl.ad
		fmt.Println("DEBUG: not waiting")
		return
	}
	// register this channel with the server to listen for updates
	k := key{cl.ad.ClusterId, cl.ad.ClusterType}
	l, ok := s.waiters[k]
	if !ok {
		l = []chan *warden.ClusterAdvertisement{wait}
	} else {
		l = append(l, wait)
	}
	fmt.Println("DEBUG: waiters for", k, ":", l)
	s.waiters[k] = l
	return wait
}

func (s *wardenServer) processRequest(req *warden.ClusterRequest) (chan *warden.ClusterAdvertisement, error) {
	s.lock.Lock()
	defer s.lock.Unlock()

	var cl *cluster
	var found bool

	// Check to see if we have already reserved a cluster for this request
	cl, found = s.lookupRequest(req)
	if !found {
		// Assign the request to an available cluster
		cl, found = s.assignRequest(req)
	}
	if !found {
		return nil, fmt.Errorf("No available clusters for req %s", req.RequestId)
	}

	// Forward the request to the agent
	err := cl.agent.Send(req)
	if err != nil {
		return nil, err
	} else {
		logAgent(cl.agent.Context(), "Sending request to", req)
	}

	// Wait for the cluster to become ready
	return s.waitForReady(cl), nil
}

func (s *wardenServer) returnCluster(cl *cluster) {
	// Note: callers must hold s.lock
	// Mark the cluster as unavailable internally in case a client asks
	k := keyFromCluster(cl)
	cl.ad.State = warden.ClusterAdvertisement_UNAVAILABLE
	s.clusters[k] = *cl

	// Build minimal request based on cluster advertisement
	req := warden.ClusterRequest{
		ClusterId:   cl.ad.ClusterId,
		ClusterType: cl.ad.ClusterType,
		RequestId:   cl.ad.RequestId,
		Type:        warden.ClusterRequest_RETURN,
	}
	logAgent(cl.agent.Context(), "Sending internal return request to", req)
	cl.agent.Send(&req)
}

func (s *wardenServer) cleanupStaleClusters() {
	for {
		s.lock.Lock()
		for _, cl := range s.clusters {
			switch cl.ad.State {
			case warden.ClusterAdvertisement_RESERVED:
				fallthrough
			case warden.ClusterAdvertisement_READY:
				info := cl.ad.ReservationInfo
				if info != nil {
					if info.Duration < 0 {
						// Reservation does not expire
						continue
					}
					end := time.Unix(info.ReservationStartTime, 0)
					end = end.Add(time.Duration(info.Duration) * time.Minute)
					if end.After(time.Now()) {
						// Reservation has not expired yet
						fmt.Println("Time remaining in seconds", cl.ad.ClusterId, cl.ad.ClusterType, end.Sub(time.Now()))
						continue
					}
				}
				fmt.Println("Cluster timed out:", cl.ad)
				s.returnCluster(&cl)
			}
		}
		s.lock.Unlock()
		time.Sleep(20 * time.Second)
	}
}

func newServer() *wardenServer {
	s := new(wardenServer)
	s.clusters = make(map[key]cluster)
	s.requests = make(map[string]key)
	s.clients = make(map[warden.ClusterClientService_ServerClustersServer]bool)
	s.agents = make(map[warden.ClusterAgentService_AgentClustersServer]bool)
	s.waiters = make(map[key][]chan *warden.ClusterAdvertisement)
	return s
}

func main() {
	lis, err := net.Listen("tcp", ":1234")
	if err != nil {
		grpclog.Fatalf("failed to listen: %v", err)
	}
	grpcServer := grpc.NewServer()
	s := newServer()
	go s.cleanupStaleClusters()
	warden.RegisterClusterClientServiceServer(grpcServer, s)
	warden.RegisterClusterAgentServiceServer(grpcServer, s)
	fmt.Println("starting to serve...")
	grpcServer.Serve(lis)
}
<commit_msg>Removing server waiter debug lines<commit_after>package main

import (
	"fmt"
	"github.com/opennetworkinglab/onos-warden/warden"
	"golang.org/x/net/context"
	"google.golang.org/grpc"
	"google.golang.org/grpc/grpclog"
	"google.golang.org/grpc/peer"
	"io"
	"net"
	"sync"
	"time"
)

type cluster struct {
	ad    *warden.ClusterAdvertisement
	agent warden.ClusterAgentService_AgentClustersServer
}

type request struct {
	req    *warden.ClusterRequest
	client warden.ClusterClientService_ServerClustersServer
}

type key struct {
	cId   string
	cType string
}

type recvAd interface {
	Send(*warden.ClusterAdvertisement) error
	Context() context.Context
}

type recvReq interface {
	Send(*warden.ClusterRequest) error
	Context() context.Context
}

type wardenServer struct {
	lock sync.Mutex

	// mapping from key(ClusterId, ClusterType) to cluster resources
	clusters map[key]cluster

	// mapping from RequestId to key(ClusterId, ClusterType)
	requests map[string]key

	// registries of client and agent streams
	clients map[warden.ClusterClientService_ServerClustersServer]bool
	agents  map[warden.ClusterAgentService_AgentClustersServer]bool

	// registries of channels waiting for a cluster to be ready
	waiters map[key][]chan *warden.ClusterAdvertisement
}

func keyFromCluster(cl *cluster) key {
	return key{cl.ad.ClusterId, cl.ad.ClusterType}
}

func logClient(ctx context.Context, t string, o interface{}) {
	logConn(ctx, "Client", t, o)
}

func logAgent(ctx context.Context, t string, o interface{}) {
	logConn(ctx, "Agent", t, o)
}

func logConn(ctx context.Context, a, t string, o interface{}) {
	var pStr string
	p, ok := peer.FromContext(ctx)
	if !ok {
		pStr = "<unknown>"
	} else {
		pStr = fmt.Sprintf("%v (auth: %v)", p.Addr, p.AuthInfo)
	}
	if o != nil {
		fmt.Printf("%s: %s %s: %v\n", a, t, pStr, o)
	} else {
		fmt.Printf("%s: %s %s\n", a, t, pStr)
	}
}

func (s *wardenServer) Request(ctx context.Context, req *warden.ClusterRequest) (ad *warden.ClusterAdvertisement, err error) {
	logClient(ctx, "New request from", req)
	wait, err := s.processRequest(req)
	if err != nil {
		fmt.Printf("Error processing request %v\n%v\n", req, err)
		return nil, err
	}
	ad = <-wait
	logClient(ctx, "Sending ad to", ad)
	return ad, nil
}

func (s *wardenServer) sendSnapshot(c recvAd) error {
	// Note: callers to this method should hold s.lock
	for _, cluster := range s.clusters {
		err := c.Send(cluster.ad)
		if err != nil {
			return err
		} else {
			logClient(c.Context(), "Sending ad (snapshot) to", cluster.ad)
		}
	}
	return nil
}

func (s *wardenServer) List(_ *warden.Empty, stream warden.ClusterClientService_ListServer) error {
	logClient(stream.Context(), "List from", nil)
	s.lock.Lock()
	defer s.lock.Unlock()
	return s.sendSnapshot(stream)
}

func (s *wardenServer) ServerClusters(stream warden.ClusterClientService_ServerClustersServer) error {
	logClient(stream.Context(), "New stream from", nil)
	s.lock.Lock()
	// register the stream so that we can send it new information to all active client
	s.clients[stream] = true

	// we can use the defer mechanism to prune the stream when the stream is closed or encounters an error
	defer func() {
		s.lock.Lock()
		defer s.lock.Unlock()
		delete(s.clients, stream)
		//FIXME revoke all duration == -1 requests if client disconnects
	}()

	err := s.sendSnapshot(stream)
	s.lock.Unlock()
	if err != nil {
		return err
	}

	// Poll for requests from the client
	for {
		req, err := stream.Recv()
		if err == io.EOF {
			logClient(stream.Context(), "EOF (close) from", nil)
			return nil
		}
		if err != nil {
			logClient(stream.Context(), "Connection error from", err)
			return err
		}
		go s.processRequest(req)
	}
	return nil
}

func (s *wardenServer) sendUpdate(ad *warden.ClusterAdvertisement) {
	// Note: callers to this method should hold s.lock
	for c := range s.clients {
		err := c.Send(ad)
		if err != nil {
			logClient(c.Context(), "Error sending update to", err)
		} else {
			logClient(c.Context(), "Sent update to", ad)
		}
	}
}

func (s *wardenServer) updateCluster(cl *cluster) {
	// Note: callers must hold s.lock
	k := keyFromCluster(cl)
	existing, ok := s.clusters[k]
	if ok && cl.ad.RequestId != existing.ad.RequestId {
		// reservation is no longer assocated with the old request; delete the mapping
		delete(s.requests, existing.ad.RequestId)
	}
	s.clusters[k] = *cl
	if cl.ad.RequestId != "" {
		// update the request mapping (this is conservative, and likely won't change anything
		s.requests[cl.ad.RequestId] = k
	}

	// If the cluster is ready, update all local waiters for this cluster
	if cl.ad.State == warden.ClusterAdvertisement_READY {
		w, ok := s.waiters[k]
		if ok {
			for _, ch := range w {
				ch <- cl.ad
			}
			// Remove the waiters, now that they have been updated
			delete(s.waiters, k)
		}
	}

	// Send update to all streaming clients
	s.sendUpdate(cl.ad)
}

func (s *wardenServer) deleteCluster(cl *cluster) {
	// Note: callers must hold s.lock
	k := keyFromCluster(cl)
	delete(s.clusters, k)
	if rId := cl.ad.RequestId; rId != "" {
		delete(s.requests, rId)
		//TODO need to send UNAVAILABLE
	}

	// Close all local waiters for this cluster
	w, ok := s.waiters[k]
	if ok {
		for _, ch := range w {
			close(ch)
		}
		// Remove the waiters, now that they have been closed
		delete(s.waiters, k)
	}

}

func (s *wardenServer) AgentClusters(stream warden.ClusterAgentService_AgentClustersServer) error {
	logAgent(stream.Context(), "New stream from", nil)

	s.lock.Lock()
	// register the stream into the inventory of active agent
	s.agents[stream] = true
	s.lock.Unlock()

	// defer mechanism to prune the inventory
	defer func() {
		s.lock.Lock()
		defer s.lock.Unlock()

		// remove cells from the warden map when agent disappears
		for _, cl := range s.clusters {
			//TODO maybe we should time these out instead? in case, the agent is coming right back
			if cl.agent == stream {
				s.deleteCluster(&cl)
			}
		}
		delete(s.agents, stream)
	}()

	// setup polling loop for receiving new cluster advertisements
	for {
		cl, err := stream.Recv()
		if err == io.EOF {
			logAgent(stream.Context(), "EOF (close) from", nil)
			return nil
		}
		if err != nil {
			logAgent(stream.Context(), "Connection error from", err)
			return err
		}
		logAgent(stream.Context(), "Update from", cl)
		s.lock.Lock()
		s.updateCluster(&cluster{cl, stream})
		s.lock.Unlock()
	}
	return nil
}

func (s *wardenServer) lookupRequest(req *warden.ClusterRequest) (*cluster, bool) {
	// Note: callers must hold s.lock
	rId := req.RequestId
	k, ok := s.requests[rId]
	if ok {
		ad, ok := s.clusters[k]
		if !ok {
			// this shouldn't happen, but we'll remove the stale request id to cluster id mapping
			delete(s.requests, rId)
			return nil, false
		}
		return &ad, true
	}
	return nil, false
}

func (s *wardenServer) assignRequest(req *warden.ClusterRequest) (*cluster, bool) {
	// Note: callers must hold s.lock
	for _, c := range s.clusters {
		if req.ClusterType != "" && req.ClusterType != c.ad.ClusterType {
			continue
		}
		if req.ClusterId != "" && req.ClusterId != c.ad.ClusterId {
			continue
		}
		// find the first one that is available
		if c.ad.State == warden.ClusterAdvertisement_AVAILABLE {
			k := key{c.ad.ClusterId, c.ad.ClusterType}
			// Update the request with cluster info
			req.ClusterType = c.ad.ClusterType
			req.ClusterId = c.ad.ClusterId

			// Mark the cluster as reserved internally so that it is not reassigned
			c.ad.State = warden.ClusterAdvertisement_RESERVED
			c.ad.RequestId = req.RequestId
			s.clusters[k] = c
			s.requests[req.RequestId] = k
			fmt.Println("Assigning cluster:", c.ad)
			return &c, true
		}
	}
	return nil, false
}

func (s *wardenServer) waitForReady(cl *cluster) (wait chan *warden.ClusterAdvertisement) {
	// Note: callers must hold s.lock
	// We allocate a buffered channel, so that we will not block if the cluster is already ready
	wait = make(chan *warden.ClusterAdvertisement, 1)
	if cl.ad.State == warden.ClusterAdvertisement_READY {
		// cluster is already ready, return immediately
		wait <- cl.ad
		return
	}
	// register this channel with the server to listen for updates
	k := key{cl.ad.ClusterId, cl.ad.ClusterType}
	l, ok := s.waiters[k]
	if !ok {
		l = []chan *warden.ClusterAdvertisement{wait}
	} else {
		l = append(l, wait)
	}
	s.waiters[k] = l
	return wait
}

func (s *wardenServer) processRequest(req *warden.ClusterRequest) (chan *warden.ClusterAdvertisement, error) {
	s.lock.Lock()
	defer s.lock.Unlock()

	var cl *cluster
	var found bool

	// Check to see if we have already reserved a cluster for this request
	cl, found = s.lookupRequest(req)
	if !found {
		// Assign the request to an available cluster
		cl, found = s.assignRequest(req)
	}
	if !found {
		return nil, fmt.Errorf("No available clusters for req %s", req.RequestId)
	}

	// Forward the request to the agent
	err := cl.agent.Send(req)
	if err != nil {
		return nil, err
	} else {
		logAgent(cl.agent.Context(), "Sending request to", req)
	}

	// Wait for the cluster to become ready
	return s.waitForReady(cl), nil
}

func (s *wardenServer) returnCluster(cl *cluster) {
	// Note: callers must hold s.lock
	// Mark the cluster as unavailable internally in case a client asks
	k := keyFromCluster(cl)
	cl.ad.State = warden.ClusterAdvertisement_UNAVAILABLE
	s.clusters[k] = *cl

	// Build minimal request based on cluster advertisement
	req := warden.ClusterRequest{
		ClusterId:   cl.ad.ClusterId,
		ClusterType: cl.ad.ClusterType,
		RequestId:   cl.ad.RequestId,
		Type:        warden.ClusterRequest_RETURN,
	}
	logAgent(cl.agent.Context(), "Sending internal return request to", req)
	cl.agent.Send(&req)
}

func (s *wardenServer) cleanupStaleClusters() {
	for {
		s.lock.Lock()
		for _, cl := range s.clusters {
			switch cl.ad.State {
			case warden.ClusterAdvertisement_RESERVED:
				fallthrough
			case warden.ClusterAdvertisement_READY:
				info := cl.ad.ReservationInfo
				if info != nil {
					if info.Duration < 0 {
						// Reservation does not expire
						continue
					}
					end := time.Unix(info.ReservationStartTime, 0)
					end = end.Add(time.Duration(info.Duration) * time.Minute)
					if end.After(time.Now()) {
						// Reservation has not expired yet
						fmt.Println("Time remaining in seconds", cl.ad.ClusterId, cl.ad.ClusterType, end.Sub(time.Now()))
						continue
					}
				}
				fmt.Println("Cluster timed out:", cl.ad)
				s.returnCluster(&cl)
			}
		}
		s.lock.Unlock()
		time.Sleep(20 * time.Second)
	}
}

func newServer() *wardenServer {
	s := new(wardenServer)
	s.clusters = make(map[key]cluster)
	s.requests = make(map[string]key)
	s.clients = make(map[warden.ClusterClientService_ServerClustersServer]bool)
	s.agents = make(map[warden.ClusterAgentService_AgentClustersServer]bool)
	s.waiters = make(map[key][]chan *warden.ClusterAdvertisement)
	return s
}

func main() {
	lis, err := net.Listen("tcp", ":1234")
	if err != nil {
		grpclog.Fatalf("failed to listen: %v", err)
	}
	grpcServer := grpc.NewServer()
	s := newServer()
	go s.cleanupStaleClusters()
	warden.RegisterClusterClientServiceServer(grpcServer, s)
	warden.RegisterClusterAgentServiceServer(grpcServer, s)
	fmt.Println("starting to serve...")
	grpcServer.Serve(lis)
}
<|endoftext|>