<commit_before>// Package torrent provides a BitTorrent client implementation.
package torrent

import (
	"context"
	"crypto/tls"
	"errors"
	"net"
	"net/http"
	"os"
	"path/filepath"
	"strings"
	"sync"
	"time"

	"github.com/boltdb/bolt"
	"github.com/cenkalti/rain/internal/bitfield"
	"github.com/cenkalti/rain/internal/blocklist"
	"github.com/cenkalti/rain/internal/logger"
	"github.com/cenkalti/rain/internal/piececache"
	"github.com/cenkalti/rain/internal/resolver"
	"github.com/cenkalti/rain/internal/resourcemanager"
	"github.com/cenkalti/rain/internal/resumer/boltdbresumer"
	"github.com/cenkalti/rain/internal/semaphore"
	"github.com/cenkalti/rain/internal/storage/filestorage"
	"github.com/cenkalti/rain/internal/tracker"
	"github.com/cenkalti/rain/internal/trackermanager"
	"github.com/juju/ratelimit"
	"github.com/mitchellh/go-homedir"
	"github.com/nictuku/dht"
)

var (
	sessionBucket         = []byte("session")
	torrentsBucket        = []byte("torrents")
	blocklistKey          = []byte("blocklist")
	blocklistTimestampKey = []byte("blocklist-timestamp")
)

// Session contains torrents, DHT node, caches and other data structures shared by multiple torrents.
type Session struct {
	config         Config
	db             *bolt.DB
	resumer        *boltdbresumer.Resumer
	log            logger.Logger
	extensions     [8]byte
	dht            *dht.DHT
	rpc            *rpcServer
	trackerManager *trackermanager.TrackerManager
	ram            *resourcemanager.ResourceManager
	pieceCache     *piececache.Cache
	webseedClient  http.Client
	createdAt      time.Time
	semWrite       *semaphore.Semaphore
	metrics        *sessionMetrics
	bucketDownload *ratelimit.Bucket
	bucketUpload   *ratelimit.Bucket
	closeC         chan struct{}

	mPeerRequests   sync.Mutex
	dhtPeerRequests map[*torrent]struct{}

	mTorrents          sync.RWMutex
	torrents           map[string]*Torrent
	torrentsByInfoHash map[dht.InfoHash][]*Torrent

	mPorts         sync.RWMutex
	availablePorts map[int]struct{}

	mBlocklist         sync.RWMutex
	blocklist          *blocklist.Blocklist
	blocklistTimestamp time.Time
}

// NewSession creates a new Session for downloading and seeding torrents.
// Returned session must be closed after use.
func NewSession(cfg Config) (*Session, error) {
	if cfg.PortBegin >= cfg.PortEnd {
		return nil, errors.New("invalid port range")
	}
	var err error
	cfg.Database, err = homedir.Expand(cfg.Database)
	if err != nil {
		return nil, err
	}
	cfg.DataDir, err = homedir.Expand(cfg.DataDir)
	if err != nil {
		return nil, err
	}
	err = os.MkdirAll(filepath.Dir(cfg.Database), 0750)
	if err != nil {
		return nil, err
	}
	l := logger.New("session")
	db, err := bolt.Open(cfg.Database, 0640, &bolt.Options{Timeout: time.Second})
	if err == bolt.ErrTimeout {
		return nil, errors.New("resume database is locked by another process")
	} else if err != nil {
		return nil, err
	}
	defer func() {
		if err != nil {
			db.Close()
		}
	}()
	var ids []string
	err = db.Update(func(tx *bolt.Tx) error {
		_, err2 := tx.CreateBucketIfNotExists(sessionBucket)
		if err2 != nil {
			return err2
		}
		b, err2 := tx.CreateBucketIfNotExists(torrentsBucket)
		if err2 != nil {
			return err2
		}
		return b.ForEach(func(k, _ []byte) error {
			ids = append(ids, string(k))
			return nil
		})
	})
	if err != nil {
		return nil, err
	}
	res, err := boltdbresumer.New(db, torrentsBucket)
	if err != nil {
		return nil, err
	}
	var dhtNode *dht.DHT
	if cfg.DHTEnabled {
		dhtConfig := dht.NewConfig()
		dhtConfig.Address = cfg.DHTHost
		dhtConfig.Port = int(cfg.DHTPort)
		dhtConfig.DHTRouters = strings.Join(cfg.DHTBootstrapNodes, ",")
		dhtConfig.SaveRoutingTable = false
		dhtNode, err = dht.New(dhtConfig)
		if err != nil {
			return nil, err
		}
		err = dhtNode.Start()
		if err != nil {
			return nil, err
		}
	}
	ports := make(map[int]struct{})
	for p := cfg.PortBegin; p < cfg.PortEnd; p++ {
		ports[int(p)] = struct{}{}
	}
	bl := blocklist.New()
	var blTracker *blocklist.Blocklist
	if cfg.BlocklistEnabledForTrackers {
		blTracker = bl
	}
	c := &Session{
		config:             cfg,
		db:                 db,
		resumer:            res,
		blocklist:          bl,
		trackerManager:     trackermanager.New(blTracker, cfg.DNSResolveTimeout, !cfg.TrackerHTTPVerifyTLS),
		log:                l,
		torrents:           make(map[string]*Torrent),
		torrentsByInfoHash: make(map[dht.InfoHash][]*Torrent),
		availablePorts:     ports,
		dht:                dhtNode,
		pieceCache:         piececache.New(cfg.ReadCacheSize, cfg.ReadCacheTTL, cfg.ParallelReads),
		ram:                resourcemanager.New(cfg.WriteCacheSize),
		createdAt:          time.Now(),
		semWrite:           semaphore.New(int(cfg.ParallelWrites)),
		closeC:             make(chan struct{}),
		webseedClient: http.Client{
			Transport: &http.Transport{
				DialContext: func(ctx context.Context, network, addr string) (net.Conn, error) {
					ip, port, err := resolver.Resolve(ctx, addr, cfg.DNSResolveTimeout, bl)
					if err != nil {
						return nil, err
					}
					var d net.Dialer
					taddr := &net.TCPAddr{IP: ip, Port: port}
					dctx, cancel := context.WithTimeout(ctx, cfg.WebseedDialTimeout)
					defer cancel()
					return d.DialContext(dctx, network, taddr.String())
				},
				TLSHandshakeTimeout:   cfg.WebseedTLSHandshakeTimeout,
				TLSClientConfig:       &tls.Config{InsecureSkipVerify: !cfg.WebseedVerifyTLS}, // nolint: gosec
				ResponseHeaderTimeout: cfg.WebseedResponseHeaderTimeout,
			},
		},
	}
	if cfg.SpeedLimitDownload > 0 {
		c.bucketDownload = ratelimit.NewBucketWithRate(float64(cfg.SpeedLimitDownload), cfg.SpeedLimitDownload)
	}
	if cfg.SpeedLimitUpload > 0 {
		c.bucketUpload = ratelimit.NewBucketWithRate(float64(cfg.SpeedLimitUpload), cfg.SpeedLimitUpload)
	}
	err = c.startBlocklistReloader()
	if err != nil {
		return nil, err
	}
	ext, err := bitfield.NewBytes(c.extensions[:], 64)
	if err != nil {
		panic(err)
	}
	ext.Set(61) // Fast Extension (BEP 6)
	ext.Set(43) // Extension Protocol (BEP 10)
	if cfg.DHTEnabled {
		ext.Set(63) // DHT Protocol (BEP 5)
		c.dhtPeerRequests = make(map[*torrent]struct{})
	}
	c.loadExistingTorrents(ids)
	if c.config.RPCEnabled {
		c.rpc = newRPCServer(c)
		err = c.rpc.Start(c.config.RPCHost, c.config.RPCPort)
		if err != nil {
			return nil, err
		}
	}
	if cfg.DHTEnabled {
		go c.processDHTResults()
	}
	go c.updateStatsLoop()
	err = c.initMetrics()
	if err != nil {
		return nil, err
	}
	return c, nil
}

func (s *Session) parseTrackers(tiers [][]string, private bool) []tracker.Tracker {
	ret := make([]tracker.Tracker, 0, len(tiers))
	for _, tier := range tiers {
		trackers := make([]tracker.Tracker, 0, len(tier))
		for _, tr := range tier {
			t, err := s.trackerManager.Get(tr, s.config.TrackerHTTPTimeout, s.getTrackerUserAgent(private), int64(s.config.TrackerHTTPMaxResponseSize))
			if err != nil {
				continue
			}
			trackers = append(trackers, t)
		}
		if len(trackers) > 0 {
			tra := tracker.NewTier(trackers)
			ret = append(ret, tra)
		}
	}
	return ret
}

func (s *Session) getTrackerUserAgent(private bool) string {
	if private {
		return s.config.TrackerHTTPPrivateUserAgent
	}
	return trackerHTTPPublicUserAgent
}

func (s *Session) Close() error {
	close(s.closeC)

	if s.config.DHTEnabled {
		s.dht.Stop()
	}

	s.updateStats()

	var wg sync.WaitGroup
	s.mTorrents.Lock()
	wg.Add(len(s.torrents))
	for _, t := range s.torrents {
		go func(t *Torrent) {
			t.torrent.Close()
			wg.Done()
		}(t)
	}
	wg.Wait()
	s.torrents = nil
	s.mTorrents.Unlock()

	if s.rpc != nil {
		err := s.rpc.Stop(s.config.RPCShutdownTimeout)
		if err != nil {
			s.log.Errorln("cannot stop RPC server:", err.Error())
		}
	}

	s.ram.Close()
	s.pieceCache.Close()
	s.metrics.Close()
	return s.db.Close()
}

func (s *Session) ListTorrents() []*Torrent {
	s.mTorrents.RLock()
	defer s.mTorrents.RUnlock()
	torrents := make([]*Torrent, 0, len(s.torrents))
	for _, t := range s.torrents {
		torrents = append(torrents, t)
	}
	return torrents
}

func (s *Session) getPort() (int, error) {
	s.mPorts.Lock()
	defer s.mPorts.Unlock()
	for p := range s.availablePorts {
		delete(s.availablePorts, p)
		return p, nil
	}
	return 0, errors.New("no free port")
}

func (s *Session) releasePort(port int) {
	s.mPorts.Lock()
	defer s.mPorts.Unlock()
	s.availablePorts[port] = struct{}{}
}

func (s *Session) GetTorrent(id string) *Torrent {
	s.mTorrents.RLock()
	defer s.mTorrents.RUnlock()
	return s.torrents[id]
}

func (s *Session) RemoveTorrent(id string) error {
	t, err := s.removeTorrentFromClient(id)
	if t != nil {
		go s.stopAndRemoveData(t)
	}
	return err
}

func (s *Session) removeTorrentFromClient(id string) (*Torrent, error) {
	s.mTorrents.Lock()
	defer s.mTorrents.Unlock()
	t, ok := s.torrents[id]
	if !ok {
		return nil, nil
	}
	t.torrent.log.Info("removing torrent")
	delete(s.torrents, id)
	delete(s.torrentsByInfoHash, dht.InfoHash(t.torrent.InfoHash()))
	return t, s.db.Update(func(tx *bolt.Tx) error {
		return tx.Bucket(torrentsBucket).DeleteBucket([]byte(id))
	})
}

func (s *Session) stopAndRemoveData(t *Torrent) {
	t.torrent.Close()
	s.releasePort(t.torrent.port)
	dest := t.torrent.storage.(*filestorage.FileStorage).Dest()
	err := os.RemoveAll(dest)
	if err != nil {
		s.log.Errorf("cannot remove torrent data. err: %s dest: %s", err, dest)
	}
}

func (s *Session) StartAll() error {
	err := s.db.Update(func(tx *bolt.Tx) error {
		tb := tx.Bucket(torrentsBucket)
		s.mTorrents.RLock()
		for _, t := range s.torrents {
			b := tb.Bucket([]byte(t.torrent.id))
			_ = b.Put([]byte("started"), []byte("1"))
		}
		defer s.mTorrents.RUnlock()
		return nil
	})
	if err != nil {
		return err
	}
	for _, t := range s.torrents {
		t.torrent.Start()
	}
	return nil
}

func (s *Session) StopAll() error {
	err := s.db.Update(func(tx *bolt.Tx) error {
		tb := tx.Bucket(torrentsBucket)
		s.mTorrents.RLock()
		for _, t := range s.torrents {
			b := tb.Bucket([]byte(t.torrent.id))
			_ = b.Put([]byte("started"), []byte("0"))
		}
		defer s.mTorrents.RUnlock()
		return nil
	})
	if err != nil {
		return err
	}
	for _, t := range s.torrents {
		t.torrent.Stop()
	}
	return nil
}
<commit_msg>init metrics before loading existing torrents<commit_after>// Package torrent provides a BitTorrent client implementation.
package torrent

import (
	"context"
	"crypto/tls"
	"errors"
	"net"
	"net/http"
	"os"
	"path/filepath"
	"strings"
	"sync"
	"time"

	"github.com/boltdb/bolt"
	"github.com/cenkalti/rain/internal/bitfield"
	"github.com/cenkalti/rain/internal/blocklist"
	"github.com/cenkalti/rain/internal/logger"
	"github.com/cenkalti/rain/internal/piececache"
	"github.com/cenkalti/rain/internal/resolver"
	"github.com/cenkalti/rain/internal/resourcemanager"
	"github.com/cenkalti/rain/internal/resumer/boltdbresumer"
	"github.com/cenkalti/rain/internal/semaphore"
	"github.com/cenkalti/rain/internal/storage/filestorage"
	"github.com/cenkalti/rain/internal/tracker"
	"github.com/cenkalti/rain/internal/trackermanager"
	"github.com/juju/ratelimit"
	"github.com/mitchellh/go-homedir"
	"github.com/nictuku/dht"
)

var (
	sessionBucket         = []byte("session")
	torrentsBucket        = []byte("torrents")
	blocklistKey          = []byte("blocklist")
	blocklistTimestampKey = []byte("blocklist-timestamp")
)

// Session contains torrents, DHT node, caches and other data structures shared by multiple torrents.
type Session struct {
	config         Config
	db             *bolt.DB
	resumer        *boltdbresumer.Resumer
	log            logger.Logger
	extensions     [8]byte
	dht            *dht.DHT
	rpc            *rpcServer
	trackerManager *trackermanager.TrackerManager
	ram            *resourcemanager.ResourceManager
	pieceCache     *piececache.Cache
	webseedClient  http.Client
	createdAt      time.Time
	semWrite       *semaphore.Semaphore
	metrics        *sessionMetrics
	bucketDownload *ratelimit.Bucket
	bucketUpload   *ratelimit.Bucket
	closeC         chan struct{}

	mPeerRequests   sync.Mutex
	dhtPeerRequests map[*torrent]struct{}

	mTorrents          sync.RWMutex
	torrents           map[string]*Torrent
	torrentsByInfoHash map[dht.InfoHash][]*Torrent

	mPorts         sync.RWMutex
	availablePorts map[int]struct{}

	mBlocklist         sync.RWMutex
	blocklist          *blocklist.Blocklist
	blocklistTimestamp time.Time
}

// NewSession creates a new Session for downloading and seeding torrents.
// Returned session must be closed after use.
func NewSession(cfg Config) (*Session, error) {
	if cfg.PortBegin >= cfg.PortEnd {
		return nil, errors.New("invalid port range")
	}
	var err error
	cfg.Database, err = homedir.Expand(cfg.Database)
	if err != nil {
		return nil, err
	}
	cfg.DataDir, err = homedir.Expand(cfg.DataDir)
	if err != nil {
		return nil, err
	}
	err = os.MkdirAll(filepath.Dir(cfg.Database), 0750)
	if err != nil {
		return nil, err
	}
	l := logger.New("session")
	db, err := bolt.Open(cfg.Database, 0640, &bolt.Options{Timeout: time.Second})
	if err == bolt.ErrTimeout {
		return nil, errors.New("resume database is locked by another process")
	} else if err != nil {
		return nil, err
	}
	defer func() {
		if err != nil {
			db.Close()
		}
	}()
	var ids []string
	err = db.Update(func(tx *bolt.Tx) error {
		_, err2 := tx.CreateBucketIfNotExists(sessionBucket)
		if err2 != nil {
			return err2
		}
		b, err2 := tx.CreateBucketIfNotExists(torrentsBucket)
		if err2 != nil {
			return err2
		}
		return b.ForEach(func(k, _ []byte) error {
			ids = append(ids, string(k))
			return nil
		})
	})
	if err != nil {
		return nil, err
	}
	res, err := boltdbresumer.New(db, torrentsBucket)
	if err != nil {
		return nil, err
	}
	var dhtNode *dht.DHT
	if cfg.DHTEnabled {
		dhtConfig := dht.NewConfig()
		dhtConfig.Address = cfg.DHTHost
		dhtConfig.Port = int(cfg.DHTPort)
		dhtConfig.DHTRouters = strings.Join(cfg.DHTBootstrapNodes, ",")
		dhtConfig.SaveRoutingTable = false
		dhtNode, err = dht.New(dhtConfig)
		if err != nil {
			return nil, err
		}
		err = dhtNode.Start()
		if err != nil {
			return nil, err
		}
	}
	ports := make(map[int]struct{})
	for p := cfg.PortBegin; p < cfg.PortEnd; p++ {
		ports[int(p)] = struct{}{}
	}
	bl := blocklist.New()
	var blTracker *blocklist.Blocklist
	if cfg.BlocklistEnabledForTrackers {
		blTracker = bl
	}
	c := &Session{
		config:             cfg,
		db:                 db,
		resumer:            res,
		blocklist:          bl,
		trackerManager:     trackermanager.New(blTracker, cfg.DNSResolveTimeout, !cfg.TrackerHTTPVerifyTLS),
		log:                l,
		torrents:           make(map[string]*Torrent),
		torrentsByInfoHash: make(map[dht.InfoHash][]*Torrent),
		availablePorts:     ports,
		dht:                dhtNode,
		pieceCache:         piececache.New(cfg.ReadCacheSize, cfg.ReadCacheTTL, cfg.ParallelReads),
		ram:                resourcemanager.New(cfg.WriteCacheSize),
		createdAt:          time.Now(),
		semWrite:           semaphore.New(int(cfg.ParallelWrites)),
		closeC:             make(chan struct{}),
		webseedClient: http.Client{
			Transport: &http.Transport{
				DialContext: func(ctx context.Context, network, addr string) (net.Conn, error) {
					ip, port, err := resolver.Resolve(ctx, addr, cfg.DNSResolveTimeout, bl)
					if err != nil {
						return nil, err
					}
					var d net.Dialer
					taddr := &net.TCPAddr{IP: ip, Port: port}
					dctx, cancel := context.WithTimeout(ctx, cfg.WebseedDialTimeout)
					defer cancel()
					return d.DialContext(dctx, network, taddr.String())
				},
				TLSHandshakeTimeout:   cfg.WebseedTLSHandshakeTimeout,
				TLSClientConfig:       &tls.Config{InsecureSkipVerify: !cfg.WebseedVerifyTLS}, // nolint: gosec
				ResponseHeaderTimeout: cfg.WebseedResponseHeaderTimeout,
			},
		},
	}
	if cfg.SpeedLimitDownload > 0 {
		c.bucketDownload = ratelimit.NewBucketWithRate(float64(cfg.SpeedLimitDownload), cfg.SpeedLimitDownload)
	}
	if cfg.SpeedLimitUpload > 0 {
		c.bucketUpload = ratelimit.NewBucketWithRate(float64(cfg.SpeedLimitUpload), cfg.SpeedLimitUpload)
	}
	err = c.startBlocklistReloader()
	if err != nil {
		return nil, err
	}
	ext, err := bitfield.NewBytes(c.extensions[:], 64)
	if err != nil {
		panic(err)
	}
	ext.Set(61) // Fast Extension (BEP 6)
	ext.Set(43) // Extension Protocol (BEP 10)
	if cfg.DHTEnabled {
		ext.Set(63) // DHT Protocol (BEP 5)
		c.dhtPeerRequests = make(map[*torrent]struct{})
	}
	err = c.initMetrics()
	if err != nil {
		return nil, err
	}
	c.loadExistingTorrents(ids)
	if c.config.RPCEnabled {
		c.rpc = newRPCServer(c)
		err = c.rpc.Start(c.config.RPCHost, c.config.RPCPort)
		if err != nil {
			return nil, err
		}
	}
	if cfg.DHTEnabled {
		go c.processDHTResults()
	}
	go c.updateStatsLoop()
	return c, nil
}

func (s *Session) parseTrackers(tiers [][]string, private bool) []tracker.Tracker {
	ret := make([]tracker.Tracker, 0, len(tiers))
	for _, tier := range tiers {
		trackers := make([]tracker.Tracker, 0, len(tier))
		for _, tr := range tier {
			t, err := s.trackerManager.Get(tr, s.config.TrackerHTTPTimeout, s.getTrackerUserAgent(private), int64(s.config.TrackerHTTPMaxResponseSize))
			if err != nil {
				continue
			}
			trackers = append(trackers, t)
		}
		if len(trackers) > 0 {
			tra := tracker.NewTier(trackers)
			ret = append(ret, tra)
		}
	}
	return ret
}

func (s *Session) getTrackerUserAgent(private bool) string {
	if private {
		return s.config.TrackerHTTPPrivateUserAgent
	}
	return trackerHTTPPublicUserAgent
}

func (s *Session) Close() error {
	close(s.closeC)

	if s.config.DHTEnabled {
		s.dht.Stop()
	}

	s.updateStats()

	var wg sync.WaitGroup
	s.mTorrents.Lock()
	wg.Add(len(s.torrents))
	for _, t := range s.torrents {
		go func(t *Torrent) {
			t.torrent.Close()
			wg.Done()
		}(t)
	}
	wg.Wait()
	s.torrents = nil
	s.mTorrents.Unlock()

	if s.rpc != nil {
		err := s.rpc.Stop(s.config.RPCShutdownTimeout)
		if err != nil {
			s.log.Errorln("cannot stop RPC server:", err.Error())
		}
	}

	s.ram.Close()
	s.pieceCache.Close()
	s.metrics.Close()
	return s.db.Close()
}

func (s *Session) ListTorrents() []*Torrent {
	s.mTorrents.RLock()
	defer s.mTorrents.RUnlock()
	torrents := make([]*Torrent, 0, len(s.torrents))
	for _, t := range s.torrents {
		torrents = append(torrents, t)
	}
	return torrents
}

func (s *Session) getPort() (int, error) {
	s.mPorts.Lock()
	defer s.mPorts.Unlock()
	for p := range s.availablePorts {
		delete(s.availablePorts, p)
		return p, nil
	}
	return 0, errors.New("no free port")
}

func (s *Session) releasePort(port int) {
	s.mPorts.Lock()
	defer s.mPorts.Unlock()
	s.availablePorts[port] = struct{}{}
}

func (s *Session) GetTorrent(id string) *Torrent {
	s.mTorrents.RLock()
	defer s.mTorrents.RUnlock()
	return s.torrents[id]
}

func (s *Session) RemoveTorrent(id string) error {
	t, err := s.removeTorrentFromClient(id)
	if t != nil {
		go s.stopAndRemoveData(t)
	}
	return err
}

func (s *Session) removeTorrentFromClient(id string) (*Torrent, error) {
	s.mTorrents.Lock()
	defer s.mTorrents.Unlock()
	t, ok := s.torrents[id]
	if !ok {
		return nil, nil
	}
	t.torrent.log.Info("removing torrent")
	delete(s.torrents, id)
	delete(s.torrentsByInfoHash, dht.InfoHash(t.torrent.InfoHash()))
	return t, s.db.Update(func(tx *bolt.Tx) error {
		return tx.Bucket(torrentsBucket).DeleteBucket([]byte(id))
	})
}

func (s *Session) stopAndRemoveData(t *Torrent) {
	t.torrent.Close()
	s.releasePort(t.torrent.port)
	dest := t.torrent.storage.(*filestorage.FileStorage).Dest()
	err := os.RemoveAll(dest)
	if err != nil {
		s.log.Errorf("cannot remove torrent data. err: %s dest: %s", err, dest)
	}
}

func (s *Session) StartAll() error {
	err := s.db.Update(func(tx *bolt.Tx) error {
		tb := tx.Bucket(torrentsBucket)
		s.mTorrents.RLock()
		for _, t := range s.torrents {
			b := tb.Bucket([]byte(t.torrent.id))
			_ = b.Put([]byte("started"), []byte("1"))
		}
		defer s.mTorrents.RUnlock()
		return nil
	})
	if err != nil {
		return err
	}
	for _, t := range s.torrents {
		t.torrent.Start()
	}
	return nil
}

func (s *Session) StopAll() error {
	err := s.db.Update(func(tx *bolt.Tx) error {
		tb := tx.Bucket(torrentsBucket)
		s.mTorrents.RLock()
		for _, t := range s.torrents {
			b := tb.Bucket([]byte(t.torrent.id))
			_ = b.Put([]byte("started"), []byte("0"))
		}
		defer s.mTorrents.RUnlock()
		return nil
	})
	if err != nil {
		return err
	}
	for _, t := range s.torrents {
		t.torrent.Stop()
	}
	return nil
}
<|endoftext|>


<commit_before>package main

import (
	"bytes"
	"encoding/json"
	"fmt"
	"io/ioutil"
	"net/http"
	"net/url"
	"os"
	"os/exec"
	"os/signal"
	"strconv"
	"strings"
	"sync"
	"syscall"
	"text/template"
	"time"

	"github.com/citadel/citadel"
	"github.com/citadel/citadel/cluster"
	"github.com/citadel/citadel/scheduler"
	"github.com/ehazlett/interlock"
)

const (
	haproxyTmpl = `# managed by interlock
global
    {{ if .Config.SyslogAddr }}log {{ .Config.SyslogAddr }} local0
    log-send-hostname{{ end }}
    maxconn {{ .Config.MaxConn }}
    pidfile {{ .Config.PidPath }}

defaults
    mode http
    retries 3
    option redispatch
    option httplog
    option dontlognull
    timeout connect {{ .Config.ConnectTimeout }}
    timeout client {{ .Config.ClientTimeout }}
    timeout server {{ .Config.ServerTimeout }}

frontend http-default
    bind *:{{ .Config.Port }}
    {{ if .Config.StatsUser }}stats realm Stats
    stats auth {{ .Config.StatsUser }}:{{ .Config.StatsPassword }}
    stats enable
    stats uri /haproxy?stats{{ end }}
    {{ range $host := .Hosts }}acl is_{{ $host.Name }} hdr_end(host) -i {{ $host.Domain }}
    use_backend {{ $host.Name }} if is_{{ $host.Name }}
    {{ end }}
{{ range $host := .Hosts }}backend {{ $host.Name }}
    balance roundrobin
    option httpclose
    option forwardfor
    {{ range $i,$up := $host.Upstreams }}server {{$host.Name}}_{{$i}} {{$up.Addr}} check
    {{ end }}
{{ end }}`
)

type (
	Manager struct {
		mux      sync.Mutex
		config   *interlock.Config
		engines  []*citadel.Engine
		cluster  *cluster.Cluster
		proxyCmd *exec.Cmd
	}
)

func NewManager(cfg *interlock.Config) (*Manager, error) {
	engines := []*citadel.Engine{}
	for _, e := range cfg.InterlockEngines {
		if err := e.Engine.Connect(nil); err != nil {
			return nil, err
		}
		engines = append(engines, e.Engine)
	}
	m := &Manager{
		config:  cfg,
		engines: engines,
	}
	if err := m.init(); err != nil {
		return nil, err
	}
	return m, nil
}

func (m *Manager) init() error {
	c, err := cluster.New(scheduler.NewResourceManager(), m.engines...)
	if err != nil {
		return err
	}
	m.cluster = c
	return nil
}

func (m *Manager) writeConfig(config *interlock.ProxyConfig) error {
	m.mux.Lock()
	defer m.mux.Unlock()
	f, err := os.OpenFile(m.config.ProxyConfigPath, os.O_WRONLY|os.O_TRUNC, 0664)
	if err != nil {
		if !os.IsNotExist(err) {
			return err
		}
		ff, fErr := os.Create(m.config.ProxyConfigPath)
		defer ff.Close()
		if fErr != nil {
			return fErr
		}
		f = ff
	}
	defer f.Close()
	t := template.New("haproxy")
	tmpl, err := t.Parse(haproxyTmpl)
	if err != nil {
		return err
	}
	var c bytes.Buffer
	if err := tmpl.Execute(&c, config); err != nil {
		return err
	}
	_, fErr := f.Write(c.Bytes())
	if fErr != nil {
		return fErr
	}
	f.Sync()
	return nil
}

func (m *Manager) GenerateProxyConfig(isKillEvent bool) (*interlock.ProxyConfig, error) {
	containers, err := m.cluster.ListContainers()
	if err != nil {
		return nil, err
	}
	var hosts []*interlock.Host
	proxyUpstreams := map[string][]*interlock.Upstream{}
	for _, cnt := range containers {
		if cnt.Image.Domainname == "" {
			continue
		}
		cntId := cnt.ID[:12]
		// load interlock data
		env := cnt.Image.Environment
		interlockData := &interlock.InterlockData{}
		if key, ok := env["INTERLOCK_DATA"]; ok {
			b := bytes.NewBufferString(key)
			if err := json.NewDecoder(b).Decode(&interlockData); err != nil {
				logger.Warnf("%s: unable to parse interlock data: %s", cntId, err)
			}
		}
		hostname := cnt.Image.Hostname
		domain := cnt.Image.Domainname
		if hostname != domain && hostname != "" {
			domain = fmt.Sprintf("%s.%s", hostname, domain)
		}
		hostAddrUrl, err := url.Parse(cnt.Engine.Addr)
		if err != nil {
			logger.Warnf("%s: unable to parse engine addr: %s", cntId, err)
			continue
		}
		host := hostAddrUrl.Host
		hostParts := strings.Split(hostAddrUrl.Host, ":")
		if len(hostParts) != 1 {
			host = hostParts[0]
		}
		if len(cnt.Ports) == 0 {
			logger.Warnf("%s: no ports exposed", cntId)
			continue
		}
		portDef := cnt.Ports[0]
		addr := fmt.Sprintf("%s:%d", host, portDef.Port)
		if interlockData.Port != 0 {
			for _, p := range cnt.Ports {
				if p.ContainerPort == interlockData.Port {
					addr = fmt.Sprintf("%s:%d", host, p.Port)
				}
			}
		}
		up := &interlock.Upstream{
			Addr: addr,
		}
		for _, alias := range interlockData.AliasDomains {
			proxyUpstreams[alias] = append(proxyUpstreams[alias], up)
		}
		proxyUpstreams[domain] = append(proxyUpstreams[domain], up)
		if !isKillEvent && interlockData.Warm {
			logger.Infof("warming %s: %s", cntId, addr)
			http.Get(fmt.Sprintf("http://%s", addr))
		}

	}
	for k, v := range proxyUpstreams {
		name := strings.Replace(k, ".", "_", -1)
		host := &interlock.Host{
			Name:      name,
			Domain:    k,
			Upstreams: v,
		}
		logger.Infof("adding host name=%s domain=%s", host.Name, host.Domain)
		hosts = append(hosts, host)
	}
	// generate config
	cfg := &interlock.ProxyConfig{
		Hosts:  hosts,
		Config: m.config,
	}
	return cfg, nil
}

func (m *Manager) UpdateConfig(e *citadel.Event) error {
	isKillEvent := false
	if e != nil && e.Type == "kill" {
		isKillEvent = true
	}
	cfg, err := m.GenerateProxyConfig(isKillEvent)
	if err != nil {
		return err
	}
	if err := m.writeConfig(cfg); err != nil {
		return err
	}
	return nil
}

func (m *Manager) getProxyPid() (int, error) {
	f, err := ioutil.ReadFile(m.config.PidPath)
	if err != nil {
		return -1, err
	}
	buf := bytes.NewBuffer(f)
	p := buf.String()
	p = strings.TrimSpace(p)
	pid, err := strconv.Atoi(p)
	if err != nil {
		return -1, err
	}
	return pid, nil
}

func (m *Manager) Reload() error {
	args := []string{"-D", "-f", m.config.ProxyConfigPath, "-p", m.config.PidPath, "-sf"}
	if m.proxyCmd != nil {
		p, err := m.getProxyPid()
		if err != nil {
			logger.Error(err)
		}
		pid := strconv.Itoa(p)
		args = append(args, pid)
	}
	//cmd := exec.Command("haproxy", "-f", m.config.ProxyConfigPath, "-p", "/var/run/haproxy.pid", pidKill)
	cmd := exec.Command("haproxy", args...)
	if err := cmd.Run(); err != nil {
		return err
	}
	// kill old process
	//if m.proxyCmd != nil {
	//	syscall.Kill(m.proxyCmd.Process.Pid, syscall.SIGKILL)
	//}
	m.proxyCmd = cmd
	logger.Info("reloaded proxy")
	return nil
}

func (m *Manager) Run() error {
	if err := m.UpdateConfig(nil); err != nil {
		return err
	}
	if err := m.cluster.Events(&EventHandler{Manager: m}); err != nil {
		return err
	}
	m.Reload()
	ch := make(chan os.Signal, 1)
	signal.Notify(ch, os.Interrupt)
	signal.Notify(ch, syscall.SIGTERM)
	go func() {
		<-ch
		if m.proxyCmd != nil {
			pid, err := m.getProxyPid()
			if err != nil {
				logger.Fatal(err)
			}
			syscall.Kill(pid, syscall.SIGTERM)
		}
		os.Exit(1)
	}()

	for {
		time.Sleep(1 * time.Second)
	}
}
<commit_msg>logging updates<commit_after>package main

import (
	"bytes"
	"encoding/json"
	"fmt"
	"io/ioutil"
	"net/http"
	"net/url"
	"os"
	"os/exec"
	"os/signal"
	"strconv"
	"strings"
	"sync"
	"syscall"
	"text/template"
	"time"

	"github.com/citadel/citadel"
	"github.com/citadel/citadel/cluster"
	"github.com/citadel/citadel/scheduler"
	"github.com/ehazlett/interlock"
)

const (
	haproxyTmpl = `# managed by interlock
global
    {{ if .Config.SyslogAddr }}log {{ .Config.SyslogAddr }} local0
    log-send-hostname{{ end }}
    maxconn {{ .Config.MaxConn }}
    pidfile {{ .Config.PidPath }}

defaults
    mode http
    retries 3
    option redispatch
    option httplog
    option dontlognull
    timeout connect {{ .Config.ConnectTimeout }}
    timeout client {{ .Config.ClientTimeout }}
    timeout server {{ .Config.ServerTimeout }}

frontend http-default
    bind *:{{ .Config.Port }}
    {{ if .Config.StatsUser }}stats realm Stats
    stats auth {{ .Config.StatsUser }}:{{ .Config.StatsPassword }}
    stats enable
    stats uri /haproxy?stats{{ end }}
    {{ range $host := .Hosts }}acl is_{{ $host.Name }} hdr_end(host) -i {{ $host.Domain }}
    use_backend {{ $host.Name }} if is_{{ $host.Name }}
    {{ end }}
{{ range $host := .Hosts }}backend {{ $host.Name }}
    balance roundrobin
    option httpclose
    option forwardfor
    {{ range $i,$up := $host.Upstreams }}server {{$host.Name}}_{{$i}} {{$up.Addr}} check
    {{ end }}
{{ end }}`
)

type (
	Manager struct {
		mux      sync.Mutex
		config   *interlock.Config
		engines  []*citadel.Engine
		cluster  *cluster.Cluster
		proxyCmd *exec.Cmd
	}
)

func NewManager(cfg *interlock.Config) (*Manager, error) {
	engines := []*citadel.Engine{}
	for _, e := range cfg.InterlockEngines {
		if err := e.Engine.Connect(nil); err != nil {
			return nil, err
		}
		engines = append(engines, e.Engine)
	}
	m := &Manager{
		config:  cfg,
		engines: engines,
	}
	if err := m.init(); err != nil {
		return nil, err
	}
	return m, nil
}

func (m *Manager) init() error {
	c, err := cluster.New(scheduler.NewResourceManager(), m.engines...)
	if err != nil {
		return err
	}
	m.cluster = c
	return nil
}

func (m *Manager) writeConfig(config *interlock.ProxyConfig) error {
	m.mux.Lock()
	defer m.mux.Unlock()
	f, err := os.OpenFile(m.config.ProxyConfigPath, os.O_WRONLY|os.O_TRUNC, 0664)
	if err != nil {
		if !os.IsNotExist(err) {
			return err
		}
		ff, fErr := os.Create(m.config.ProxyConfigPath)
		defer ff.Close()
		if fErr != nil {
			return fErr
		}
		f = ff
	}
	defer f.Close()
	t := template.New("haproxy")
	tmpl, err := t.Parse(haproxyTmpl)
	if err != nil {
		return err
	}
	var c bytes.Buffer
	if err := tmpl.Execute(&c, config); err != nil {
		return err
	}
	_, fErr := f.Write(c.Bytes())
	if fErr != nil {
		return fErr
	}
	f.Sync()
	return nil
}

func (m *Manager) GenerateProxyConfig(isKillEvent bool) (*interlock.ProxyConfig, error) {
	containers, err := m.cluster.ListContainers()
	if err != nil {
		return nil, err
	}
	var hosts []*interlock.Host
	proxyUpstreams := map[string][]*interlock.Upstream{}
	for _, cnt := range containers {
		if cnt.Image.Domainname == "" {
			continue
		}
		cntId := cnt.ID[:12]
		// load interlock data
		env := cnt.Image.Environment
		interlockData := &interlock.InterlockData{}
		if key, ok := env["INTERLOCK_DATA"]; ok {
			b := bytes.NewBufferString(key)
			if err := json.NewDecoder(b).Decode(&interlockData); err != nil {
				logger.Warnf("%s: unable to parse interlock data: %s", cntId, err)
			}
		}
		hostname := cnt.Image.Hostname
		domain := cnt.Image.Domainname
		if hostname != domain && hostname != "" {
			domain = fmt.Sprintf("%s.%s", hostname, domain)
		}
		hostAddrUrl, err := url.Parse(cnt.Engine.Addr)
		if err != nil {
			logger.Warnf("%s: unable to parse engine addr: %s", cntId, err)
			continue
		}
		host := hostAddrUrl.Host
		hostParts := strings.Split(hostAddrUrl.Host, ":")
		if len(hostParts) != 1 {
			host = hostParts[0]
		}
		if len(cnt.Ports) == 0 {
			logger.Warnf("%s: no ports exposed", cntId)
			continue
		}
		portDef := cnt.Ports[0]
		addr := fmt.Sprintf("%s:%d", host, portDef.Port)
		if interlockData.Port != 0 {
			for _, p := range cnt.Ports {
				if p.ContainerPort == interlockData.Port {
					addr = fmt.Sprintf("%s:%d", host, p.Port)
				}
			}
		}
		up := &interlock.Upstream{
			Addr: addr,
		}
		for _, alias := range interlockData.AliasDomains {
			logger.Infof("adding alias %s for %s", alias, cntId)
			proxyUpstreams[alias] = append(proxyUpstreams[alias], up)
		}
		proxyUpstreams[domain] = append(proxyUpstreams[domain], up)
		if !isKillEvent && interlockData.Warm {
			logger.Infof("warming %s: %s", cntId, addr)
			http.Get(fmt.Sprintf("http://%s", addr))
		}

	}
	for k, v := range proxyUpstreams {
		name := strings.Replace(k, ".", "_", -1)
		host := &interlock.Host{
			Name:      name,
			Domain:    k,
			Upstreams: v,
		}
		logger.Infof("adding host name=%s domain=%s", host.Name, host.Domain)
		hosts = append(hosts, host)
	}
	// generate config
	cfg := &interlock.ProxyConfig{
		Hosts:  hosts,
		Config: m.config,
	}
	return cfg, nil
}

func (m *Manager) UpdateConfig(e *citadel.Event) error {
	isKillEvent := false
	if e != nil && e.Type == "kill" {
		isKillEvent = true
	}
	cfg, err := m.GenerateProxyConfig(isKillEvent)
	if err != nil {
		return err
	}
	if err := m.writeConfig(cfg); err != nil {
		return err
	}
	return nil
}

func (m *Manager) getProxyPid() (int, error) {
	f, err := ioutil.ReadFile(m.config.PidPath)
	if err != nil {
		return -1, err
	}
	buf := bytes.NewBuffer(f)
	p := buf.String()
	p = strings.TrimSpace(p)
	pid, err := strconv.Atoi(p)
	if err != nil {
		return -1, err
	}
	return pid, nil
}

func (m *Manager) Reload() error {
	args := []string{"-D", "-f", m.config.ProxyConfigPath, "-p", m.config.PidPath, "-sf"}
	if m.proxyCmd != nil {
		p, err := m.getProxyPid()
		if err != nil {
			logger.Error(err)
		}
		pid := strconv.Itoa(p)
		args = append(args, pid)
	}
	//cmd := exec.Command("haproxy", "-f", m.config.ProxyConfigPath, "-p", "/var/run/haproxy.pid", pidKill)
	cmd := exec.Command("haproxy", args...)
	if err := cmd.Run(); err != nil {
		return err
	}
	// kill old process
	//if m.proxyCmd != nil {
	//	syscall.Kill(m.proxyCmd.Process.Pid, syscall.SIGKILL)
	//}
	m.proxyCmd = cmd
	logger.Info("reloaded proxy")
	return nil
}

func (m *Manager) Run() error {
	if err := m.UpdateConfig(nil); err != nil {
		return err
	}
	if err := m.cluster.Events(&EventHandler{Manager: m}); err != nil {
		return err
	}
	m.Reload()
	ch := make(chan os.Signal, 1)
	signal.Notify(ch, os.Interrupt)
	signal.Notify(ch, syscall.SIGTERM)
	go func() {
		<-ch
		if m.proxyCmd != nil {
			pid, err := m.getProxyPid()
			if err != nil {
				logger.Fatal(err)
			}
			syscall.Kill(pid, syscall.SIGTERM)
		}
		os.Exit(1)
	}()

	for {
		time.Sleep(1 * time.Second)
	}
}
<|endoftext|>


<commit_before>package flake

import (
	"encoding/binary"
	"os"
)

const (
	// DefaultSequenceBits is the default # of bits used for sequence #'s
	DefaultSequenceBits uint64 = 16
)

//  ---------------------------------------------------------------------------
//  Layout - Big Endian
//  ---------------------------------------------------------------------------
//
//  [0:6]   48 bits | Upper 48 bits of timestamp (milliseconds since the epoch)
//  [6:8]   16 bits | a per-interval sequence # (interval == 1 millisecond)
//  [8:14]  48 bits | a hardware id
//  [14:16] 16 bits | process ID
//
//  ---------------------------------------------------------------------------
//  | 0 | 1 | 2 | 3 | 4 | 5 |  6  |  7  |  8  | 9 | A | B | C | D |  E  |  F  |
//  ---------------------------------------------------------------------------
//  |           48 bits     |  16 bits  |         48 bits         |  16 bits  |
//  ---------------------------------------------------------------------------
//  |          timestamp    |  sequence |        HardwareID       | ProcessID |
//  ---------------------------------------------------------------------------
//  Notes
//  ---------------------------------------------------------------------------
//  The time bits are the most significant bits because they have the primary
//  impact on the sort order of ids. The sequence # is next most significant
//  as it is the tie-breaker when the time portions are equivalent.
//
//  Note that the lower 64 bits are basically random and not specifically
//  useful for ordering, although they play their part when the upper 64-bits
//  are equivalent between two ids. Again, the ordering outcome in this
//  situation is somewhat random, but generally somewhat repeatable (hardware
//  id should be consistent and stable a vast majority of the time).
//  ---------------------------------------------------------------------------

type overtFlakeIDSynthesizer struct {
	epoch        int64
	sequenceBits uint64
	sequenceMask uint64
	hardwareID   HardwareID
	processID    int
	machineID    uint64
}

// NewOvertFlakeIDSynthesizer creates an instance of generator (which implements Generator) and
// allows the # of sequence bits to be specified (16 is standard)
func NewOvertFlakeIDSynthesizer(epoch int64, sequenceBits uint64, hardwareID HardwareID, processID int) OvertFlakeIDGenerator {
	// binary.BigEndian.Uint64 won't work on a []byte < len(8) so we need to
	// copy our 6-byte hardwareID into the most-signficant bits
	tempBytes := make([]byte, 8)
	copy(tempBytes[0:6], hardwareID[0:6])

	return &overtFlakeIDSynthesizer{
		epoch:        epoch,
		sequenceBits: sequenceBits,
		sequenceMask: uint64(int64(-1) ^ (int64(-1) << sequenceBits)),
		hardwareID:   hardwareID,
		processID:    processID & 0xFFFF,
		machineID:    binary.BigEndian.Uint64(tempBytes) | uint64(processID&0xFFFF),
	}
}

// NewOvertFlakeGeneratorWithBits creates an instance of generator (which implements Generator.) this
// "constructor" allows the # of sequence bits to be specified
func NewOvertFlakeGeneratorWithBits(epoch int64, hardwareID HardwareID, processID int, waitForTime int64, seqBits uint64) Generator {
	// binary.BigEndian.Uint64 won't work on a []byte < len(8) so we need to
	// copy our 6-byte hardwareID into the most-signficant bits
	tempBytes := make([]byte, 8)
	copy(tempBytes[0:6], hardwareID[0:6])

	return &generator{
		idGen:    NewOvertFlakeIDSynthesizer(epoch, seqBits, hardwareID, processID),
		lastTime: waitForTime,
	}
}

// NewOvertFlakeGenerator creates an instance of generator which implements Generator
func NewOvertFlakeGenerator(epoch int64, hardwareID HardwareID, processID int, waitForTime int64) Generator {
	return NewOvertFlakeGeneratorWithBits(epoch, hardwareID, processID, waitForTime, DefaultSequenceBits)
}

// NewOvertoneEpochGenerator creates an instance of generator using the Overtone Epoch
func NewOvertoneEpochGenerator(hardwareID HardwareID) Generator {
	return NewOvertFlakeGeneratorWithBits(OvertoneEpochMs, hardwareID, os.Getpid(), 0, DefaultSequenceBits)
}

func (ofid *overtFlakeIDSynthesizer) HardwareID() HardwareID {
	return ofid.hardwareID
}

func (ofid *overtFlakeIDSynthesizer) ProcessID() int {
	return ofid.processID
}

func (ofid *overtFlakeIDSynthesizer) IDSize() int {
	return OvertFlakeIDLength
}

func (ofid *overtFlakeIDSynthesizer) SequenceBitCount() uint64 {
	return ofid.sequenceBits
}

func (ofid *overtFlakeIDSynthesizer) SequenceBitMask() uint64 {
	return ofid.sequenceMask
}

func (ofid *overtFlakeIDSynthesizer) MaxSequenceNumber() uint64 {
	return ofid.sequenceMask
}

func (ofid *overtFlakeIDSynthesizer) Epoch() int64 {
	return ofid.epoch
}

func (ofid *overtFlakeIDSynthesizer) SynthesizeID(buffer []byte, index int, time int64, sequence uint64) int {
	// time is Unix Epoch (note that this is inefficient in that delta has to be calculated for
	// each id, when the generator could do the calculation + shift 1 time per allocate)
	delta := uint64(time - ofid.epoch)

	// upper 32 are time | sequence
	var upper = (delta << ofid.sequenceBits) | (sequence & ofid.sequenceMask)

	// Write the id
	binary.BigEndian.PutUint64(buffer[index:index+8], upper)
	binary.BigEndian.PutUint64(buffer[index+8:index+16], ofid.machineID)

	// return the length of the id
	return OvertFlakeIDLength
}
<commit_msg>added some dox<commit_after>package flake

import (
	"encoding/binary"
	"os"
)

const (
	// DefaultSequenceBits is the default # of bits used for sequence #'s
	DefaultSequenceBits uint64 = 16
)

//  ---------------------------------------------------------------------------
//  Layout - Big Endian
//  ---------------------------------------------------------------------------
//
//  [0:6]   48 bits | Upper 48 bits of timestamp (milliseconds since the epoch)
//  [6:8]   16 bits | a per-interval sequence # (interval == 1 millisecond)
//  [8:14]  48 bits | a hardware id
//  [14:16] 16 bits | process ID
//
//  ---------------------------------------------------------------------------
//  | 0 | 1 | 2 | 3 | 4 | 5 |  6  |  7  |  8  | 9 | A | B | C | D |  E  |  F  |
//  ---------------------------------------------------------------------------
//  |           48 bits     |  16 bits  |         48 bits         |  16 bits  |
//  ---------------------------------------------------------------------------
//  |          timestamp    |  sequence |        HardwareID       | ProcessID |
//  ---------------------------------------------------------------------------
//  Notes
//  ---------------------------------------------------------------------------
//  The time bits are the most significant bits because they have the primary
//  impact on the sort order of ids. The sequence # is next most significant
//  as it is the tie-breaker when the time portions are equivalent.
//
//  Note that the lower 64 bits are basically random and not specifically
//  useful for ordering, although they play their part when the upper 64-bits
//  are equivalent between two ids. Again, the ordering outcome in this
//  situation is somewhat random, but generally somewhat repeatable (hardware
//  id should be consistent and stable a vast majority of the time).
//  ---------------------------------------------------------------------------

type overtFlakeIDSynthesizer struct {
	epoch        int64
	sequenceBits uint64
	sequenceMask uint64
	hardwareID   HardwareID
	processID    int
	machineID    uint64
}

// NewOvertFlakeIDSynthesizer creates an instance of generator (which implements Generator) and
// allows the # of sequence bits to be specified (16 is standard)
//
// Notes
//
// Setting a value of sequenceBits > 22 will result in unacceptable time truncation
func NewOvertFlakeIDSynthesizer(epoch int64, sequenceBits uint64, hardwareID HardwareID, processID int) OvertFlakeIDGenerator {
	// binary.BigEndian.Uint64 won't work on a []byte < len(8) so we need to
	// copy our 6-byte hardwareID into the most-signficant bits
	tempBytes := make([]byte, 8)
	copy(tempBytes[0:6], hardwareID[0:6])

	return &overtFlakeIDSynthesizer{
		epoch:        epoch,
		sequenceBits: sequenceBits,
		sequenceMask: uint64(int64(-1) ^ (int64(-1) << sequenceBits)),
		hardwareID:   hardwareID,
		processID:    processID & 0xFFFF,
		machineID:    binary.BigEndian.Uint64(tempBytes) | uint64(processID&0xFFFF),
	}
}

// NewOvertFlakeGeneratorWithBits creates an instance of generator (which implements Generator.) this
// "constructor" allows the # of sequence bits to be specified
//
// Notes
//
// Setting a value of seqBits > 22 will result in unacceptable time truncation
func NewOvertFlakeGeneratorWithBits(epoch int64, hardwareID HardwareID, processID int, waitForTime int64, seqBits uint64) Generator {
	// binary.BigEndian.Uint64 won't work on a []byte < len(8) so we need to
	// copy our 6-byte hardwareID into the most-signficant bits
	tempBytes := make([]byte, 8)
	copy(tempBytes[0:6], hardwareID[0:6])

	return &generator{
		idGen:    NewOvertFlakeIDSynthesizer(epoch, seqBits, hardwareID, processID),
		lastTime: waitForTime,
	}
}

// NewOvertFlakeGenerator creates an instance of generator which implements Generator
func NewOvertFlakeGenerator(epoch int64, hardwareID HardwareID, processID int, waitForTime int64) Generator {
	return NewOvertFlakeGeneratorWithBits(epoch, hardwareID, processID, waitForTime, DefaultSequenceBits)
}

// NewOvertoneEpochGenerator creates an instance of generator using the Overtone Epoch
func NewOvertoneEpochGenerator(hardwareID HardwareID) Generator {
	return NewOvertFlakeGeneratorWithBits(OvertoneEpochMs, hardwareID, os.Getpid(), 0, DefaultSequenceBits)
}

func (ofid *overtFlakeIDSynthesizer) HardwareID() HardwareID {
	return ofid.hardwareID
}

func (ofid *overtFlakeIDSynthesizer) ProcessID() int {
	return ofid.processID
}

func (ofid *overtFlakeIDSynthesizer) IDSize() int {
	return OvertFlakeIDLength
}

func (ofid *overtFlakeIDSynthesizer) SequenceBitCount() uint64 {
	return ofid.sequenceBits
}

func (ofid *overtFlakeIDSynthesizer) SequenceBitMask() uint64 {
	return ofid.sequenceMask
}

func (ofid *overtFlakeIDSynthesizer) MaxSequenceNumber() uint64 {
	return ofid.sequenceMask
}

func (ofid *overtFlakeIDSynthesizer) Epoch() int64 {
	return ofid.epoch
}

func (ofid *overtFlakeIDSynthesizer) SynthesizeID(buffer []byte, index int, time int64, sequence uint64) int {
	// time is Unix Epoch (note that this is inefficient in that delta has to be calculated for
	// each id, when the generator could do the calculation + shift 1 time per allocate)
	delta := uint64(time - ofid.epoch)

	// upper 32 are time | sequence
	var upper = (delta << ofid.sequenceBits) | (sequence & ofid.sequenceMask)

	// Write the id
	binary.BigEndian.PutUint64(buffer[index:index+8], upper)
	binary.BigEndian.PutUint64(buffer[index+8:index+16], ofid.machineID)

	// return the length of the id
	return OvertFlakeIDLength
}
<|endoftext|>


<commit_before>package scrabble_score

import "testing"

var tests = []struct {
	input    string
	expected int
}{
	{"", 0},
	{" \t\n", 0},
	{"a", 1},
	{"f", 4},
	{"street", 6},
	{"quirky", 22},
	{"MULTIBILLIONAIRE", 20},
	{"alacrity", 13},
}

func TestScore(t *testing.T) {
	for _, test := range tests {
		if actual := Score(test.input); actual != test.expected {
			t.Errorf("Score(%q) expected %d, Actual %d", test.input, test.expected, actual)
		}
	}
}

func BenchmarkScore(b *testing.B) {
	for i := 0; i < b.N; i++ {
		for _, test := range tests {
			Score(test.input)
		}
	}
}
<commit_msg>Switch MultiBilliionaire for a vaid word<commit_after>package scrabble_score

import "testing"

var tests = []struct {
	input    string
	expected int
}{
	{"", 0},
	{" \t\n", 0},
	{"a", 1},
	{"f", 4},
	{"street", 6},
	{"quirky", 22},
	{"oxyphenbutazone", 41},
	{"alacrity", 13},
}

func TestScore(t *testing.T) {
	for _, test := range tests {
		if actual := Score(test.input); actual != test.expected {
			t.Errorf("Score(%q) expected %d, Actual %d", test.input, test.expected, actual)
		}
	}
}

func BenchmarkScore(b *testing.B) {
	for i := 0; i < b.N; i++ {
		for _, test := range tests {
			Score(test.input)
		}
	}
}
<|endoftext|>


<commit_before>package main

import (
	"fmt"
	"os"
	"sort"

	"github.com/coreos/fleet/job"
	"github.com/coreos/fleet/unit"
)

var cmdListUnits = &Command{
	Name:    "list-units",
	Summary: "Enumerate units loaded in the cluster",
	Usage:   "[--no-legend] [-l|--full]",
	Description: `Lists all units submitted or started on the cluster.

For easily parsable output, you can remove the column headers:
	fleetctl list-units --no-legend

Output the list without ellipses:
	fleetctl list-units --full`,
	Run: runListUnits,
}

func init() {
	cmdListUnits.Flags.BoolVar(&sharedFlags.Full, "full", false, "Do not ellipsize fields on output")
	cmdListUnits.Flags.BoolVar(&sharedFlags.Full, "l", false, "Shorthand for --full")
	cmdListUnits.Flags.BoolVar(&sharedFlags.NoLegend, "no-legend", false, "Do not print a legend (column headers)")
}

func runListUnits(args []string) (exit int) {
	jobs, sortable, err := findAllUnits()

	if err != nil {
		fmt.Fprintf(os.Stderr, "Error retrieving list of units from repository: %v\n", err)
		return -1
	}

	if !sharedFlags.NoLegend {
		fmt.Fprintln(out, "UNIT\tSTATE\tLOAD\tACTIVE\tSUB\tDESC\tMACHINE")
	}

	for _, name := range sortable {
		j := jobs[name]
		printUnitState(name, j.Unit.Description(), j.State, j.UnitState, sharedFlags.Full)
	}

	out.Flush()
	return
}

func findAllUnits() (jobs map[string]job.Job, sortable sort.StringSlice, err error) {
	jobs = make(map[string]job.Job, 0)
	sortable = make(sort.StringSlice, 0)

	jj, err := registryCtl.GetAllJobs()
	if err != nil {
		return
	}

	for _, j := range jj {
		jobs[j.Name] = j
		sortable = append(sortable, j.Name)
	}

	sortable.Sort()

	return
}

func printUnitState(name, description string, js *job.JobState, us *unit.UnitState, full bool) {
	jobState := "-"
	loadState := "-"
	activeState := "-"
	subState := "-"
	mach := "-"

	if description == "" {
		description = "-"
	}

	if js != nil {
		jobState = string(*js)
	}

	if us != nil {
		loadState = us.LoadState
		activeState = us.ActiveState
		subState = us.SubState

		if us.MachineState != nil {
			mach = machineFullLegend(*us.MachineState, full)
		}
	}

	fmt.Fprintf(out, "%s\t%s\t%s\t%s\t%s\t%s\t%s\n", name, jobState, loadState, activeState, subState, description, mach)
}
<commit_msg>s/-1/1/<commit_after>package main

import (
	"fmt"
	"os"
	"sort"

	"github.com/coreos/fleet/job"
	"github.com/coreos/fleet/unit"
)

var cmdListUnits = &Command{
	Name:    "list-units",
	Summary: "Enumerate units loaded in the cluster",
	Usage:   "[--no-legend] [-l|--full]",
	Description: `Lists all units submitted or started on the cluster.

For easily parsable output, you can remove the column headers:
	fleetctl list-units --no-legend

Output the list without ellipses:
	fleetctl list-units --full`,
	Run: runListUnits,
}

func init() {
	cmdListUnits.Flags.BoolVar(&sharedFlags.Full, "full", false, "Do not ellipsize fields on output")
	cmdListUnits.Flags.BoolVar(&sharedFlags.Full, "l", false, "Shorthand for --full")
	cmdListUnits.Flags.BoolVar(&sharedFlags.NoLegend, "no-legend", false, "Do not print a legend (column headers)")
}

func runListUnits(args []string) (exit int) {
	jobs, sortable, err := findAllUnits()

	if err != nil {
		fmt.Fprintf(os.Stderr, "Error retrieving list of units from repository: %v\n", err)
		return 1
	}

	if !sharedFlags.NoLegend {
		fmt.Fprintln(out, "UNIT\tSTATE\tLOAD\tACTIVE\tSUB\tDESC\tMACHINE")
	}

	for _, name := range sortable {
		j := jobs[name]
		printUnitState(name, j.Unit.Description(), j.State, j.UnitState, sharedFlags.Full)
	}

	out.Flush()
	return
}

func findAllUnits() (jobs map[string]job.Job, sortable sort.StringSlice, err error) {
	jobs = make(map[string]job.Job, 0)
	sortable = make(sort.StringSlice, 0)

	jj, err := registryCtl.GetAllJobs()
	if err != nil {
		return
	}

	for _, j := range jj {
		jobs[j.Name] = j
		sortable = append(sortable, j.Name)
	}

	sortable.Sort()

	return
}

func printUnitState(name, description string, js *job.JobState, us *unit.UnitState, full bool) {
	jobState := "-"
	loadState := "-"
	activeState := "-"
	subState := "-"
	mach := "-"

	if description == "" {
		description = "-"
	}

	if js != nil {
		jobState = string(*js)
	}

	if us != nil {
		loadState = us.LoadState
		activeState = us.ActiveState
		subState = us.SubState

		if us.MachineState != nil {
			mach = machineFullLegend(*us.MachineState, full)
		}
	}

	fmt.Fprintf(out, "%s\t%s\t%s\t%s\t%s\t%s\t%s\n", name, jobState, loadState, activeState, subState, description, mach)
}
<|endoftext|>


<commit_before>package library

import (
	"github.com/nytlabs/streamtools/st/blocks" // blocks
)

type Join struct {
	blocks.Block
	inA   chan interface{}
	inB   chan interface{}
	clear chan interface{}
	out   chan interface{}
	quit  chan interface{}
}

func NewJoin() blocks.BlockInterface {
	return &Join{}
}

func (b *Join) Setup() {
	b.Kind = "Join"
	b.inA = b.InRoute("inA")
	b.inB = b.InRoute("inB")
	b.clear = b.InRoute("clear")
	b.quit = b.Quit()
	b.out = b.Broadcast()
}

func (b *Join) Run() {
	A := make(chan interface{}, 1000)
	B := make(chan interface{}, 1000)
	for {
		select {
		case <-b.quit:
			return
		case msg := <-b.inA:
			A <- msg
		case msg := <-b.inB:
			B <- msg
		case <-b.clear:
			go func() {
				for {
					select {
					case <-A:
					case <-B:
					default:
						return
					}
				}
			}()
		}
		for len(A) > 0 && len(B) > 0 {
			b.out <- map[string]interface{}{
				"A": <-A,
				"B": <-B,
			}
		}
	}
}
<commit_msg>removed superfluous go routine<commit_after>package library

import (
	"github.com/nytlabs/streamtools/st/blocks" // blocks
)

type Join struct {
	blocks.Block
	inA   chan interface{}
	inB   chan interface{}
	clear chan interface{}
	out   chan interface{}
	quit  chan interface{}
}

func NewJoin() blocks.BlockInterface {
	return &Join{}
}

func (b *Join) Setup() {
	b.Kind = "Join"
	b.inA = b.InRoute("inA")
	b.inB = b.InRoute("inB")
	b.clear = b.InRoute("clear")
	b.quit = b.Quit()
	b.out = b.Broadcast()
}

func (b *Join) Run() {
	A := make(chan interface{}, 1000)
	B := make(chan interface{}, 1000)
	for {
		select {
		case <-b.quit:
			return
		case msg := <-b.inA:
			A <- msg
		case msg := <-b.inB:
			B <- msg
		case <-b.clear:
			for {
				select {
				case <-A:
				case <-B:
				default:
					break
				}
			}
		}
		for len(A) > 0 && len(B) > 0 {
			b.out <- map[string]interface{}{
				"A": <-A,
				"B": <-B,
			}
		}
	}
}
<|endoftext|>


<commit_before>// Copyright 2017 Xiaomi, 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 service

import (
	"fmt"
	"log"
	"time"
)

// FIX ME: too many JOIN
func GetHostsFromGroup(grpName string) map[string]int {
	hosts := make(map[string]int)

	now := time.Now().Unix()
	q := fmt.Sprintf("SELECT host.id, host.hostname FROM grp_host AS gh "+
		" INNER JOIN host ON host.id=gh.host_id AND (host.maintain_begin > %d OR host.maintain_end < %d)"+
		" INNER JOIN grp ON grp.id=gh.grp_id AND grp.grp_name='%s'", now, now, grpName)

	dbConn, err := GetDbConn("nodata.host")
	if err != nil {
		log.Println("db.get_conn error, host", err)
		return hosts
	}

	rows, err := dbConn.Query(q)
	if err != nil {
		log.Println("[ERROR]", err)
		return hosts
	}

	defer rows.Close()
	for rows.Next() {
		hid := -1
		hostname := ""
		err = rows.Scan(&hid, &hostname)
		if err != nil {
			log.Println("[ERROR]", err)
			continue
		}
		if hid < 0 || hostname == "" {
			continue
		}

		hosts[hostname] = hid
	}

	return hosts
}
<commit_msg>fix-nodata-sql-injection (#954)<commit_after>// Copyright 2017 Xiaomi, 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 service

import (
	"log"
	"time"
)

// FIX ME: too many JOIN
func GetHostsFromGroup(grpName string) map[string]int {
	hosts := make(map[string]int)

	now := time.Now().Unix()
	sqlStatement := "SELECT host.id, host.hostname FROM grp_host AS gh " +
		" INNER JOIN host ON host.id=gh.host_id AND (host.maintain_begin > ? OR host.maintain_end < ?)" +
		" INNER JOIN grp ON grp.id=gh.grp_id AND grp.grp_name=?"

	dbConn, err := GetDbConn("nodata.host")
	if err != nil {
		log.Println("db.get_conn error, host", err)
		return hosts
	}

	rows, err := dbConn.Query(sqlStatement, now, now, grpName)
	if err != nil {
		log.Println("[ERROR]", err)
		return hosts
	}

	defer rows.Close()
	for rows.Next() {
		hid := -1
		hostname := ""
		err = rows.Scan(&hid, &hostname)
		if err != nil {
			log.Println("[ERROR]", err)
			continue
		}
		if hid < 0 || hostname == "" {
			continue
		}

		hosts[hostname] = hid
	}

	return hosts
}
<|endoftext|>


<commit_before>// +build linux

package overlay

import (
	"bufio"
	"fmt"
	"io/ioutil"
	"os"
	"os/exec"
	"path"
	"sync"
	"syscall"

	log "github.com/Sirupsen/logrus"
	"github.com/docker/docker/daemon/graphdriver"
	"github.com/docker/docker/pkg/archive"
	"github.com/docker/docker/pkg/chrootarchive"
	"github.com/docker/libcontainer/label"
)

// This is a small wrapper over the NaiveDiffWriter that lets us have a custom
// implementation of ApplyDiff()

var (
	ErrApplyDiffFallback = fmt.Errorf("Fall back to normal ApplyDiff")
)

type ApplyDiffProtoDriver interface {
	graphdriver.ProtoDriver
	ApplyDiff(id, parent string, diff archive.ArchiveReader) (size int64, err error)
}

type naiveDiffDriverWithApply struct {
	graphdriver.Driver
	applyDiff ApplyDiffProtoDriver
}

func NaiveDiffDriverWithApply(driver ApplyDiffProtoDriver) graphdriver.Driver {
	return &naiveDiffDriverWithApply{
		Driver:    graphdriver.NaiveDiffDriver(driver),
		applyDiff: driver,
	}
}

func (d *naiveDiffDriverWithApply) ApplyDiff(id, parent string, diff archive.ArchiveReader) (int64, error) {
	b, err := d.applyDiff.ApplyDiff(id, parent, diff)
	if err == ErrApplyDiffFallback {
		return d.Driver.ApplyDiff(id, parent, diff)
	}
	return b, err
}

// This backend uses the overlay union filesystem for containers
// plus hard link file sharing for images.

// Each container/image can have a "root" subdirectory which is a plain
// filesystem hierarchy, or they can use overlay.

// If they use overlay there is a "upper" directory and a "lower-id"
// file, as well as "merged" and "work" directories. The "upper"
// directory has the upper layer of the overlay, and "lower-id" contains
// the id of the parent whose "root" directory shall be used as the lower
// layer in the overlay. The overlay itself is mounted in the "merged"
// directory, and the "work" dir is needed for overlay to work.

// When a overlay layer is created there are two cases, either the
// parent has a "root" dir, then we start out with a empty "upper"
// directory overlaid on the parents root. This is typically the
// case with the init layer of a container which is based on an image.
// If there is no "root" in the parent, we inherit the lower-id from
// the parent and start by making a copy if the parents "upper" dir.
// This is typically the case for a container layer which copies
// its parent -init upper layer.

// Additionally we also have a custom implementation of ApplyLayer
// which makes a recursive copy of the parent "root" layer using
// hardlinks to share file data, and then applies the layer on top
// of that. This means all child images share file (but not directory)
// data with the parent.

type ActiveMount struct {
	count   int
	path    string
	mounted bool
}
type Driver struct {
	home       string
	sync.Mutex // Protects concurrent modification to active
	active     map[string]*ActiveMount
}

func init() {
	graphdriver.Register("overlay", Init)
}

func Init(home string, options []string) (graphdriver.Driver, error) {
	if err := supportsOverlay(); err != nil {
		return nil, graphdriver.ErrNotSupported
	}

	// Create the driver home dir
	if err := os.MkdirAll(home, 0755); err != nil && !os.IsExist(err) {
		return nil, err
	}

	d := &Driver{
		home:   home,
		active: make(map[string]*ActiveMount),
	}

	return NaiveDiffDriverWithApply(d), nil
}

func supportsOverlay() error {
	// We can try to modprobe overlay first before looking at
	// proc/filesystems for when overlay is supported
	exec.Command("modprobe", "overlay").Run()

	f, err := os.Open("/proc/filesystems")
	if err != nil {
		return err
	}
	defer f.Close()

	s := bufio.NewScanner(f)
	for s.Scan() {
		if s.Text() == "nodev\toverlay" {
			return nil
		}
	}
	log.Error("'overlay' not found as a supported filesystem on this host. Please ensure kernel is new enough and has overlay support loaded.")
	return graphdriver.ErrNotSupported
}

func (d *Driver) String() string {
	return "overlay"
}

func (d *Driver) Status() [][2]string {
	return nil
}

func (d *Driver) Cleanup() error {
	return nil
}

func (d *Driver) Create(id string, parent string) (retErr error) {
	dir := d.dir(id)
	if err := os.MkdirAll(path.Dir(dir), 0700); err != nil {
		return err
	}
	if err := os.Mkdir(dir, 0700); err != nil {
		return err
	}

	defer func() {
		// Clean up on failure
		if retErr != nil {
			os.RemoveAll(dir)
		}
	}()

	// Toplevel images are just a "root" dir
	if parent == "" {
		if err := os.Mkdir(path.Join(dir, "root"), 0755); err != nil {
			return err
		}
		return nil
	}

	parentDir := d.dir(parent)

	// Ensure parent exists
	if _, err := os.Lstat(parentDir); err != nil {
		return err
	}

	// If parent has a root, just do a overlay to it
	parentRoot := path.Join(parentDir, "root")

	if s, err := os.Lstat(parentRoot); err == nil {
		if err := os.Mkdir(path.Join(dir, "upper"), s.Mode()); err != nil {
			return err
		}
		if err := os.Mkdir(path.Join(dir, "work"), 0700); err != nil {
			return err
		}
		if err := os.Mkdir(path.Join(dir, "merged"), 0700); err != nil {
			return err
		}
		if err := ioutil.WriteFile(path.Join(dir, "lower-id"), []byte(parent), 0666); err != nil {
			return err
		}
		return nil
	}

	// Otherwise, copy the upper and the lower-id from the parent

	lowerId, err := ioutil.ReadFile(path.Join(parentDir, "lower-id"))
	if err != nil {
		return err
	}

	if err := ioutil.WriteFile(path.Join(dir, "lower-id"), lowerId, 0666); err != nil {
		return err
	}

	parentUpperDir := path.Join(parentDir, "upper")
	s, err := os.Lstat(parentUpperDir)
	if err != nil {
		return err
	}

	upperDir := path.Join(dir, "upper")
	if err := os.Mkdir(upperDir, s.Mode()); err != nil {
		return err
	}
	if err := os.Mkdir(path.Join(dir, "work"), 0700); err != nil {
		return err
	}
	if err := os.Mkdir(path.Join(dir, "merged"), 0700); err != nil {
		return err
	}

	return copyDir(parentUpperDir, upperDir, 0)
}

func (d *Driver) dir(id string) string {
	return path.Join(d.home, id)
}

func (d *Driver) Remove(id string) error {
	dir := d.dir(id)
	if _, err := os.Stat(dir); err != nil {
		return err
	}
	return os.RemoveAll(dir)
}

func (d *Driver) Get(id string, mountLabel string) (string, error) {
	// Protect the d.active from concurrent access
	d.Lock()
	defer d.Unlock()

	mount := d.active[id]
	if mount != nil {
		mount.count++
		return mount.path, nil
	} else {
		mount = &ActiveMount{count: 1}
	}

	dir := d.dir(id)
	if _, err := os.Stat(dir); err != nil {
		return "", err
	}

	// If id has a root, just return it
	rootDir := path.Join(dir, "root")
	if _, err := os.Stat(rootDir); err == nil {
		mount.path = rootDir
		d.active[id] = mount
		return mount.path, nil
	}

	lowerId, err := ioutil.ReadFile(path.Join(dir, "lower-id"))
	if err != nil {
		return "", err
	}
	lowerDir := path.Join(d.dir(string(lowerId)), "root")
	upperDir := path.Join(dir, "upper")
	workDir := path.Join(dir, "work")
	mergedDir := path.Join(dir, "merged")

	opts := fmt.Sprintf("lowerdir=%s,upperdir=%s,workdir=%s", lowerDir, upperDir, workDir)
	if err := syscall.Mount("overlay", mergedDir, "overlay", 0, label.FormatMountLabel(opts, mountLabel)); err != nil {
		return "", err
	}
	mount.path = mergedDir
	mount.mounted = true
	d.active[id] = mount

	return mount.path, nil
}

func (d *Driver) Put(id string) {
	// Protect the d.active from concurrent access
	d.Lock()
	defer d.Unlock()

	mount := d.active[id]
	if mount == nil {
		log.Debugf("Put on a non-mounted device %s", id)
		return
	}

	mount.count--
	if mount.count > 0 {
		return
	}

	if mount.mounted {
		if err := syscall.Unmount(mount.path, 0); err != nil {
			log.Debugf("Failed to unmount %s overlay: %v", id, err)
		}
	}

	delete(d.active, id)
}

func (d *Driver) ApplyDiff(id string, parent string, diff archive.ArchiveReader) (size int64, err error) {
	dir := d.dir(id)

	if parent == "" {
		return 0, ErrApplyDiffFallback
	}

	parentRootDir := path.Join(d.dir(parent), "root")
	if _, err := os.Stat(parentRootDir); err != nil {
		return 0, ErrApplyDiffFallback
	}

	// We now know there is a parent, and it has a "root" directory containing
	// the full root filesystem. We can just hardlink it and apply the
	// layer. This relies on two things:
	// 1) ApplyDiff is only run once on a clean (no writes to upper layer) container
	// 2) ApplyDiff doesn't do any in-place writes to files (would break hardlinks)
	// These are all currently true and are not expected to break

	tmpRootDir, err := ioutil.TempDir(dir, "tmproot")
	if err != nil {
		return 0, err
	}
	defer func() {
		if err != nil {
			os.RemoveAll(tmpRootDir)
		} else {
			os.RemoveAll(path.Join(dir, "upper"))
			os.RemoveAll(path.Join(dir, "work"))
			os.RemoveAll(path.Join(dir, "merged"))
			os.RemoveAll(path.Join(dir, "lower-id"))
		}
	}()

	if err = copyDir(parentRootDir, tmpRootDir, CopyHardlink); err != nil {
		return 0, err
	}

	if size, err = chrootarchive.ApplyLayer(tmpRootDir, diff); err != nil {
		return 0, err
	}

	rootDir := path.Join(dir, "root")
	if err := os.Rename(tmpRootDir, rootDir); err != nil {
		return 0, err
	}

	return
}

func (d *Driver) Exists(id string) bool {
	_, err := os.Stat(d.dir(id))
	return err == nil
}
<commit_msg>Add error when running overlay over btrfs.<commit_after>// +build linux

package overlay

import (
	"bufio"
	"fmt"
	"io/ioutil"
	"os"
	"os/exec"
	"path"
	"sync"
	"syscall"

	log "github.com/Sirupsen/logrus"
	"github.com/docker/docker/daemon/graphdriver"
	"github.com/docker/docker/pkg/archive"
	"github.com/docker/docker/pkg/chrootarchive"
	"github.com/docker/libcontainer/label"
)

// This is a small wrapper over the NaiveDiffWriter that lets us have a custom
// implementation of ApplyDiff()

var (
	ErrApplyDiffFallback = fmt.Errorf("Fall back to normal ApplyDiff")
)

type ApplyDiffProtoDriver interface {
	graphdriver.ProtoDriver
	ApplyDiff(id, parent string, diff archive.ArchiveReader) (size int64, err error)
}

type naiveDiffDriverWithApply struct {
	graphdriver.Driver
	applyDiff ApplyDiffProtoDriver
}

func NaiveDiffDriverWithApply(driver ApplyDiffProtoDriver) graphdriver.Driver {
	return &naiveDiffDriverWithApply{
		Driver:    graphdriver.NaiveDiffDriver(driver),
		applyDiff: driver,
	}
}

func (d *naiveDiffDriverWithApply) ApplyDiff(id, parent string, diff archive.ArchiveReader) (int64, error) {
	b, err := d.applyDiff.ApplyDiff(id, parent, diff)
	if err == ErrApplyDiffFallback {
		return d.Driver.ApplyDiff(id, parent, diff)
	}
	return b, err
}

// This backend uses the overlay union filesystem for containers
// plus hard link file sharing for images.

// Each container/image can have a "root" subdirectory which is a plain
// filesystem hierarchy, or they can use overlay.

// If they use overlay there is a "upper" directory and a "lower-id"
// file, as well as "merged" and "work" directories. The "upper"
// directory has the upper layer of the overlay, and "lower-id" contains
// the id of the parent whose "root" directory shall be used as the lower
// layer in the overlay. The overlay itself is mounted in the "merged"
// directory, and the "work" dir is needed for overlay to work.

// When a overlay layer is created there are two cases, either the
// parent has a "root" dir, then we start out with a empty "upper"
// directory overlaid on the parents root. This is typically the
// case with the init layer of a container which is based on an image.
// If there is no "root" in the parent, we inherit the lower-id from
// the parent and start by making a copy if the parents "upper" dir.
// This is typically the case for a container layer which copies
// its parent -init upper layer.

// Additionally we also have a custom implementation of ApplyLayer
// which makes a recursive copy of the parent "root" layer using
// hardlinks to share file data, and then applies the layer on top
// of that. This means all child images share file (but not directory)
// data with the parent.

type ActiveMount struct {
	count   int
	path    string
	mounted bool
}
type Driver struct {
	home       string
	sync.Mutex // Protects concurrent modification to active
	active     map[string]*ActiveMount
}

func init() {
	graphdriver.Register("overlay", Init)
}

func Init(home string, options []string) (graphdriver.Driver, error) {
	if err := supportsOverlay(); err != nil {
		return nil, graphdriver.ErrNotSupported
	}

	// check if they are running over btrfs
	var buf syscall.Statfs_t
	if err := syscall.Statfs(path.Dir(home), &buf); err != nil {
		return nil, err
	}

	switch graphdriver.FsMagic(buf.Type) {
	case graphdriver.FsMagicBtrfs:
		log.Error("'overlay' is not supported over btrfs.")
		return nil, graphdriver.ErrIncompatibleFS
	case graphdriver.FsMagicAufs:
		log.Error("'overlay' is not supported over aufs.")
		return nil, graphdriver.ErrIncompatibleFS
	}

	// Create the driver home dir
	if err := os.MkdirAll(home, 0755); err != nil && !os.IsExist(err) {
		return nil, err
	}

	d := &Driver{
		home:   home,
		active: make(map[string]*ActiveMount),
	}

	return NaiveDiffDriverWithApply(d), nil
}

func supportsOverlay() error {
	// We can try to modprobe overlay first before looking at
	// proc/filesystems for when overlay is supported
	exec.Command("modprobe", "overlay").Run()

	f, err := os.Open("/proc/filesystems")
	if err != nil {
		return err
	}
	defer f.Close()

	s := bufio.NewScanner(f)
	for s.Scan() {
		if s.Text() == "nodev\toverlay" {
			return nil
		}
	}
	log.Error("'overlay' not found as a supported filesystem on this host. Please ensure kernel is new enough and has overlay support loaded.")
	return graphdriver.ErrNotSupported
}

func (d *Driver) String() string {
	return "overlay"
}

func (d *Driver) Status() [][2]string {
	return nil
}

func (d *Driver) Cleanup() error {
	return nil
}

func (d *Driver) Create(id string, parent string) (retErr error) {
	dir := d.dir(id)
	if err := os.MkdirAll(path.Dir(dir), 0700); err != nil {
		return err
	}
	if err := os.Mkdir(dir, 0700); err != nil {
		return err
	}

	defer func() {
		// Clean up on failure
		if retErr != nil {
			os.RemoveAll(dir)
		}
	}()

	// Toplevel images are just a "root" dir
	if parent == "" {
		if err := os.Mkdir(path.Join(dir, "root"), 0755); err != nil {
			return err
		}
		return nil
	}

	parentDir := d.dir(parent)

	// Ensure parent exists
	if _, err := os.Lstat(parentDir); err != nil {
		return err
	}

	// If parent has a root, just do a overlay to it
	parentRoot := path.Join(parentDir, "root")

	if s, err := os.Lstat(parentRoot); err == nil {
		if err := os.Mkdir(path.Join(dir, "upper"), s.Mode()); err != nil {
			return err
		}
		if err := os.Mkdir(path.Join(dir, "work"), 0700); err != nil {
			return err
		}
		if err := os.Mkdir(path.Join(dir, "merged"), 0700); err != nil {
			return err
		}
		if err := ioutil.WriteFile(path.Join(dir, "lower-id"), []byte(parent), 0666); err != nil {
			return err
		}
		return nil
	}

	// Otherwise, copy the upper and the lower-id from the parent

	lowerId, err := ioutil.ReadFile(path.Join(parentDir, "lower-id"))
	if err != nil {
		return err
	}

	if err := ioutil.WriteFile(path.Join(dir, "lower-id"), lowerId, 0666); err != nil {
		return err
	}

	parentUpperDir := path.Join(parentDir, "upper")
	s, err := os.Lstat(parentUpperDir)
	if err != nil {
		return err
	}

	upperDir := path.Join(dir, "upper")
	if err := os.Mkdir(upperDir, s.Mode()); err != nil {
		return err
	}
	if err := os.Mkdir(path.Join(dir, "work"), 0700); err != nil {
		return err
	}
	if err := os.Mkdir(path.Join(dir, "merged"), 0700); err != nil {
		return err
	}

	return copyDir(parentUpperDir, upperDir, 0)
}

func (d *Driver) dir(id string) string {
	return path.Join(d.home, id)
}

func (d *Driver) Remove(id string) error {
	dir := d.dir(id)
	if _, err := os.Stat(dir); err != nil {
		return err
	}
	return os.RemoveAll(dir)
}

func (d *Driver) Get(id string, mountLabel string) (string, error) {
	// Protect the d.active from concurrent access
	d.Lock()
	defer d.Unlock()

	mount := d.active[id]
	if mount != nil {
		mount.count++
		return mount.path, nil
	} else {
		mount = &ActiveMount{count: 1}
	}

	dir := d.dir(id)
	if _, err := os.Stat(dir); err != nil {
		return "", err
	}

	// If id has a root, just return it
	rootDir := path.Join(dir, "root")
	if _, err := os.Stat(rootDir); err == nil {
		mount.path = rootDir
		d.active[id] = mount
		return mount.path, nil
	}

	lowerId, err := ioutil.ReadFile(path.Join(dir, "lower-id"))
	if err != nil {
		return "", err
	}
	lowerDir := path.Join(d.dir(string(lowerId)), "root")
	upperDir := path.Join(dir, "upper")
	workDir := path.Join(dir, "work")
	mergedDir := path.Join(dir, "merged")

	opts := fmt.Sprintf("lowerdir=%s,upperdir=%s,workdir=%s", lowerDir, upperDir, workDir)
	if err := syscall.Mount("overlay", mergedDir, "overlay", 0, label.FormatMountLabel(opts, mountLabel)); err != nil {
		return "", err
	}
	mount.path = mergedDir
	mount.mounted = true
	d.active[id] = mount

	return mount.path, nil
}

func (d *Driver) Put(id string) {
	// Protect the d.active from concurrent access
	d.Lock()
	defer d.Unlock()

	mount := d.active[id]
	if mount == nil {
		log.Debugf("Put on a non-mounted device %s", id)
		return
	}

	mount.count--
	if mount.count > 0 {
		return
	}

	if mount.mounted {
		if err := syscall.Unmount(mount.path, 0); err != nil {
			log.Debugf("Failed to unmount %s overlay: %v", id, err)
		}
	}

	delete(d.active, id)
}

func (d *Driver) ApplyDiff(id string, parent string, diff archive.ArchiveReader) (size int64, err error) {
	dir := d.dir(id)

	if parent == "" {
		return 0, ErrApplyDiffFallback
	}

	parentRootDir := path.Join(d.dir(parent), "root")
	if _, err := os.Stat(parentRootDir); err != nil {
		return 0, ErrApplyDiffFallback
	}

	// We now know there is a parent, and it has a "root" directory containing
	// the full root filesystem. We can just hardlink it and apply the
	// layer. This relies on two things:
	// 1) ApplyDiff is only run once on a clean (no writes to upper layer) container
	// 2) ApplyDiff doesn't do any in-place writes to files (would break hardlinks)
	// These are all currently true and are not expected to break

	tmpRootDir, err := ioutil.TempDir(dir, "tmproot")
	if err != nil {
		return 0, err
	}
	defer func() {
		if err != nil {
			os.RemoveAll(tmpRootDir)
		} else {
			os.RemoveAll(path.Join(dir, "upper"))
			os.RemoveAll(path.Join(dir, "work"))
			os.RemoveAll(path.Join(dir, "merged"))
			os.RemoveAll(path.Join(dir, "lower-id"))
		}
	}()

	if err = copyDir(parentRootDir, tmpRootDir, CopyHardlink); err != nil {
		return 0, err
	}

	if size, err = chrootarchive.ApplyLayer(tmpRootDir, diff); err != nil {
		return 0, err
	}

	rootDir := path.Join(dir, "root")
	if err := os.Rename(tmpRootDir, rootDir); err != nil {
		return 0, err
	}

	return
}

func (d *Driver) Exists(id string) bool {
	_, err := os.Stat(d.dir(id))
	return err == nil
}
<|endoftext|>


<commit_before>// Copyright 2019 PingCAP, 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,
// See the License for the specific language governing permissions and
// limitations under the License.

package expression

import (
	"sync/atomic"
	"testing"
	"time"

	. "github.com/pingcap/check"
	"github.com/pingcap/parser/ast"
	"github.com/pingcap/tidb/types"
	"github.com/pingcap/tidb/util/chunk"
	"github.com/pingcap/tidb/util/mock"
)

var vecBuiltinMiscellaneousCases = map[string][]vecExprBenchCase{
	ast.Inet6Aton: {
		{retEvalType: types.ETString, childrenTypes: []types.EvalType{types.ETString}, geners: []dataGenerator{&ipv6StrGener{newDefaultRandGen()}}},
	},
	ast.IsIPv6: {
		{retEvalType: types.ETInt, childrenTypes: []types.EvalType{types.ETString}},
	},
	ast.Sleep: {
		{retEvalType: types.ETInt, childrenTypes: []types.EvalType{types.ETReal}, geners: []dataGenerator{
			newSelectRealGener([]float64{0, 0.000001}),
		}},
	},
	ast.UUID: {},
	ast.Inet6Ntoa: {
		{retEvalType: types.ETString, childrenTypes: []types.EvalType{types.ETString}, geners: []dataGenerator{
			newSelectStringGener(
				[]string{
					"192.168.0.1",
					"2001:db8::68", //ipv6
				},
			)}},
	},
	ast.InetAton: {
		{retEvalType: types.ETInt, childrenTypes: []types.EvalType{types.ETString}, geners: []dataGenerator{&ipv4StrGener{newDefaultRandGen()}}},
		{retEvalType: types.ETInt, childrenTypes: []types.EvalType{types.ETString}},
		{retEvalType: types.ETInt, childrenTypes: []types.EvalType{types.ETString}, geners: []dataGenerator{
			newSelectStringGener(
				[]string{
					"11.11.11.11.",    // last char is .
					"266.266.266.266", // int in string exceed 255
					"127",
					".122",
					".123.123",
					"127.255",
					"127.2.1",
				},
			)}},
	},
	ast.IsIPv4Mapped: {
		{retEvalType: types.ETInt, childrenTypes: []types.EvalType{types.ETString}, geners: []dataGenerator{&ipv4MappedByteGener{newDefaultRandGen()}}},
		{retEvalType: types.ETInt, childrenTypes: []types.EvalType{types.ETString}, geners: []dataGenerator{&ipv6ByteGener{newDefaultRandGen()}}},
		{retEvalType: types.ETInt, childrenTypes: []types.EvalType{types.ETString}, geners: []dataGenerator{&ipv4ByteGener{newDefaultRandGen()}}},
		{retEvalType: types.ETInt, childrenTypes: []types.EvalType{types.ETString}, geners: []dataGenerator{newDefaultGener(1.0, types.ETString)}},
	},
	ast.IsIPv4Compat: {
		{retEvalType: types.ETInt, childrenTypes: []types.EvalType{types.ETString}, geners: []dataGenerator{&ipv4CompatByteGener{newDefaultRandGen()}}},
		{retEvalType: types.ETInt, childrenTypes: []types.EvalType{types.ETString}, geners: []dataGenerator{&ipv6ByteGener{newDefaultRandGen()}}},
		{retEvalType: types.ETInt, childrenTypes: []types.EvalType{types.ETString}, geners: []dataGenerator{&ipv4ByteGener{newDefaultRandGen()}}},
		{retEvalType: types.ETInt, childrenTypes: []types.EvalType{types.ETString}, geners: []dataGenerator{newDefaultGener(1.0, types.ETString)}},
	},
	ast.InetNtoa: {
		{retEvalType: types.ETString, childrenTypes: []types.EvalType{types.ETInt}},
	},
	ast.IsIPv4: {
		{retEvalType: types.ETInt, childrenTypes: []types.EvalType{types.ETString}},
	},
	ast.AnyValue: {
		{retEvalType: types.ETDuration, childrenTypes: []types.EvalType{types.ETDuration}},
		{retEvalType: types.ETInt, childrenTypes: []types.EvalType{types.ETInt}},
		{retEvalType: types.ETDecimal, childrenTypes: []types.EvalType{types.ETDecimal}},
		{retEvalType: types.ETTimestamp, childrenTypes: []types.EvalType{types.ETTimestamp}},
		{retEvalType: types.ETReal, childrenTypes: []types.EvalType{types.ETReal}},
		{retEvalType: types.ETString, childrenTypes: []types.EvalType{types.ETString}},
		{retEvalType: types.ETJson, childrenTypes: []types.EvalType{types.ETJson}},
	},
	ast.NameConst: {
		{retEvalType: types.ETDuration, childrenTypes: []types.EvalType{types.ETString, types.ETDuration}},
		{retEvalType: types.ETString, childrenTypes: []types.EvalType{types.ETString, types.ETString}},
		{retEvalType: types.ETDecimal, childrenTypes: []types.EvalType{types.ETString, types.ETDecimal}},
		{retEvalType: types.ETInt, childrenTypes: []types.EvalType{types.ETString, types.ETInt}},
		{retEvalType: types.ETReal, childrenTypes: []types.EvalType{types.ETString, types.ETReal}},
		{retEvalType: types.ETJson, childrenTypes: []types.EvalType{types.ETString, types.ETJson}},
		{retEvalType: types.ETTimestamp, childrenTypes: []types.EvalType{types.ETString, types.ETTimestamp}},
	},
}

func (s *testEvaluatorSuite) TestVectorizedBuiltinMiscellaneousEvalOneVec(c *C) {
	testVectorizedEvalOneVec(c, vecBuiltinMiscellaneousCases)
}

func (s *testEvaluatorSuite) TestVectorizedBuiltinMiscellaneousFunc(c *C) {
	testVectorizedBuiltinFunc(c, vecBuiltinMiscellaneousCases)
}

func BenchmarkVectorizedBuiltinMiscellaneousEvalOneVec(b *testing.B) {
	benchmarkVectorizedEvalOneVec(b, vecBuiltinMiscellaneousCases)
}

func BenchmarkVectorizedBuiltinMiscellaneousFunc(b *testing.B) {
	benchmarkVectorizedBuiltinFunc(b, vecBuiltinMiscellaneousCases)
}

type counter struct {
	count int
}

func (c *counter) add(diff int) int {
	c.count += diff
	return c.count
}

func (s *testEvaluatorSuite) TestSleepVectorized(c *C) {
	ctx := mock.NewContext()
	sessVars := ctx.GetSessionVars()

	fc := funcs[ast.Sleep]
	ft := eType2FieldType(types.ETReal)
	col0 := &Column{RetType: ft, Index: 0}
	f, err := fc.getFunction(ctx, []Expression{col0})
	c.Assert(err, IsNil)
	input := chunk.NewChunkWithCapacity([]*types.FieldType{ft}, 1024)
	result := chunk.NewColumn(ft, 1024)
	warnCnt := counter{}

	// non-strict model
	sessVars.StrictSQLMode = false
	input.AppendFloat64(0, 1)
	err = f.vecEvalInt(input, result)
	c.Assert(err, IsNil)
	c.Assert(result.GetInt64(0), Equals, int64(0))
	c.Assert(sessVars.StmtCtx.WarningCount(), Equals, uint16(warnCnt.add(0)))

	input.Reset()
	input.AppendFloat64(0, -1)
	err = f.vecEvalInt(input, result)
	c.Assert(err, IsNil)
	c.Assert(result.GetInt64(0), Equals, int64(0))
	c.Assert(sessVars.StmtCtx.WarningCount(), Equals, uint16(warnCnt.add(1)))

	input.Reset()
	input.AppendNull(0)
	err = f.vecEvalInt(input, result)
	c.Assert(err, IsNil)
	c.Assert(result.GetInt64(0), Equals, int64(0))
	c.Assert(sessVars.StmtCtx.WarningCount(), Equals, uint16(warnCnt.add(1)))

	input.Reset()
	input.AppendNull(0)
	input.AppendFloat64(0, 1)
	input.AppendFloat64(0, -1)
	err = f.vecEvalInt(input, result)
	c.Assert(err, IsNil)
	c.Assert(result.GetInt64(0), Equals, int64(0))
	c.Assert(result.GetInt64(1), Equals, int64(0))
	c.Assert(result.GetInt64(2), Equals, int64(0))
	c.Assert(sessVars.StmtCtx.WarningCount(), Equals, uint16(warnCnt.add(2)))

	// for error case under the strict model
	sessVars.StrictSQLMode = true
	input.Reset()
	input.AppendNull(0)
	err = f.vecEvalInt(input, result)
	c.Assert(err, NotNil)
	c.Assert(result.GetInt64(0), Equals, int64(0))

	sessVars.StmtCtx.SetWarnings(nil)
	input.Reset()
	input.AppendFloat64(0, -2.5)
	err = f.vecEvalInt(input, result)
	c.Assert(err, NotNil)
	c.Assert(result.GetInt64(0), Equals, int64(0))

	//// strict model
	input.Reset()
	input.AppendFloat64(0, 0.5)
	start := time.Now()
	err = f.vecEvalInt(input, result)
	c.Assert(err, IsNil)
	c.Assert(result.GetInt64(0), Equals, int64(0))
	sub := time.Since(start)
	c.Assert(sub.Nanoseconds(), GreaterEqual, int64(0.5*1e9))

	input.Reset()
	input.AppendFloat64(0, 0.01)
	input.AppendFloat64(0, 1)
	input.AppendFloat64(0, 2)
	start = time.Now()
	go func() {
		time.Sleep(1 * time.Second)
		atomic.CompareAndSwapUint32(&ctx.GetSessionVars().Killed, 0, 1)
	}()
	err = f.vecEvalInt(input, result)
	sub = time.Since(start)
	c.Assert(err, IsNil)
	c.Assert(result.GetInt64(0), Equals, int64(0))
	c.Assert(result.GetInt64(1), Equals, int64(1))
	c.Assert(result.GetInt64(2), Equals, int64(1))
	c.Assert(sub.Nanoseconds(), LessEqual, int64(2*1e9))
	c.Assert(sub.Nanoseconds(), GreaterEqual, int64(1*1e9))
}
<commit_msg>test: stabilize test TestSleepVectorized (#17910)<commit_after>// Copyright 2019 PingCAP, 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,
// See the License for the specific language governing permissions and
// limitations under the License.

package expression

import (
	"sync/atomic"
	"testing"
	"time"

	. "github.com/pingcap/check"
	"github.com/pingcap/parser/ast"
	"github.com/pingcap/tidb/types"
	"github.com/pingcap/tidb/util/chunk"
	"github.com/pingcap/tidb/util/mock"
)

var vecBuiltinMiscellaneousCases = map[string][]vecExprBenchCase{
	ast.Inet6Aton: {
		{retEvalType: types.ETString, childrenTypes: []types.EvalType{types.ETString}, geners: []dataGenerator{&ipv6StrGener{newDefaultRandGen()}}},
	},
	ast.IsIPv6: {
		{retEvalType: types.ETInt, childrenTypes: []types.EvalType{types.ETString}},
	},
	ast.Sleep: {
		{retEvalType: types.ETInt, childrenTypes: []types.EvalType{types.ETReal}, geners: []dataGenerator{
			newSelectRealGener([]float64{0, 0.000001}),
		}},
	},
	ast.UUID: {},
	ast.Inet6Ntoa: {
		{retEvalType: types.ETString, childrenTypes: []types.EvalType{types.ETString}, geners: []dataGenerator{
			newSelectStringGener(
				[]string{
					"192.168.0.1",
					"2001:db8::68", //ipv6
				},
			)}},
	},
	ast.InetAton: {
		{retEvalType: types.ETInt, childrenTypes: []types.EvalType{types.ETString}, geners: []dataGenerator{&ipv4StrGener{newDefaultRandGen()}}},
		{retEvalType: types.ETInt, childrenTypes: []types.EvalType{types.ETString}},
		{retEvalType: types.ETInt, childrenTypes: []types.EvalType{types.ETString}, geners: []dataGenerator{
			newSelectStringGener(
				[]string{
					"11.11.11.11.",    // last char is .
					"266.266.266.266", // int in string exceed 255
					"127",
					".122",
					".123.123",
					"127.255",
					"127.2.1",
				},
			)}},
	},
	ast.IsIPv4Mapped: {
		{retEvalType: types.ETInt, childrenTypes: []types.EvalType{types.ETString}, geners: []dataGenerator{&ipv4MappedByteGener{newDefaultRandGen()}}},
		{retEvalType: types.ETInt, childrenTypes: []types.EvalType{types.ETString}, geners: []dataGenerator{&ipv6ByteGener{newDefaultRandGen()}}},
		{retEvalType: types.ETInt, childrenTypes: []types.EvalType{types.ETString}, geners: []dataGenerator{&ipv4ByteGener{newDefaultRandGen()}}},
		{retEvalType: types.ETInt, childrenTypes: []types.EvalType{types.ETString}, geners: []dataGenerator{newDefaultGener(1.0, types.ETString)}},
	},
	ast.IsIPv4Compat: {
		{retEvalType: types.ETInt, childrenTypes: []types.EvalType{types.ETString}, geners: []dataGenerator{&ipv4CompatByteGener{newDefaultRandGen()}}},
		{retEvalType: types.ETInt, childrenTypes: []types.EvalType{types.ETString}, geners: []dataGenerator{&ipv6ByteGener{newDefaultRandGen()}}},
		{retEvalType: types.ETInt, childrenTypes: []types.EvalType{types.ETString}, geners: []dataGenerator{&ipv4ByteGener{newDefaultRandGen()}}},
		{retEvalType: types.ETInt, childrenTypes: []types.EvalType{types.ETString}, geners: []dataGenerator{newDefaultGener(1.0, types.ETString)}},
	},
	ast.InetNtoa: {
		{retEvalType: types.ETString, childrenTypes: []types.EvalType{types.ETInt}},
	},
	ast.IsIPv4: {
		{retEvalType: types.ETInt, childrenTypes: []types.EvalType{types.ETString}},
	},
	ast.AnyValue: {
		{retEvalType: types.ETDuration, childrenTypes: []types.EvalType{types.ETDuration}},
		{retEvalType: types.ETInt, childrenTypes: []types.EvalType{types.ETInt}},
		{retEvalType: types.ETDecimal, childrenTypes: []types.EvalType{types.ETDecimal}},
		{retEvalType: types.ETTimestamp, childrenTypes: []types.EvalType{types.ETTimestamp}},
		{retEvalType: types.ETReal, childrenTypes: []types.EvalType{types.ETReal}},
		{retEvalType: types.ETString, childrenTypes: []types.EvalType{types.ETString}},
		{retEvalType: types.ETJson, childrenTypes: []types.EvalType{types.ETJson}},
	},
	ast.NameConst: {
		{retEvalType: types.ETDuration, childrenTypes: []types.EvalType{types.ETString, types.ETDuration}},
		{retEvalType: types.ETString, childrenTypes: []types.EvalType{types.ETString, types.ETString}},
		{retEvalType: types.ETDecimal, childrenTypes: []types.EvalType{types.ETString, types.ETDecimal}},
		{retEvalType: types.ETInt, childrenTypes: []types.EvalType{types.ETString, types.ETInt}},
		{retEvalType: types.ETReal, childrenTypes: []types.EvalType{types.ETString, types.ETReal}},
		{retEvalType: types.ETJson, childrenTypes: []types.EvalType{types.ETString, types.ETJson}},
		{retEvalType: types.ETTimestamp, childrenTypes: []types.EvalType{types.ETString, types.ETTimestamp}},
	},
}

func (s *testEvaluatorSuite) TestVectorizedBuiltinMiscellaneousEvalOneVec(c *C) {
	testVectorizedEvalOneVec(c, vecBuiltinMiscellaneousCases)
}

func (s *testEvaluatorSuite) TestVectorizedBuiltinMiscellaneousFunc(c *C) {
	testVectorizedBuiltinFunc(c, vecBuiltinMiscellaneousCases)
}

func BenchmarkVectorizedBuiltinMiscellaneousEvalOneVec(b *testing.B) {
	benchmarkVectorizedEvalOneVec(b, vecBuiltinMiscellaneousCases)
}

func BenchmarkVectorizedBuiltinMiscellaneousFunc(b *testing.B) {
	benchmarkVectorizedBuiltinFunc(b, vecBuiltinMiscellaneousCases)
}

type counter struct {
	count int
}

func (c *counter) add(diff int) int {
	c.count += diff
	return c.count
}

func (s *testEvaluatorSuite) TestSleepVectorized(c *C) {
	ctx := mock.NewContext()
	sessVars := ctx.GetSessionVars()

	fc := funcs[ast.Sleep]
	ft := eType2FieldType(types.ETReal)
	col0 := &Column{RetType: ft, Index: 0}
	f, err := fc.getFunction(ctx, []Expression{col0})
	c.Assert(err, IsNil)
	input := chunk.NewChunkWithCapacity([]*types.FieldType{ft}, 1024)
	result := chunk.NewColumn(ft, 1024)
	warnCnt := counter{}

	// non-strict model
	sessVars.StrictSQLMode = false
	input.AppendFloat64(0, 1)
	err = f.vecEvalInt(input, result)
	c.Assert(err, IsNil)
	c.Assert(result.GetInt64(0), Equals, int64(0))
	c.Assert(sessVars.StmtCtx.WarningCount(), Equals, uint16(warnCnt.add(0)))

	input.Reset()
	input.AppendFloat64(0, -1)
	err = f.vecEvalInt(input, result)
	c.Assert(err, IsNil)
	c.Assert(result.GetInt64(0), Equals, int64(0))
	c.Assert(sessVars.StmtCtx.WarningCount(), Equals, uint16(warnCnt.add(1)))

	input.Reset()
	input.AppendNull(0)
	err = f.vecEvalInt(input, result)
	c.Assert(err, IsNil)
	c.Assert(result.GetInt64(0), Equals, int64(0))
	c.Assert(sessVars.StmtCtx.WarningCount(), Equals, uint16(warnCnt.add(1)))

	input.Reset()
	input.AppendNull(0)
	input.AppendFloat64(0, 1)
	input.AppendFloat64(0, -1)
	err = f.vecEvalInt(input, result)
	c.Assert(err, IsNil)
	c.Assert(result.GetInt64(0), Equals, int64(0))
	c.Assert(result.GetInt64(1), Equals, int64(0))
	c.Assert(result.GetInt64(2), Equals, int64(0))
	c.Assert(sessVars.StmtCtx.WarningCount(), Equals, uint16(warnCnt.add(2)))

	// for error case under the strict model
	sessVars.StrictSQLMode = true
	input.Reset()
	input.AppendNull(0)
	err = f.vecEvalInt(input, result)
	c.Assert(err, NotNil)
	c.Assert(result.GetInt64(0), Equals, int64(0))

	sessVars.StmtCtx.SetWarnings(nil)
	input.Reset()
	input.AppendFloat64(0, -2.5)
	err = f.vecEvalInt(input, result)
	c.Assert(err, NotNil)
	c.Assert(result.GetInt64(0), Equals, int64(0))

	//// strict model
	input.Reset()
	input.AppendFloat64(0, 0.5)
	start := time.Now()
	err = f.vecEvalInt(input, result)
	c.Assert(err, IsNil)
	c.Assert(result.GetInt64(0), Equals, int64(0))
	sub := time.Since(start)
	c.Assert(sub.Nanoseconds(), GreaterEqual, int64(0.5*1e9))

	input.Reset()
	input.AppendFloat64(0, 0.01)
	input.AppendFloat64(0, 2)
	input.AppendFloat64(0, 2)
	start = time.Now()
	go func() {
		time.Sleep(1 * time.Second)
		atomic.CompareAndSwapUint32(&ctx.GetSessionVars().Killed, 0, 1)
	}()
	err = f.vecEvalInt(input, result)
	sub = time.Since(start)
	c.Assert(err, IsNil)
	c.Assert(result.GetInt64(0), Equals, int64(0))
	c.Assert(result.GetInt64(1), Equals, int64(1))
	c.Assert(result.GetInt64(2), Equals, int64(1))
	c.Assert(sub.Nanoseconds(), LessEqual, int64(2*1e9))
	c.Assert(sub.Nanoseconds(), GreaterEqual, int64(1*1e9))
}
<|endoftext|>


<commit_before>package commands

import (
	"fmt"
	"io"

	cmds "github.com/jbenet/go-ipfs/commands"
	uio "github.com/jbenet/go-ipfs/unixfs/io"
)

var cat = &cmds.Command{
	Help: "TODO",
	Run: func(req cmds.Request, res cmds.Response) {
		node := req.Context().Node
		fmt.Println(node.Resolver)
		readers := make([]io.Reader, 0, len(req.Arguments()))

		for _, path := range req.Arguments() {
			dagnode, err := node.Resolver.ResolvePath(path)
			if err != nil {
				res.SetError(err, cmds.ErrNormal)
				return
			}

			read, err := uio.NewDagReader(dagnode, node.DAG)
			if err != nil {
				res.SetError(err, cmds.ErrNormal)
				return
			}

			readers = append(readers, read)
		}

		reader := io.MultiReader(readers...)
		res.SetValue(reader)
	},
}
<commit_msg>fix(commands2/cat) rm extraneous print in output<commit_after>package commands

import (
	"io"

	cmds "github.com/jbenet/go-ipfs/commands"
	uio "github.com/jbenet/go-ipfs/unixfs/io"
)

var cat = &cmds.Command{
	Help: "TODO",
	Run: func(req cmds.Request, res cmds.Response) {
		node := req.Context().Node
		readers := make([]io.Reader, 0, len(req.Arguments()))

		for _, path := range req.Arguments() {
			dagnode, err := node.Resolver.ResolvePath(path)
			if err != nil {
				res.SetError(err, cmds.ErrNormal)
				return
			}

			read, err := uio.NewDagReader(dagnode, node.DAG)
			if err != nil {
				res.SetError(err, cmds.ErrNormal)
				return
			}

			readers = append(readers, read)
		}

		reader := io.MultiReader(readers...)
		res.SetValue(reader)
	},
}
<|endoftext|>