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

package store

import (
	"strconv"

	"github.com/dgraph-io/badger-bench/rdb"
	"github.com/dgraph-io/dgraph/x"
)

// Store contains some handles to RocksDB.
type Store struct {
	db       *rdb.DB
	opt      *rdb.Options // Contains blockopt.
	blockopt *rdb.BlockBasedTableOptions
	ropt     *rdb.ReadOptions
	wopt     *rdb.WriteOptions
}

func (s *Store) setOpts() {
	s.opt = rdb.NewDefaultOptions()
	s.blockopt = rdb.NewDefaultBlockBasedTableOptions()
	s.opt.SetBlockBasedTableFactory(s.blockopt)

	// If you want to access blockopt.blockCache, you need to grab handles to them
	// as well. Otherwise, they will be nil. However, for now, we do not really need
	// to do this.
	// s.blockopt.SetBlockCache(rocksdb.NewLRUCache(blockCacheSize))
	// s.blockopt.SetBlockCacheCompressed(rocksdb.NewLRUCache(blockCacheSize))

	s.opt.SetCreateIfMissing(true)
	fp := rdb.NewBloomFilter(16)
	s.blockopt.SetFilterPolicy(fp)

	s.ropt = rdb.NewDefaultReadOptions()
	s.wopt = rdb.NewDefaultWriteOptions()
	s.wopt.SetSync(false) // We don't need to do synchronous writes.
}

// NewStore constructs a Store object at filepath, given some options.
func NewStore(filepath string) (*Store, error) {
	s := &Store{}
	s.setOpts()
	var err error
	s.db, err = rdb.OpenDb(s.opt, filepath)
	return s, x.Wrap(err)
}

func NewSyncStore(filepath string) (*Store, error) {
	s := &Store{}
	s.setOpts()
	s.wopt.SetSync(true) // Do synchronous writes.
	var err error
	s.db, err = rdb.OpenDb(s.opt, filepath)
	return s, x.Wrap(err)
}

// NewReadOnlyStore constructs a readonly Store object at filepath, given options.
func NewReadOnlyStore(filepath string) (*Store, error) {
	s := &Store{}
	s.setOpts()
	var err error
	s.db, err = rdb.OpenDbForReadOnly(s.opt, filepath, false)
	return s, x.Wrap(err)
}

// Get returns the value given a key for RocksDB.
func (s *Store) Get(key []byte) (*rdb.Slice, error) {
	valSlice, err := s.db.Get(s.ropt, key)
	if err != nil {
		return nil, x.Wrapf(err, "Key: %v", key)
	}

	return valSlice, nil
}

// SetOne adds a key-value to data store.
func (s *Store) SetOne(k []byte, val []byte) error { return s.db.Put(s.wopt, k, val) }

// Delete deletes a key from data store.
func (s *Store) Delete(k []byte) error { return s.db.Delete(s.wopt, k) }

// NewIterator initializes a new iterator and returns it.
func (s *Store) NewIterator() *rdb.Iterator {
	ro := rdb.NewDefaultReadOptions()
	// SetFillCache should be set to false for bulk reads to avoid caching data
	// while doing bulk scans.
	ro.SetFillCache(false)
	return s.db.NewIterator(ro)
}

// Close closes our data store.
func (s *Store) Close() { s.db.Close() }

// Memtable returns the memtable size.
func (s *Store) MemtableSize() uint64 {
	memTableSize, _ := strconv.ParseUint(s.db.GetProperty("rocksdb.cur-size-all-mem-tables"), 10, 64)
	return memTableSize
}

// IndexFilterblockSize returns the filter block size.
func (s *Store) IndexFilterblockSize() uint64 {
	blockSize, _ := strconv.ParseUint(s.db.GetProperty("rocksdb.estimate-table-readers-mem"), 10, 64)
	return blockSize
}

// NewWriteBatch creates a new WriteBatch object and returns a pointer to it.
func (s *Store) NewWriteBatch() *rdb.WriteBatch { return rdb.NewWriteBatch() }

// WriteBatch does a batch write to RocksDB from the data in WriteBatch object.
func (s *Store) WriteBatch(wb *rdb.WriteBatch) error {
	return x.Wrap(s.db.Write(s.wopt, wb))
}

// NewCheckpoint creates new checkpoint from current store.
func (s *Store) NewCheckpoint() (*rdb.Checkpoint, error) { return s.db.NewCheckpoint() }

// NewSnapshot creates new snapshot from current store.
func (s *Store) NewSnapshot() *rdb.Snapshot { return s.db.NewSnapshot() }

// SetSnapshot updates default read options to use the given snapshot.
func (s *Store) SetSnapshot(snapshot *rdb.Snapshot) { s.ropt.SetSnapshot(snapshot) }

// GetStats returns stats of our data store.
func (s *Store) GetStats() string { return s.db.GetStats() }
<commit_msg>Remove dependency on dgraph/x package<commit_after>/*
 * Copyright (C) 2017 Dgraph Labs, Inc. and Contributors
 *
 * This program is free software: you can redistribute it and/or modify
 * it under the terms of the GNU Affero General Public License as published by
 * the Free Software Foundation, either version 3 of the License, or
 * (at your option) any later version.
 *
 * This program is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 * GNU Affero General Public License for more details.
 *
 * You should have received a copy of the GNU Affero General Public License
 * along with this program.  If not, see <http://www.gnu.org/licenses/>.
 */

package store

import (
	"strconv"

	"github.com/dgraph-io/badger-bench/rdb"
)

// Store contains some handles to RocksDB.
type Store struct {
	db       *rdb.DB
	opt      *rdb.Options // Contains blockopt.
	blockopt *rdb.BlockBasedTableOptions
	ropt     *rdb.ReadOptions
	wopt     *rdb.WriteOptions
}

func (s *Store) setOpts() {
	s.opt = rdb.NewDefaultOptions()
	s.blockopt = rdb.NewDefaultBlockBasedTableOptions()
	s.opt.SetBlockBasedTableFactory(s.blockopt)

	// If you want to access blockopt.blockCache, you need to grab handles to them
	// as well. Otherwise, they will be nil. However, for now, we do not really need
	// to do this.
	// s.blockopt.SetBlockCache(rocksdb.NewLRUCache(blockCacheSize))
	// s.blockopt.SetBlockCacheCompressed(rocksdb.NewLRUCache(blockCacheSize))

	s.opt.SetCreateIfMissing(true)
	fp := rdb.NewBloomFilter(16)
	s.blockopt.SetFilterPolicy(fp)

	s.ropt = rdb.NewDefaultReadOptions()
	s.wopt = rdb.NewDefaultWriteOptions()
	s.wopt.SetSync(false) // We don't need to do synchronous writes.
}

// NewStore constructs a Store object at filepath, given some options.
func NewStore(filepath string) (*Store, error) {
	s := &Store{}
	s.setOpts()
	var err error
	s.db, err = rdb.OpenDb(s.opt, filepath)
	return s, err
}

func NewSyncStore(filepath string) (*Store, error) {
	s := &Store{}
	s.setOpts()
	s.wopt.SetSync(true) // Do synchronous writes.
	var err error
	s.db, err = rdb.OpenDb(s.opt, filepath)
	return s, err
}

// NewReadOnlyStore constructs a readonly Store object at filepath, given options.
func NewReadOnlyStore(filepath string) (*Store, error) {
	s := &Store{}
	s.setOpts()
	var err error
	s.db, err = rdb.OpenDbForReadOnly(s.opt, filepath, false)
	return s, err
}

// Get returns the value given a key for RocksDB.
func (s *Store) Get(key []byte) (*rdb.Slice, error) {
	valSlice, err := s.db.Get(s.ropt, key)
	if err != nil {
		return nil, err
	}

	return valSlice, nil
}

// SetOne adds a key-value to data store.
func (s *Store) SetOne(k []byte, val []byte) error { return s.db.Put(s.wopt, k, val) }

// Delete deletes a key from data store.
func (s *Store) Delete(k []byte) error { return s.db.Delete(s.wopt, k) }

// NewIterator initializes a new iterator and returns it.
func (s *Store) NewIterator() *rdb.Iterator {
	ro := rdb.NewDefaultReadOptions()
	// SetFillCache should be set to false for bulk reads to avoid caching data
	// while doing bulk scans.
	ro.SetFillCache(false)
	return s.db.NewIterator(ro)
}

// Close closes our data store.
func (s *Store) Close() { s.db.Close() }

// Memtable returns the memtable size.
func (s *Store) MemtableSize() uint64 {
	memTableSize, _ := strconv.ParseUint(s.db.GetProperty("rocksdb.cur-size-all-mem-tables"), 10, 64)
	return memTableSize
}

// IndexFilterblockSize returns the filter block size.
func (s *Store) IndexFilterblockSize() uint64 {
	blockSize, _ := strconv.ParseUint(s.db.GetProperty("rocksdb.estimate-table-readers-mem"), 10, 64)
	return blockSize
}

// NewWriteBatch creates a new WriteBatch object and returns a pointer to it.
func (s *Store) NewWriteBatch() *rdb.WriteBatch { return rdb.NewWriteBatch() }

// WriteBatch does a batch write to RocksDB from the data in WriteBatch object.
func (s *Store) WriteBatch(wb *rdb.WriteBatch) error {
	return s.db.Write(s.wopt, wb)
}

// NewCheckpoint creates new checkpoint from current store.
func (s *Store) NewCheckpoint() (*rdb.Checkpoint, error) { return s.db.NewCheckpoint() }

// NewSnapshot creates new snapshot from current store.
func (s *Store) NewSnapshot() *rdb.Snapshot { return s.db.NewSnapshot() }

// SetSnapshot updates default read options to use the given snapshot.
func (s *Store) SetSnapshot(snapshot *rdb.Snapshot) { s.ropt.SetSnapshot(snapshot) }

// GetStats returns stats of our data store.
func (s *Store) GetStats() string { return s.db.GetStats() }
<|endoftext|>


<commit_before>// Package store provides a distributed SQLite instance.
//
// Distributed consensus is provided via the Raft algorithm.
package store

import (
	"bytes"
	"encoding/json"
	"fmt"
	"io"
	"io/ioutil"
	"log"
	"net"
	"os"
	"path/filepath"
	"sync"
	"time"

	"github.com/hashicorp/raft"
	"github.com/hashicorp/raft-boltdb"
	sql "github.com/otoolep/rqlite/db"
)

const (
	retainSnapshotCount = 2
	raftTimeout         = 10 * time.Second
	sqliteFile          = "db.sqlite"
	leaderWaitDelay     = 100 * time.Millisecond
	appliedWaitDelay    = 100 * time.Millisecond
)

type command struct {
	Tx      bool     `json:"tx,omitempty"`
	Queries []string `json:"queries,omitempty"`
}

// Store is a SQLite database, where all changes are made via Raft consensus.
type Store struct {
	raftDir  string
	raftBind string

	mu sync.RWMutex // Sync access between queries and snapshots.

	ln     *networkLayer // Raft network between nodes.
	raft   *raft.Raft    // The consensus mechanism.
	dbConf *sql.Config   // SQLite database config.
	dbPath string        // Path to database file.
	db     *sql.DB       // The underlying SQLite store.

	logger *log.Logger
}

// New returns a new Store.
func New(dbConf *sql.Config, dir, bind string) *Store {
	dbPath := filepath.Join(dir, sqliteFile)
	if dbConf.Memory {
		dbPath = ":memory:"
	}

	return &Store{
		raftDir:  dir,
		raftBind: bind,
		dbConf:   dbConf,
		dbPath:   dbPath,
		logger:   log.New(os.Stderr, "[store] ", log.LstdFlags),
	}
}

// Open opens the store. If enableSingle is set, and there are no existing peers,
// then this node becomes the first node, and therefore leader, of the cluster.
func (s *Store) Open(enableSingle bool) error {
	// Create the database. Unless it's a memory-based database, tt must be deleted
	if !s.dbConf.Memory {
		// as it will be rebuilt from (possibly) a snapshot and committed log entries.
		if err := os.Remove(s.dbPath); err != nil && !os.IsNotExist(err) {
			return err
		}
	}

	db, err := sql.OpenWithConfiguration(s.dbPath, s.dbConf)
	if err != nil {
		return err
	}
	s.db = db
	s.logger.Println("SQLite database opened at", s.dbConf.FQDSN(s.dbPath))

	// Setup Raft configuration.
	config := raft.DefaultConfig()

	// Check for any existing peers.
	peers, err := readPeersJSON(filepath.Join(s.raftDir, "peers.json"))
	if err != nil {
		return err
	}

	// Allow the node to entry single-mode, potentially electing itself, if
	// explicitly enabled and there is only 1 node in the cluster already.
	if enableSingle && len(peers) <= 1 {
		s.logger.Println("enabling single-node mode")
		config.EnableSingleNode = true
		config.DisableBootstrapAfterElect = false
	}

	// Setup Raft communication.
	ln, err := net.Listen("tcp", s.raftBind)
	if err != nil {
		return err
	}
	s.ln = newnetworkLayer(ln)
	transport := raft.NewNetworkTransport(s.ln, 3, 10*time.Second, os.Stderr)

	// Create peer storage.
	peerStore := raft.NewJSONPeers(s.raftDir, transport)

	// Create the snapshot store. This allows Raft to truncate the log.
	snapshots, err := raft.NewFileSnapshotStore(s.raftDir, retainSnapshotCount, os.Stderr)
	if err != nil {
		return fmt.Errorf("file snapshot store: %s", err)
	}

	// Create the log store and stable store.
	logStore, err := raftboltdb.NewBoltStore(filepath.Join(s.raftDir, "raft.db"))
	if err != nil {
		return fmt.Errorf("new bolt store: %s", err)
	}

	// Instantiate the Raft system.
	ra, err := raft.NewRaft(config, (*fsm)(s), logStore, logStore, snapshots, peerStore, transport)
	if err != nil {
		return fmt.Errorf("new raft: %s", err)
	}
	s.raft = ra

	return nil
}

// Close closes the store.
func (s *Store) Close() error {
	if err := s.db.Close(); err != nil {
		return err
	}
	f := s.raft.Shutdown()
	if e := f.(raft.Future); e.Error() != nil {
		return e.Error()
	}
	return nil
}

// Path returns the path to the store's storage directory.
func (s *Store) Path() string {
	return s.raftDir
}

func (s *Store) Addr() net.Addr {
	return s.ln.Addr()
}

// Leader returns the current leader. Returns a blank string if there is
// no leader.
func (s *Store) Leader() string {
	return s.raft.Leader()
}

// WaitForLeader blocks until a leader is detected, or the timeout expires.
func (s *Store) WaitForLeader(timeout time.Duration) (string, error) {
	tck := time.NewTicker(leaderWaitDelay)
	defer tck.Stop()
	tmr := time.NewTimer(timeout)
	defer tmr.Stop()

	for {
		select {
		case <-tck.C:
			l := s.Leader()
			if l != "" {
				return l, nil
			}
		case <-tmr.C:
			return "", fmt.Errorf("timeout expired")
		}
	}
}

// WaitForAppliedIndex blocks until a given log index has been applied,
// or the timeout expires.
func (s *Store) WaitForAppliedIndex(idx uint64, timeout time.Duration) error {
	tck := time.NewTicker(appliedWaitDelay)
	defer tck.Stop()
	tmr := time.NewTimer(timeout)
	defer tmr.Stop()

	for {
		select {
		case <-tck.C:
			if s.raft.AppliedIndex() >= idx {
				return nil
			}
		case <-tmr.C:
			return fmt.Errorf("timeout expired")
		}
	}
}

// Stats returns stats for the store.
func (s *Store) Stats() (map[string]interface{}, error) {
	dbStatus := map[string]interface{}{
		"path": s.dbPath,
		"dns":  s.dbConf.FQDSN(s.dbPath),
	}
	status := map[string]interface{}{
		"raft":    s.raft.Stats(),
		"addr":    s.Addr(),
		"sqlite3": dbStatus,
	}
	return status, nil
}

// Execute executes queries that return no rows, but do modify the database.
func (s *Store) Execute(queries []string, tx bool) ([]*sql.Result, error) {
	if s.raft.State() != raft.Leader {
		return nil, fmt.Errorf("not leader")
	}

	c := &command{
		Tx:      tx,
		Queries: queries,
	}
	b, err := json.Marshal(c)
	if err != nil {
		return nil, err
	}

	f := s.raft.Apply(b, raftTimeout)
	if e := f.(raft.Future); e.Error() != nil {
		return nil, e.Error()
	}

	r := f.Response().(*fsmResponse)
	return r.results, r.error
}

// Backup return a consistent snapshot of the underlying database.
func (s *Store) Backup(leader bool) ([]byte, error) {
	if leader && s.raft.State() != raft.Leader {
		return nil, fmt.Errorf("not leader")
	}
	return s.db.Backup()
}

// Query executes queries that return rows, and do not modify the database.
func (s *Store) Query(queries []string, tx, leader, verify bool) ([]*sql.Rows, error) {
	// Allow concurrent queries.
	s.mu.RLock()
	defer s.mu.RUnlock()

	if leader && s.raft.State() != raft.Leader {
		return nil, fmt.Errorf("not leader")
	}

	if verify {
		f := s.raft.VerifyLeader()
		if e := f.(raft.Future); e.Error() != nil {
			return nil, e.Error()
		}
	}

	r, err := s.db.Query(queries, tx, true)
	return r, err
}

// Join joins a node, located at addr, to this store. The node must be ready to
// respond to Raft communications at that address.
func (s *Store) Join(addr string) error {
	s.logger.Printf("received request to join node at %s", addr)

	f := s.raft.AddPeer(addr)
	if f.Error() != nil {
		return f.Error()
	}
	s.logger.Printf("node at %s joined successfully", addr)
	return nil
}

type fsm Store

type fsmResponse struct {
	results []*sql.Result
	error   error
}

// Apply applies a Raft log entry to the database.
func (f *fsm) Apply(l *raft.Log) interface{} {
	var c command
	if err := json.Unmarshal(l.Data, &c); err != nil {
		panic(fmt.Sprintf("failed to unmarshal command: %s", err.Error()))
	}

	r, err := f.db.Execute(c.Queries, c.Tx, true)

	return &fsmResponse{results: r, error: err}
}

// Snapshot returns a snapshot of the database. The caller must ensure that
// no transaction is taking place during this call. Hashsicorp Raft guarantees
// that this function will not be called concurrently with Apply.
//
// http://sqlite.org/howtocorrupt.html states it is safe to do this
// as long as no transaction is in progress.
func (f *fsm) Snapshot() (raft.FSMSnapshot, error) {
	// Ensure only one snapshot can take place at once, and block all queries.
	f.mu.Lock()
	defer f.mu.Unlock()

	b, err := ioutil.ReadFile(f.dbPath)
	if err != nil {
		log.Printf("Failed to generate snapshot: %s", err.Error())
		return nil, err
	}
	return &fsmSnapshot{data: b}, nil
}

// Restore restores the database to a previous state.
func (f *fsm) Restore(rc io.ReadCloser) error {
	if err := os.Remove(f.dbPath); err != nil && !os.IsNotExist(err) {
		return err
	}

	b, err := ioutil.ReadAll(rc)
	if err != nil {
		return err
	}

	if err := ioutil.WriteFile(f.dbPath, b, 0660); err != nil {
		return err
	}

	db, err := sql.OpenWithConfiguration(f.dbPath, f.dbConf)
	if err != nil {
		return err
	}
	f.db = db

	return nil
}

type fsmSnapshot struct {
	data []byte
}

// Persist writes the snapshot to the give sink.
func (f *fsmSnapshot) Persist(sink raft.SnapshotSink) error {
	err := func() error {
		// Write data to sink.
		if _, err := sink.Write(f.data); err != nil {
			return err
		}

		// Close the sink.
		if err := sink.Close(); err != nil {
			return err
		}

		return nil
	}()

	if err != nil {
		sink.Cancel()
		return err
	}

	return nil
}

// Release is a no-op.
func (f *fsmSnapshot) Release() {}

// readPeersJSON reads the peers from the path.
func readPeersJSON(path string) ([]string, error) {
	b, err := ioutil.ReadFile(path)
	if err != nil && !os.IsNotExist(err) {
		return nil, err
	}

	if len(b) == 0 {
		return nil, nil
	}

	var peers []string
	dec := json.NewDecoder(bytes.NewReader(b))
	if err := dec.Decode(&peers); err != nil {
		return nil, err
	}

	return peers, nil
}
<commit_msg>Prettier store address in status<commit_after>// Package store provides a distributed SQLite instance.
//
// Distributed consensus is provided via the Raft algorithm.
package store

import (
	"bytes"
	"encoding/json"
	"fmt"
	"io"
	"io/ioutil"
	"log"
	"net"
	"os"
	"path/filepath"
	"sync"
	"time"

	"github.com/hashicorp/raft"
	"github.com/hashicorp/raft-boltdb"
	sql "github.com/otoolep/rqlite/db"
)

const (
	retainSnapshotCount = 2
	raftTimeout         = 10 * time.Second
	sqliteFile          = "db.sqlite"
	leaderWaitDelay     = 100 * time.Millisecond
	appliedWaitDelay    = 100 * time.Millisecond
)

type command struct {
	Tx      bool     `json:"tx,omitempty"`
	Queries []string `json:"queries,omitempty"`
}

// Store is a SQLite database, where all changes are made via Raft consensus.
type Store struct {
	raftDir  string
	raftBind string

	mu sync.RWMutex // Sync access between queries and snapshots.

	ln     *networkLayer // Raft network between nodes.
	raft   *raft.Raft    // The consensus mechanism.
	dbConf *sql.Config   // SQLite database config.
	dbPath string        // Path to database file.
	db     *sql.DB       // The underlying SQLite store.

	logger *log.Logger
}

// New returns a new Store.
func New(dbConf *sql.Config, dir, bind string) *Store {
	dbPath := filepath.Join(dir, sqliteFile)
	if dbConf.Memory {
		dbPath = ":memory:"
	}

	return &Store{
		raftDir:  dir,
		raftBind: bind,
		dbConf:   dbConf,
		dbPath:   dbPath,
		logger:   log.New(os.Stderr, "[store] ", log.LstdFlags),
	}
}

// Open opens the store. If enableSingle is set, and there are no existing peers,
// then this node becomes the first node, and therefore leader, of the cluster.
func (s *Store) Open(enableSingle bool) error {
	// Create the database. Unless it's a memory-based database, tt must be deleted
	if !s.dbConf.Memory {
		// as it will be rebuilt from (possibly) a snapshot and committed log entries.
		if err := os.Remove(s.dbPath); err != nil && !os.IsNotExist(err) {
			return err
		}
	}

	db, err := sql.OpenWithConfiguration(s.dbPath, s.dbConf)
	if err != nil {
		return err
	}
	s.db = db
	s.logger.Println("SQLite database opened at", s.dbConf.FQDSN(s.dbPath))

	// Setup Raft configuration.
	config := raft.DefaultConfig()

	// Check for any existing peers.
	peers, err := readPeersJSON(filepath.Join(s.raftDir, "peers.json"))
	if err != nil {
		return err
	}

	// Allow the node to entry single-mode, potentially electing itself, if
	// explicitly enabled and there is only 1 node in the cluster already.
	if enableSingle && len(peers) <= 1 {
		s.logger.Println("enabling single-node mode")
		config.EnableSingleNode = true
		config.DisableBootstrapAfterElect = false
	}

	// Setup Raft communication.
	ln, err := net.Listen("tcp", s.raftBind)
	if err != nil {
		return err
	}
	s.ln = newnetworkLayer(ln)
	transport := raft.NewNetworkTransport(s.ln, 3, 10*time.Second, os.Stderr)

	// Create peer storage.
	peerStore := raft.NewJSONPeers(s.raftDir, transport)

	// Create the snapshot store. This allows Raft to truncate the log.
	snapshots, err := raft.NewFileSnapshotStore(s.raftDir, retainSnapshotCount, os.Stderr)
	if err != nil {
		return fmt.Errorf("file snapshot store: %s", err)
	}

	// Create the log store and stable store.
	logStore, err := raftboltdb.NewBoltStore(filepath.Join(s.raftDir, "raft.db"))
	if err != nil {
		return fmt.Errorf("new bolt store: %s", err)
	}

	// Instantiate the Raft system.
	ra, err := raft.NewRaft(config, (*fsm)(s), logStore, logStore, snapshots, peerStore, transport)
	if err != nil {
		return fmt.Errorf("new raft: %s", err)
	}
	s.raft = ra

	return nil
}

// Close closes the store.
func (s *Store) Close() error {
	if err := s.db.Close(); err != nil {
		return err
	}
	f := s.raft.Shutdown()
	if e := f.(raft.Future); e.Error() != nil {
		return e.Error()
	}
	return nil
}

// Path returns the path to the store's storage directory.
func (s *Store) Path() string {
	return s.raftDir
}

func (s *Store) Addr() net.Addr {
	return s.ln.Addr()
}

// Leader returns the current leader. Returns a blank string if there is
// no leader.
func (s *Store) Leader() string {
	return s.raft.Leader()
}

// WaitForLeader blocks until a leader is detected, or the timeout expires.
func (s *Store) WaitForLeader(timeout time.Duration) (string, error) {
	tck := time.NewTicker(leaderWaitDelay)
	defer tck.Stop()
	tmr := time.NewTimer(timeout)
	defer tmr.Stop()

	for {
		select {
		case <-tck.C:
			l := s.Leader()
			if l != "" {
				return l, nil
			}
		case <-tmr.C:
			return "", fmt.Errorf("timeout expired")
		}
	}
}

// WaitForAppliedIndex blocks until a given log index has been applied,
// or the timeout expires.
func (s *Store) WaitForAppliedIndex(idx uint64, timeout time.Duration) error {
	tck := time.NewTicker(appliedWaitDelay)
	defer tck.Stop()
	tmr := time.NewTimer(timeout)
	defer tmr.Stop()

	for {
		select {
		case <-tck.C:
			if s.raft.AppliedIndex() >= idx {
				return nil
			}
		case <-tmr.C:
			return fmt.Errorf("timeout expired")
		}
	}
}

// Stats returns stats for the store.
func (s *Store) Stats() (map[string]interface{}, error) {
	dbStatus := map[string]interface{}{
		"path": s.dbPath,
		"dns":  s.dbConf.FQDSN(s.dbPath),
	}
	status := map[string]interface{}{
		"raft":    s.raft.Stats(),
		"addr":    s.Addr().String(),
		"sqlite3": dbStatus,
	}
	return status, nil
}

// Execute executes queries that return no rows, but do modify the database.
func (s *Store) Execute(queries []string, tx bool) ([]*sql.Result, error) {
	if s.raft.State() != raft.Leader {
		return nil, fmt.Errorf("not leader")
	}

	c := &command{
		Tx:      tx,
		Queries: queries,
	}
	b, err := json.Marshal(c)
	if err != nil {
		return nil, err
	}

	f := s.raft.Apply(b, raftTimeout)
	if e := f.(raft.Future); e.Error() != nil {
		return nil, e.Error()
	}

	r := f.Response().(*fsmResponse)
	return r.results, r.error
}

// Backup return a consistent snapshot of the underlying database.
func (s *Store) Backup(leader bool) ([]byte, error) {
	if leader && s.raft.State() != raft.Leader {
		return nil, fmt.Errorf("not leader")
	}
	return s.db.Backup()
}

// Query executes queries that return rows, and do not modify the database.
func (s *Store) Query(queries []string, tx, leader, verify bool) ([]*sql.Rows, error) {
	// Allow concurrent queries.
	s.mu.RLock()
	defer s.mu.RUnlock()

	if leader && s.raft.State() != raft.Leader {
		return nil, fmt.Errorf("not leader")
	}

	if verify {
		f := s.raft.VerifyLeader()
		if e := f.(raft.Future); e.Error() != nil {
			return nil, e.Error()
		}
	}

	r, err := s.db.Query(queries, tx, true)
	return r, err
}

// Join joins a node, located at addr, to this store. The node must be ready to
// respond to Raft communications at that address.
func (s *Store) Join(addr string) error {
	s.logger.Printf("received request to join node at %s", addr)

	f := s.raft.AddPeer(addr)
	if f.Error() != nil {
		return f.Error()
	}
	s.logger.Printf("node at %s joined successfully", addr)
	return nil
}

type fsm Store

type fsmResponse struct {
	results []*sql.Result
	error   error
}

// Apply applies a Raft log entry to the database.
func (f *fsm) Apply(l *raft.Log) interface{} {
	var c command
	if err := json.Unmarshal(l.Data, &c); err != nil {
		panic(fmt.Sprintf("failed to unmarshal command: %s", err.Error()))
	}

	r, err := f.db.Execute(c.Queries, c.Tx, true)

	return &fsmResponse{results: r, error: err}
}

// Snapshot returns a snapshot of the database. The caller must ensure that
// no transaction is taking place during this call. Hashsicorp Raft guarantees
// that this function will not be called concurrently with Apply.
//
// http://sqlite.org/howtocorrupt.html states it is safe to do this
// as long as no transaction is in progress.
func (f *fsm) Snapshot() (raft.FSMSnapshot, error) {
	// Ensure only one snapshot can take place at once, and block all queries.
	f.mu.Lock()
	defer f.mu.Unlock()

	b, err := ioutil.ReadFile(f.dbPath)
	if err != nil {
		log.Printf("Failed to generate snapshot: %s", err.Error())
		return nil, err
	}
	return &fsmSnapshot{data: b}, nil
}

// Restore restores the database to a previous state.
func (f *fsm) Restore(rc io.ReadCloser) error {
	if err := os.Remove(f.dbPath); err != nil && !os.IsNotExist(err) {
		return err
	}

	b, err := ioutil.ReadAll(rc)
	if err != nil {
		return err
	}

	if err := ioutil.WriteFile(f.dbPath, b, 0660); err != nil {
		return err
	}

	db, err := sql.OpenWithConfiguration(f.dbPath, f.dbConf)
	if err != nil {
		return err
	}
	f.db = db

	return nil
}

type fsmSnapshot struct {
	data []byte
}

// Persist writes the snapshot to the give sink.
func (f *fsmSnapshot) Persist(sink raft.SnapshotSink) error {
	err := func() error {
		// Write data to sink.
		if _, err := sink.Write(f.data); err != nil {
			return err
		}

		// Close the sink.
		if err := sink.Close(); err != nil {
			return err
		}

		return nil
	}()

	if err != nil {
		sink.Cancel()
		return err
	}

	return nil
}

// Release is a no-op.
func (f *fsmSnapshot) Release() {}

// readPeersJSON reads the peers from the path.
func readPeersJSON(path string) ([]string, error) {
	b, err := ioutil.ReadFile(path)
	if err != nil && !os.IsNotExist(err) {
		return nil, err
	}

	if len(b) == 0 {
		return nil, nil
	}

	var peers []string
	dec := json.NewDecoder(bytes.NewReader(b))
	if err := dec.Decode(&peers); err != nil {
		return nil, err
	}

	return peers, nil
}
<|endoftext|>


<commit_before>package main

import (
	"encoding/json"
	"fmt"
	"log"
	"os"
	"strings"
	"time"
)

type ReadableBirthday struct {
	time.Time
}

// TODO: support short date format in addition to long, e.g. "Jan 2"
const ctLayout = "January 2"

func (ct *ReadableBirthday) UnmarshalJSON(b []byte) (err error) {
	s := strings.Trim(string(b), "\"")
	if s == "null" {
		ct.Time = time.Time{}
		return
	}
	ct.Time, err = time.Parse(ctLayout, s)
	return
}

func main() {

	type Config struct {
		Name     string
		Birthday ReadableBirthday
		Twitter  string
		Location string
	}

	if len(os.Args) == 1 {
		fmt.Println("I need a .json file to read from, please add this to your go run birthday.go command.")
		os.Exit(1)
	}

	if len(os.Args) > 2 {
		fmt.Println("You gave me too many .json files to read from, please only use one .json in your go run birthday.go command.")
		os.Exit(1)
	}

	var filepath string
	filepath = os.Args[1]

	file, err := os.Open(filepath)

	if err != nil {
		log.Fatal(err)
	}
	decoder := json.NewDecoder(file)
	configs := []Config{}
	err = decoder.Decode(&configs)
	if err != nil {
		fmt.Println("error:", err)
		os.Exit(1)
	}

	// TODO: display birthdays this week

	// TODO: look into `range`
	for i := 0; i < (len(configs)); i++ {
		fmt.Println(configs[i].Name, configs[i].Birthday.Format(ctLayout), configs[i].Twitter, configs[i].Location) // output
	}

}
<commit_msg>make const name more explicit<commit_after>package main

import (
	"encoding/json"
	"fmt"
	"log"
	"os"
	"strings"
	"time"
)

type ReadableBirthday struct {
	time.Time
}

// TODO: support short date format in addition to long, e.g. "Jan 2"
const readableBirthdayLayout = "January 2"

func (ct *ReadableBirthday) UnmarshalJSON(b []byte) (err error) {
	s := strings.Trim(string(b), "\"")
	if s == "null" {
		ct.Time = time.Time{}
		return
	}
	ct.Time, err = time.Parse(readableBirthdayLayout, s)
	return
}

func main() {

	type Config struct {
		Name     string
		Birthday ReadableBirthday
		Twitter  string
		Location string
	}

	if len(os.Args) == 1 {
		fmt.Println("I need a .json file to read from, please add this to your go run birthday.go command.")
		os.Exit(1)
	}

	if len(os.Args) > 2 {
		fmt.Println("You gave me too many .json files to read from, please only use one .json in your go run birthday.go command.")
		os.Exit(1)
	}

	var filepath string
	filepath = os.Args[1]

	file, err := os.Open(filepath)

	if err != nil {
		log.Fatal(err)
	}
	decoder := json.NewDecoder(file)
	configs := []Config{}
	err = decoder.Decode(&configs)
	if err != nil {
		fmt.Println("error:", err)
		os.Exit(1)
	}

	// TODO: display birthdays this week

	// TODO: look into `range`
	for i := 0; i < (len(configs)); i++ {
		fmt.Println(configs[i].Name, configs[i].Birthday.Format(readableBirthdayLayout), configs[i].Twitter, configs[i].Location) // output
	}

}
<|endoftext|>


<commit_before>package main

import (
	"io/ioutil"
	"os"
	"path/filepath"
	"regexp"
	"sort"
	"strings"

	"github.com/LK4D4/vndr/build"
	"github.com/LK4D4/vndr/godl"
)

func isCDir(fis []os.FileInfo) bool {
	var hFound bool
	for _, fi := range fis {
		ext := filepath.Ext(fi.Name())
		if ext == ".cc" || ext == ".cpp" || ext == ".py" {
			return false
		}
		if ext == ".h" {
			hFound = true
		}
	}
	return hFound
}

func isPBDir(fis []os.FileInfo) bool {
	var pbFound bool
	for _, fi := range fis {
		if fi.IsDir() {
			continue
		}
		ext := filepath.Ext(fi.Name())
		if ext != ".proto" {
			return false
		}
		pbFound = true
	}
	return pbFound
}

func isInterestingDir(path string) bool {
	fis, err := ioutil.ReadDir(path)
	if err != nil {
		return false
	}
	return isCDir(fis) || isPBDir(fis)
}

func isGoFile(path string) bool {
	ext := filepath.Ext(path)
	return ext == ".go" || ext == ".c" || ext == ".h" || ext == ".s" || ext == ".proto"
}

// licenseFilesRegexp is a regexp of file names that are likely to contain licensing information
var licenseFilesRegexp = regexp.MustCompile(`(?i).*(LICENSE|COPYING|PATENT|NOTICE|README).*`)

func isLicenseFile(path string) bool {
	return licenseFilesRegexp.MatchString(path) && !isGoFile(path)
}

func isVendorConfFile(path string) bool {
	return filepath.Base(path) == "vendor.conf"
}

// cleanVendor removes files from unused packages and non-go files
func cleanVendor(vendorDir string, realDeps []*build.Package) error {
	realPaths := make(map[string]bool)
	ignoredGoFiles := []string{}
	for _, pkg := range realDeps {
		realPaths[pkg.Dir] = true
		for _, f := range pkg.IgnoredGoFiles {
			ignoredGoFiles = append(ignoredGoFiles, filepath.Join(pkg.Dir, f))
		}
	}
	var paths []string
	err := filepath.Walk(vendorDir, func(path string, i os.FileInfo, err error) error {
		if path == vendorDir {
			return nil
		}
		if err != nil {
			return nil
		}

		// Make sure we don't delete anything that matches the whitelist. The
		// whitelist is relative to the vendor directory.
		relPath, err := filepath.Rel(vendorDir, path)
		if err != nil {
			return err
		}
		if cleanWhitelist.matchString(relPath) {
			return nil
		}

		if strings.HasPrefix(i.Name(), ".") || strings.HasPrefix(i.Name(), "_") {
			return os.RemoveAll(path)
		}
		if i.IsDir() {
			if i.Name() == "testdata" {
				return os.RemoveAll(path)
			}
			if isInterestingDir(path) {
				realPaths[path] = true
				return nil
			}
			if !realPaths[path] {
				paths = append(paths, path)
			}
			return nil
		}

		// keep files for licenses and recursive vendoring
		if isLicenseFile(i.Name()) || isVendorConfFile(i.Name()) {
			return nil
		}
		// remove files from non-deps, non-go files and test files
		if !realPaths[filepath.Dir(path)] || !isGoFile(path) || strings.HasSuffix(path, "_test.go") {
			return os.Remove(path)
		}
		// remove ignored go files
		for _, f := range ignoredGoFiles {
			if f == path {
				return os.Remove(path)
			}
		}
		return nil
	})
	if err != nil {
		return err
	}
	sort.Sort(sort.Reverse(sort.StringSlice(paths)))
	// iterate over paths (longest first)
	for _, p := range paths {
		// at this point we cleaned all files from unused deps dirs
		lst, err := ioutil.ReadDir(p)
		if err != nil {
			return err
		}
		// remove whole dir if there is only non-go files
		// keep whitelisted, though
		var onlyNonGoFile = true
		for _, fi := range lst {
			relPath, err := filepath.Rel(vendorDir, filepath.Join(p, fi.Name()))
			if err != nil {
				return err
			}
			if cleanWhitelist.matchString(relPath) || fi.IsDir() || isGoFile(fi.Name()) {
				onlyNonGoFile = false
				break
			}
		}
		if !onlyNonGoFile {
			continue
		}
		// remove all directories if they're not in dependency paths
		if err := os.RemoveAll(p); err != nil {
			return err
		}
	}
	return nil
}

func cleanVCS(v *godl.VCS) error {
	if err := os.RemoveAll(filepath.Join(v.Root, "."+v.Type)); err != nil {
		return err
	}
	return filepath.Walk(v.Root, func(path string, i os.FileInfo, err error) error {
		if path == vendorDir {
			return nil
		}
		if !i.IsDir() {
			return nil
		}
		name := i.Name()
		if name == "vendor" || name == "Godeps" || name == "_vendor" {
			if err := os.RemoveAll(path); err != nil {
				return err
			}
			return filepath.SkipDir
		}
		return nil
	})
}
<commit_msg>preserve directories with .c files<commit_after>package main

import (
	"io/ioutil"
	"os"
	"path/filepath"
	"regexp"
	"sort"
	"strings"

	"github.com/LK4D4/vndr/build"
	"github.com/LK4D4/vndr/godl"
)

func isCDir(fis []os.FileInfo) bool {
	var hFound bool
	for _, fi := range fis {
		ext := filepath.Ext(fi.Name())
		if ext == ".cc" || ext == ".cpp" || ext == ".py" {
			return false
		}
		if ext == ".h" || ext == ".c" {
			hFound = true
		}
	}
	return hFound
}

func isPBDir(fis []os.FileInfo) bool {
	var pbFound bool
	for _, fi := range fis {
		if fi.IsDir() {
			continue
		}
		ext := filepath.Ext(fi.Name())
		if ext != ".proto" {
			return false
		}
		pbFound = true
	}
	return pbFound
}

func isInterestingDir(path string) bool {
	fis, err := ioutil.ReadDir(path)
	if err != nil {
		return false
	}
	return isCDir(fis) || isPBDir(fis)
}

func isGoFile(path string) bool {
	ext := filepath.Ext(path)
	return ext == ".go" || ext == ".c" || ext == ".h" || ext == ".s" || ext == ".proto"
}

// licenseFilesRegexp is a regexp of file names that are likely to contain licensing information
var licenseFilesRegexp = regexp.MustCompile(`(?i).*(LICENSE|COPYING|PATENT|NOTICE|README).*`)

func isLicenseFile(path string) bool {
	return licenseFilesRegexp.MatchString(path) && !isGoFile(path)
}

func isVendorConfFile(path string) bool {
	return filepath.Base(path) == "vendor.conf"
}

// cleanVendor removes files from unused packages and non-go files
func cleanVendor(vendorDir string, realDeps []*build.Package) error {
	realPaths := make(map[string]bool)
	ignoredGoFiles := []string{}
	for _, pkg := range realDeps {
		realPaths[pkg.Dir] = true
		for _, f := range pkg.IgnoredGoFiles {
			ignoredGoFiles = append(ignoredGoFiles, filepath.Join(pkg.Dir, f))
		}
	}
	var paths []string
	err := filepath.Walk(vendorDir, func(path string, i os.FileInfo, err error) error {
		if path == vendorDir {
			return nil
		}
		if err != nil {
			return nil
		}

		// Make sure we don't delete anything that matches the whitelist. The
		// whitelist is relative to the vendor directory.
		relPath, err := filepath.Rel(vendorDir, path)
		if err != nil {
			return err
		}
		if cleanWhitelist.matchString(relPath) {
			return nil
		}

		if strings.HasPrefix(i.Name(), ".") || strings.HasPrefix(i.Name(), "_") {
			return os.RemoveAll(path)
		}
		if i.IsDir() {
			if i.Name() == "testdata" {
				return os.RemoveAll(path)
			}
			if isInterestingDir(path) {
				realPaths[path] = true
				return nil
			}
			if !realPaths[path] {
				paths = append(paths, path)
			}
			return nil
		}

		// keep files for licenses and recursive vendoring
		if isLicenseFile(i.Name()) || isVendorConfFile(i.Name()) {
			return nil
		}
		// remove files from non-deps, non-go files and test files
		if !realPaths[filepath.Dir(path)] || !isGoFile(path) || strings.HasSuffix(path, "_test.go") {
			return os.Remove(path)
		}
		// remove ignored go files
		for _, f := range ignoredGoFiles {
			if f == path {
				return os.Remove(path)
			}
		}
		return nil
	})
	if err != nil {
		return err
	}
	sort.Sort(sort.Reverse(sort.StringSlice(paths)))
	// iterate over paths (longest first)
	for _, p := range paths {
		// at this point we cleaned all files from unused deps dirs
		lst, err := ioutil.ReadDir(p)
		if err != nil {
			return err
		}
		// remove whole dir if there is only non-go files
		// keep whitelisted, though
		var onlyNonGoFile = true
		for _, fi := range lst {
			relPath, err := filepath.Rel(vendorDir, filepath.Join(p, fi.Name()))
			if err != nil {
				return err
			}
			if cleanWhitelist.matchString(relPath) || fi.IsDir() || isGoFile(fi.Name()) {
				onlyNonGoFile = false
				break
			}
		}
		if !onlyNonGoFile {
			continue
		}
		// remove all directories if they're not in dependency paths
		if err := os.RemoveAll(p); err != nil {
			return err
		}
	}
	return nil
}

func cleanVCS(v *godl.VCS) error {
	if err := os.RemoveAll(filepath.Join(v.Root, "."+v.Type)); err != nil {
		return err
	}
	return filepath.Walk(v.Root, func(path string, i os.FileInfo, err error) error {
		if path == vendorDir {
			return nil
		}
		if !i.IsDir() {
			return nil
		}
		name := i.Name()
		if name == "vendor" || name == "Godeps" || name == "_vendor" {
			if err := os.RemoveAll(path); err != nil {
				return err
			}
			return filepath.SkipDir
		}
		return nil
	})
}
<|endoftext|>


<commit_before>// Clock counts down to or up from a target time.
package main

import (
	"fmt"
	"os"
	"time"
)

func main() {
	target := time.Date(2018, 8, 16, 0, 0, 0, 0, time.Local)
	motto := "Just Go"
	printTargetTime(target, motto)
	exitOnEnterKey()

	var previous time.Time
	for {
		now := time.Now().Truncate(time.Second)
		if now != previous {
			previous = now
			countdown := now.Sub(target) // Negative times are before the target
			printCountdown(now, countdown)
		}
		time.Sleep(100 * time.Millisecond)
	}
}

func exitOnEnterKey() {
	go func() {
		buf := make([]byte, 1)
		_, _ = os.Stdin.Read(buf)
		os.Exit(0)
	}()
}

const (
	indent         = "\t"
	highlightStart = "\x1b[1;35m"
	highlightEnd   = "\x1b[0m"
)

func printTargetTime(target time.Time, motto string) {
	fmt.Print(indent, highlightStart, motto, highlightEnd, "\n")
	fmt.Print(indent, target.Format(time.UnixDate), "\n")
}

func printCountdown(now time.Time, countdown time.Duration) {
	var sign string
	if countdown >= 0 {
		sign = "+"
	} else {
		sign = "-"
		countdown = -countdown
	}

	days := int(countdown / (24 * time.Hour))
	countdown = countdown % (24 * time.Hour)

	fmt.Print(indent, now.Format(time.UnixDate), "  ", sign)
	if days > 0 {
		fmt.Print(days, "d")
	}
	fmt.Print(countdown, "          \r")
	os.Stdout.Sync()
}
<commit_msg>Go 1.11<commit_after>// Clock counts down to or up from a target time.
package main

import (
	"fmt"
	"os"
	"time"
)

func main() {
	target := time.Date(2018, 8, 25, 0, 0, 0, 0, time.Local)
	motto := "Just Go"
	printTargetTime(target, motto)
	exitOnEnterKey()

	var previous time.Time
	for {
		now := time.Now().Truncate(time.Second)
		if now != previous {
			previous = now
			countdown := now.Sub(target) // Negative times are before the target
			printCountdown(now, countdown)
		}
		time.Sleep(100 * time.Millisecond)
	}
}

func exitOnEnterKey() {
	go func() {
		buf := make([]byte, 1)
		_, _ = os.Stdin.Read(buf)
		os.Exit(0)
	}()
}

const (
	indent         = "\t"
	highlightStart = "\x1b[1;35m"
	highlightEnd   = "\x1b[0m"
)

func printTargetTime(target time.Time, motto string) {
	fmt.Print(indent, highlightStart, motto, highlightEnd, "\n")
	fmt.Print(indent, target.Format(time.UnixDate), "\n")
}

func printCountdown(now time.Time, countdown time.Duration) {
	var sign string
	if countdown >= 0 {
		sign = "+"
	} else {
		sign = "-"
		countdown = -countdown
	}

	days := int(countdown / (24 * time.Hour))
	countdown = countdown % (24 * time.Hour)

	fmt.Print(indent, now.Format(time.UnixDate), "  ", sign)
	if days > 0 {
		fmt.Print(days, "d")
	}
	fmt.Print(countdown, "          \r")
	os.Stdout.Sync()
}
<|endoftext|>


<commit_before>package multiaddr

import (
	"encoding/binary"
	"fmt"
	"net"
	"strconv"
	"strings"
)

func StringToBytes(s string) ([]byte, error) {
	b := []byte{}
	sp := strings.Split(s, "/")

	// consume first empty elem
	sp = sp[1:]

	for len(sp) > 0 {
		p := ProtocolWithName(sp[0])
		if p == nil {
			return nil, fmt.Errorf("no protocol with name %s", sp[0])
		}
		b = append(b, byte(p.Code))

		a := AddressStringToBytes(p, sp[1])
		b = append(b, a...)

		sp = sp[2:]
	}
	return b, nil
}

func BytesToString(b []byte) (ret string, err error) {
	// panic handler, in case we try accessing bytes incorrectly.
	defer func() {
		if e := recover(); e != nil {
			ret = ""
			err = e.(error)
		}
	}()

	s := ""

	for len(b) > 0 {
		p := ProtocolWithCode(int(b[0]))
		if p == nil {
			return "", fmt.Errorf("no protocol with code %d", b[0])
		}
		s = strings.Join([]string{s, "/", p.Name}, "")
		b = b[1:]

		a := AddressBytesToString(p, b[:(p.Size/8)])
		if len(a) > 0 {
			s = strings.Join([]string{s, "/", a}, "")
		}
		b = b[(p.Size / 8):]
	}

	return s, nil
}

func AddressStringToBytes(p *Protocol, s string) []byte {
	switch p.Code {

	// ipv4,6
	case 4, 41:
		return net.ParseIP(s).To4()

	// tcp udp dccp sctp
	case 6, 17, 33, 132:
		b := make([]byte, 2)
		i, err := strconv.Atoi(s)
		if err == nil {
			binary.BigEndian.PutUint16(b, uint16(i))
		}
		return b
	}

	return []byte{}
}

func AddressBytesToString(p *Protocol, b []byte) string {
	switch p.Code {

	// ipv4,6
	case 4, 41:
		return net.IP(b).String()

	// tcp udp dccp sctp
	case 6, 17, 33, 132:
		i := binary.BigEndian.Uint16(b)
		return strconv.Itoa(int(i))
	}

	return ""
}
<commit_msg>do not accept invalid-start multiaddrs<commit_after>package multiaddr

import (
	"encoding/binary"
	"fmt"
	"net"
	"strconv"
	"strings"
)

func StringToBytes(s string) ([]byte, error) {
	b := []byte{}
	sp := strings.Split(s, "/")

	if sp[0] != "" {
		return nil, fmt.Errorf("invalid multiaddr, must begin with /")
	}

	// consume first empty elem
	sp = sp[1:]

	for len(sp) > 0 {
		p := ProtocolWithName(sp[0])
		if p == nil {
			return nil, fmt.Errorf("no protocol with name %s", sp[0])
		}
		b = append(b, byte(p.Code))

		a := AddressStringToBytes(p, sp[1])
		b = append(b, a...)

		sp = sp[2:]
	}
	return b, nil
}

func BytesToString(b []byte) (ret string, err error) {
	// panic handler, in case we try accessing bytes incorrectly.
	defer func() {
		if e := recover(); e != nil {
			ret = ""
			err = e.(error)
		}
	}()

	s := ""

	for len(b) > 0 {
		p := ProtocolWithCode(int(b[0]))
		if p == nil {
			return "", fmt.Errorf("no protocol with code %d", b[0])
		}
		s = strings.Join([]string{s, "/", p.Name}, "")
		b = b[1:]

		a := AddressBytesToString(p, b[:(p.Size/8)])
		if len(a) > 0 {
			s = strings.Join([]string{s, "/", a}, "")
		}
		b = b[(p.Size / 8):]
	}

	return s, nil
}

func AddressStringToBytes(p *Protocol, s string) []byte {
	switch p.Code {

	// ipv4,6
	case 4, 41:
		return net.ParseIP(s).To4()

	// tcp udp dccp sctp
	case 6, 17, 33, 132:
		b := make([]byte, 2)
		i, err := strconv.Atoi(s)
		if err == nil {
			binary.BigEndian.PutUint16(b, uint16(i))
		}
		return b
	}

	return []byte{}
}

func AddressBytesToString(p *Protocol, b []byte) string {
	switch p.Code {

	// ipv4,6
	case 4, 41:
		return net.IP(b).String()

	// tcp udp dccp sctp
	case 6, 17, 33, 132:
		i := binary.BigEndian.Uint16(b)
		return strconv.Itoa(int(i))
	}

	return ""
}
<|endoftext|>


<commit_before>package tsdb

import (
	"io/ioutil"
	"os"
	"path/filepath"
	"strings"
	"testing"
	"time"

	"github.com/influxdb/influxdb/influxql"
	"github.com/influxdb/influxdb/meta"
)

var shardID = uint64(1)

func TestWritePointsAndExecuteQuery(t *testing.T) {
	store, executor := testStoreAndExecutor()
	defer os.RemoveAll(store.path)

	pt := NewPoint(
		"cpu",
		map[string]string{"host": "server"},
		map[string]interface{}{"value": 1.0},
		time.Unix(1, 2),
	)

	err := store.WriteToShard(shardID, []Point{pt})
	if err != nil {
		t.Fatalf(err.Error())
	}

	pt.SetTime(time.Unix(2, 3))
	err = store.WriteToShard(shardID, []Point{pt})
	if err != nil {
		t.Fatalf(err.Error())
	}

	got := executeAndGetJSON("select * from cpu", executor)
	exepected := `[{"series":[{"name":"cpu","tags":{"host":"server"},"columns":["time","value"],"values":[["1970-01-01T00:00:01.000000002Z",1],["1970-01-01T00:00:02.000000003Z",1]]}]}]`
	if exepected != got {
		t.Fatalf("exp: %s\ngot: %s", exepected, got)
	}

	store.Close()
	store = NewStore(store.path)
	err = store.Open()
	if err != nil {
		t.Fatalf(err.Error())
	}
	executor.store = store

	got = executeAndGetJSON("select * from cpu", executor)
	if exepected != got {
		t.Fatalf("exp: %s\ngot: %s", exepected, got)
	}
}

func TestDropSeriesStatement(t *testing.T) {
	store, executor := testStoreAndExecutor()
	defer os.RemoveAll(store.path)

	pt := NewPoint(
		"cpu",
		map[string]string{"host": "server"},
		map[string]interface{}{"value": 1.0},
		time.Unix(1, 2),
	)

	err := store.WriteToShard(shardID, []Point{pt})
	if err != nil {
		t.Fatalf(err.Error())
	}

	got := executeAndGetJSON("select * from cpu", executor)
	exepected := `[{"series":[{"name":"cpu","tags":{"host":"server"},"columns":["time","value"],"values":[["1970-01-01T00:00:01.000000002Z",1]]}]}]`
	if exepected != got {
		t.Fatalf("exp: %s\ngot: %s", exepected, got)
	}

	got = executeAndGetJSON("drop series from cpu", executor)

	got = executeAndGetJSON("select * from cpu", executor)
	exepected = `[{}]`
	if exepected != got {
		t.Fatalf("exp: %s\ngot: %s", exepected, got)
	}

	got = executeAndGetJSON("show tag keys from cpu", executor)
	exepected = `[{"series":[{"name":"cpu","columns":["tagKey"]}]}]`
	if exepected != got {
		t.Fatalf("exp: %s\ngot: %s", exepected, got)
	}

	store.Close()
	store = NewStore(store.path)
	store.Open()
	executor.store = store

	got = executeAndGetJSON("select * from cpu", executor)
	exepected = `[{}]`
	if exepected != got {
		t.Fatalf("exp: %s\ngot: %s", exepected, got)
	}

	got = executeAndGetJSON("show tag keys from cpu", executor)
	exepected = `[{"series":[{"name":"cpu","columns":["tagKey"]}]}]`
	if exepected != got {
		t.Fatalf("exp: %s\ngot: %s", exepected, got)
	}
}

func TestDropMeasurementStatement(t *testing.T) {
	store, executor := testStoreAndExecutor()
	defer os.RemoveAll(store.path)

	pt := NewPoint(
		"cpu",
		map[string]string{"host": "server"},
		map[string]interface{}{"value": 1.0},
		time.Unix(1, 2),
	)
	pt2 := NewPoint(
		"memory",
		map[string]string{"host": "server"},
		map[string]interface{}{"value": 1.0},
		time.Unix(1, 2),
	)

	err := store.WriteToShard(shardID, []Point{pt, pt2})
	if err != nil {
		t.Fatalf(err.Error())
	}

	got := executeAndGetJSON("show series", executor)
	exepected := `[{"series":[{"name":"cpu","columns":["_key","host"],"values":[["cpu,host=server","server"]]},{"name":"memory","columns":["_key","host"],"values":[["memory,host=server","server"]]}]}]`
	if exepected != got {
		t.Fatalf("exp: %s\ngot: %s", exepected, got)
	}

	got = executeAndGetJSON("drop measurement memory", executor)
	exepected = `[{}]`
	if exepected != got {
		t.Fatalf("exp: %s\ngot: %s", exepected, got)
	}

	validateDrop := func() {
		got = executeAndGetJSON("show series", executor)
		exepected = `[{"series":[{"name":"cpu","columns":["_key","host"],"values":[["cpu,host=server","server"]]}]}]`
		if exepected != got {
			t.Fatalf("exp: %s\ngot: %s", exepected, got)
		}
		got = executeAndGetJSON("show measurements", executor)
		exepected = `[{"series":[{"name":"measurements","columns":["name"],"values":[["cpu"]]}]}]`
		if exepected != got {
			t.Fatalf("exp: %s\ngot: %s", exepected, got)
		}
		got = executeAndGetJSON("select * from memory", executor)
		exepected = `[{"error":"measurement not found: \"foo\".\"foo\".memory"}]`
		if exepected != got {
			t.Fatalf("exp: %s\ngot: %s", exepected, got)
		}
	}

	validateDrop()
	store.Close()
	store = NewStore(store.path)
	store.Open()
	executor.store = store
	validateDrop()
}

// mock for the metaExecutor
type metaExec struct {
	fn func(stmt influxql.Statement) *influxql.Result
}

func (m *metaExec) ExecuteStatement(stmt influxql.Statement) *influxql.Result {
	return m.fn(stmt)
}

func TestDropDatabase(t *testing.T) {
	store, executor := testStoreAndExecutor()
	defer os.RemoveAll(store.path)

	pt := NewPoint(
		"cpu",
		map[string]string{"host": "server"},
		map[string]interface{}{"value": 1.0},
		time.Unix(1, 2),
	)

	err := store.WriteToShard(shardID, []Point{pt})
	if err != nil {
		t.Fatalf(err.Error())
	}

	got := executeAndGetJSON("select * from cpu", executor)
	expected := `[{"series":[{"name":"cpu","tags":{"host":"server"},"columns":["time","value"],"values":[["1970-01-01T00:00:01.000000002Z",1]]}]}]`
	if expected != got {
		t.Fatalf("exp: %s\ngot: %s", expected, got)
	}

	var name string
	me := &metaExec{fn: func(stmt influxql.Statement) *influxql.Result {
		name = stmt.(*influxql.DropDatabaseStatement).Name
		return &influxql.Result{}
	}}
	executor.MetaStatementExecutor = me

	// verify the database is there on disk
	dbPath := filepath.Join(store.path, "foo")
	if _, err := os.Stat(dbPath); err != nil {
		t.Fatalf("execpted database dir %s to exist", dbPath)
	}

	got = executeAndGetJSON("drop database foo", executor)
	expected = `[{}]`
	if got != expected {
		t.Fatalf("exp: %s\ngot: %s", expected, got)
	}

	if name != "foo" {
		t.Fatalf("expected the MetaStatementExecutor to be called with database name foo, but got %s", name)
	}

	if _, err := os.Stat(dbPath); !os.IsNotExist(err) {
		t.Fatalf("expected database dir %s to be gone", dbPath)
	}

	store.Close()
	store = NewStore(store.path)
	store.Open()
	executor.store = store

	err = store.WriteToShard(shardID, []Point{pt})
	if err == nil || err.Error() != "shard not found" {
		t.Fatalf("expected shard to not be found")
	}
}

// ensure that authenticate doesn't return an error if the user count is zero and they're attempting
// to create a user.
func TestAuthenticateIfUserCountZeroAndCreateUser(t *testing.T) {
	store, executor := testStoreAndExecutor()
	defer os.RemoveAll(store.path)
	ms := &testMetastore{userCount: 0}
	executor.MetaStore = ms

	if err := executor.Authorize(nil, mustParseQuery("create user foo with password 'asdf' with all privileges"), ""); err != nil {
		t.Fatalf("should have authenticated if no users and attempting to create a user but got error: %s", err.Error())
	}

	if executor.Authorize(nil, mustParseQuery("create user foo with password 'asdf'"), "") == nil {
		t.Fatalf("should have failed authentication if no user given and no users exist for create user query that doesn't grant all privileges")
	}

	if executor.Authorize(nil, mustParseQuery("select * from foo"), "") == nil {
		t.Fatalf("should have failed authentication if no user given and no users exist for any query other than create user")
	}

	ms.userCount = 1

	if executor.Authorize(nil, mustParseQuery("create user foo with password 'asdf'"), "") == nil {
		t.Fatalf("should have failed authentication if no user given and users exist")
	}

	if executor.Authorize(nil, mustParseQuery("select * from foo"), "") == nil {
		t.Fatalf("should have failed authentication if no user given and users exist")
	}
}

func testStoreAndExecutor() (*Store, *QueryExecutor) {
	path, _ := ioutil.TempDir("", "")

	store := NewStore(path)
	err := store.Open()
	if err != nil {
		panic(err)
	}
	database := "foo"
	retentionPolicy := "bar"
	shardID := uint64(1)
	store.CreateShard(database, retentionPolicy, shardID)

	executor := NewQueryExecutor(store)
	executor.MetaStore = &testMetastore{}

	return store, executor
}

func executeAndGetJSON(query string, executor *QueryExecutor) string {
	ch, err := executor.ExecuteQuery(mustParseQuery(query), "foo", 20)
	if err != nil {
		panic(err.Error())
	}

	var results []*influxql.Result
	for r := range ch {
		results = append(results, r)
	}
	return string(mustMarshalJSON(results))
}

type testMetastore struct {
	userCount int
}

func (t *testMetastore) Database(name string) (*meta.DatabaseInfo, error) {
	return &meta.DatabaseInfo{
		Name: name,
		DefaultRetentionPolicy: "foo",
		RetentionPolicies: []meta.RetentionPolicyInfo{
			{
				Name: "bar",
				ShardGroups: []meta.ShardGroupInfo{
					{
						ID:        uint64(1),
						StartTime: time.Now().Add(-time.Hour),
						EndTime:   time.Now().Add(time.Hour),
						Shards: []meta.ShardInfo{
							{
								ID:       uint64(1),
								OwnerIDs: []uint64{1},
							},
						},
					},
				},
			},
		},
	}, nil
}

func (t *testMetastore) Databases() ([]meta.DatabaseInfo, error) {
	db, _ := t.Database("foo")
	return []meta.DatabaseInfo{*db}, nil
}

func (t *testMetastore) User(name string) (*meta.UserInfo, error) { return nil, nil }

func (t *testMetastore) AdminUserExists() (bool, error) { return false, nil }

func (t *testMetastore) Authenticate(username, password string) (*meta.UserInfo, error) {
	return nil, nil
}

func (t *testMetastore) RetentionPolicy(database, name string) (rpi *meta.RetentionPolicyInfo, err error) {
	return &meta.RetentionPolicyInfo{
		Name: "bar",
		ShardGroups: []meta.ShardGroupInfo{
			{
				ID:        uint64(1),
				StartTime: time.Now().Add(-time.Hour),
				EndTime:   time.Now().Add(time.Hour),
				Shards: []meta.ShardInfo{
					{
						ID:       uint64(1),
						OwnerIDs: []uint64{1},
					},
				},
			},
		},
	}, nil
}

func (t *testMetastore) UserCount() (int, error) {
	return t.userCount, nil
}

// MustParseQuery parses an InfluxQL query. Panic on error.
func mustParseQuery(s string) *influxql.Query {
	q, err := influxql.NewParser(strings.NewReader(s)).ParseQuery()
	if err != nil {
		panic(err.Error())
	}
	return q
}
<commit_msg>Unit test returning empty set when no database<commit_after>package tsdb

import (
	"io/ioutil"
	"os"
	"path/filepath"
	"strings"
	"testing"
	"time"

	"github.com/influxdb/influxdb/influxql"
	"github.com/influxdb/influxdb/meta"
)

var shardID = uint64(1)

func TestWritePointsAndExecuteQuery(t *testing.T) {
	store, executor := testStoreAndExecutor()
	defer os.RemoveAll(store.path)

	pt := NewPoint(
		"cpu",
		map[string]string{"host": "server"},
		map[string]interface{}{"value": 1.0},
		time.Unix(1, 2),
	)

	err := store.WriteToShard(shardID, []Point{pt})
	if err != nil {
		t.Fatalf(err.Error())
	}

	pt.SetTime(time.Unix(2, 3))
	err = store.WriteToShard(shardID, []Point{pt})
	if err != nil {
		t.Fatalf(err.Error())
	}

	got := executeAndGetJSON("select * from cpu", executor)
	exepected := `[{"series":[{"name":"cpu","tags":{"host":"server"},"columns":["time","value"],"values":[["1970-01-01T00:00:01.000000002Z",1],["1970-01-01T00:00:02.000000003Z",1]]}]}]`
	if exepected != got {
		t.Fatalf("exp: %s\ngot: %s", exepected, got)
	}

	store.Close()
	store = NewStore(store.path)
	err = store.Open()
	if err != nil {
		t.Fatalf(err.Error())
	}
	executor.store = store

	got = executeAndGetJSON("select * from cpu", executor)
	if exepected != got {
		t.Fatalf("exp: %s\ngot: %s", exepected, got)
	}
}

func TestDropSeriesStatement(t *testing.T) {
	store, executor := testStoreAndExecutor()
	defer os.RemoveAll(store.path)

	pt := NewPoint(
		"cpu",
		map[string]string{"host": "server"},
		map[string]interface{}{"value": 1.0},
		time.Unix(1, 2),
	)

	err := store.WriteToShard(shardID, []Point{pt})
	if err != nil {
		t.Fatalf(err.Error())
	}

	got := executeAndGetJSON("select * from cpu", executor)
	exepected := `[{"series":[{"name":"cpu","tags":{"host":"server"},"columns":["time","value"],"values":[["1970-01-01T00:00:01.000000002Z",1]]}]}]`
	if exepected != got {
		t.Fatalf("exp: %s\ngot: %s", exepected, got)
	}

	got = executeAndGetJSON("drop series from cpu", executor)

	got = executeAndGetJSON("select * from cpu", executor)
	exepected = `[{}]`
	if exepected != got {
		t.Fatalf("exp: %s\ngot: %s", exepected, got)
	}

	got = executeAndGetJSON("show tag keys from cpu", executor)
	exepected = `[{"series":[{"name":"cpu","columns":["tagKey"]}]}]`
	if exepected != got {
		t.Fatalf("exp: %s\ngot: %s", exepected, got)
	}

	store.Close()
	store = NewStore(store.path)
	store.Open()
	executor.store = store

	got = executeAndGetJSON("select * from cpu", executor)
	exepected = `[{}]`
	if exepected != got {
		t.Fatalf("exp: %s\ngot: %s", exepected, got)
	}

	got = executeAndGetJSON("show tag keys from cpu", executor)
	exepected = `[{"series":[{"name":"cpu","columns":["tagKey"]}]}]`
	if exepected != got {
		t.Fatalf("exp: %s\ngot: %s", exepected, got)
	}
}

func TestDropMeasurementStatement(t *testing.T) {
	store, executor := testStoreAndExecutor()
	defer os.RemoveAll(store.path)

	pt := NewPoint(
		"cpu",
		map[string]string{"host": "server"},
		map[string]interface{}{"value": 1.0},
		time.Unix(1, 2),
	)
	pt2 := NewPoint(
		"memory",
		map[string]string{"host": "server"},
		map[string]interface{}{"value": 1.0},
		time.Unix(1, 2),
	)

	err := store.WriteToShard(shardID, []Point{pt, pt2})
	if err != nil {
		t.Fatalf(err.Error())
	}

	got := executeAndGetJSON("show series", executor)
	exepected := `[{"series":[{"name":"cpu","columns":["_key","host"],"values":[["cpu,host=server","server"]]},{"name":"memory","columns":["_key","host"],"values":[["memory,host=server","server"]]}]}]`
	if exepected != got {
		t.Fatalf("exp: %s\ngot: %s", exepected, got)
	}

	got = executeAndGetJSON("drop measurement memory", executor)
	exepected = `[{}]`
	if exepected != got {
		t.Fatalf("exp: %s\ngot: %s", exepected, got)
	}

	validateDrop := func() {
		got = executeAndGetJSON("show series", executor)
		exepected = `[{"series":[{"name":"cpu","columns":["_key","host"],"values":[["cpu,host=server","server"]]}]}]`
		if exepected != got {
			t.Fatalf("exp: %s\ngot: %s", exepected, got)
		}
		got = executeAndGetJSON("show measurements", executor)
		exepected = `[{"series":[{"name":"measurements","columns":["name"],"values":[["cpu"]]}]}]`
		if exepected != got {
			t.Fatalf("exp: %s\ngot: %s", exepected, got)
		}
		got = executeAndGetJSON("select * from memory", executor)
		exepected = `[{"error":"measurement not found: \"foo\".\"foo\".memory"}]`
		if exepected != got {
			t.Fatalf("exp: %s\ngot: %s", exepected, got)
		}
	}

	validateDrop()
	store.Close()
	store = NewStore(store.path)
	store.Open()
	executor.store = store
	validateDrop()
}

// mock for the metaExecutor
type metaExec struct {
	fn func(stmt influxql.Statement) *influxql.Result
}

func (m *metaExec) ExecuteStatement(stmt influxql.Statement) *influxql.Result {
	return m.fn(stmt)
}

func TestDropDatabase(t *testing.T) {
	store, executor := testStoreAndExecutor()
	defer os.RemoveAll(store.path)

	pt := NewPoint(
		"cpu",
		map[string]string{"host": "server"},
		map[string]interface{}{"value": 1.0},
		time.Unix(1, 2),
	)

	err := store.WriteToShard(shardID, []Point{pt})
	if err != nil {
		t.Fatalf(err.Error())
	}

	got := executeAndGetJSON("select * from cpu", executor)
	expected := `[{"series":[{"name":"cpu","tags":{"host":"server"},"columns":["time","value"],"values":[["1970-01-01T00:00:01.000000002Z",1]]}]}]`
	if expected != got {
		t.Fatalf("exp: %s\ngot: %s", expected, got)
	}

	var name string
	me := &metaExec{fn: func(stmt influxql.Statement) *influxql.Result {
		name = stmt.(*influxql.DropDatabaseStatement).Name
		return &influxql.Result{}
	}}
	executor.MetaStatementExecutor = me

	// verify the database is there on disk
	dbPath := filepath.Join(store.path, "foo")
	if _, err := os.Stat(dbPath); err != nil {
		t.Fatalf("execpted database dir %s to exist", dbPath)
	}

	got = executeAndGetJSON("drop database foo", executor)
	expected = `[{}]`
	if got != expected {
		t.Fatalf("exp: %s\ngot: %s", expected, got)
	}

	if name != "foo" {
		t.Fatalf("expected the MetaStatementExecutor to be called with database name foo, but got %s", name)
	}

	if _, err := os.Stat(dbPath); !os.IsNotExist(err) {
		t.Fatalf("expected database dir %s to be gone", dbPath)
	}

	store.Close()
	store = NewStore(store.path)
	store.Open()
	executor.store = store

	err = store.WriteToShard(shardID, []Point{pt})
	if err == nil || err.Error() != "shard not found" {
		t.Fatalf("expected shard to not be found")
	}
}

// Ensure that queries for which there is no data result in an empty set.
func TestQueryNoData(t *testing.T) {
	store, executor := testStoreAndExecutor()
	defer os.RemoveAll(store.path)

	got := executeAndGetJSON("select * from /.*/", executor)
	expected := `[{}]`
	if expected != got {
		t.Fatalf("exp: %s\ngot: %s", expected, got)
	}

	got = executeAndGetJSON("show series", executor)
	expected = `[{}]`
	if expected != got {
		t.Fatalf("exp: %s\ngot: %s", expected, got)
	}

	store.Close()
}

// ensure that authenticate doesn't return an error if the user count is zero and they're attempting
// to create a user.
func TestAuthenticateIfUserCountZeroAndCreateUser(t *testing.T) {
	store, executor := testStoreAndExecutor()
	defer os.RemoveAll(store.path)
	ms := &testMetastore{userCount: 0}
	executor.MetaStore = ms

	if err := executor.Authorize(nil, mustParseQuery("create user foo with password 'asdf' with all privileges"), ""); err != nil {
		t.Fatalf("should have authenticated if no users and attempting to create a user but got error: %s", err.Error())
	}

	if executor.Authorize(nil, mustParseQuery("create user foo with password 'asdf'"), "") == nil {
		t.Fatalf("should have failed authentication if no user given and no users exist for create user query that doesn't grant all privileges")
	}

	if executor.Authorize(nil, mustParseQuery("select * from foo"), "") == nil {
		t.Fatalf("should have failed authentication if no user given and no users exist for any query other than create user")
	}

	ms.userCount = 1

	if executor.Authorize(nil, mustParseQuery("create user foo with password 'asdf'"), "") == nil {
		t.Fatalf("should have failed authentication if no user given and users exist")
	}

	if executor.Authorize(nil, mustParseQuery("select * from foo"), "") == nil {
		t.Fatalf("should have failed authentication if no user given and users exist")
	}
}

func testStoreAndExecutor() (*Store, *QueryExecutor) {
	path, _ := ioutil.TempDir("", "")

	store := NewStore(path)
	err := store.Open()
	if err != nil {
		panic(err)
	}
	database := "foo"
	retentionPolicy := "bar"
	shardID := uint64(1)
	store.CreateShard(database, retentionPolicy, shardID)

	executor := NewQueryExecutor(store)
	executor.MetaStore = &testMetastore{}

	return store, executor
}

func executeAndGetJSON(query string, executor *QueryExecutor) string {
	ch, err := executor.ExecuteQuery(mustParseQuery(query), "foo", 20)
	if err != nil {
		panic(err.Error())
	}

	var results []*influxql.Result
	for r := range ch {
		results = append(results, r)
	}
	return string(mustMarshalJSON(results))
}

type testMetastore struct {
	userCount int
}

func (t *testMetastore) Database(name string) (*meta.DatabaseInfo, error) {
	return &meta.DatabaseInfo{
		Name: name,
		DefaultRetentionPolicy: "foo",
		RetentionPolicies: []meta.RetentionPolicyInfo{
			{
				Name: "bar",
				ShardGroups: []meta.ShardGroupInfo{
					{
						ID:        uint64(1),
						StartTime: time.Now().Add(-time.Hour),
						EndTime:   time.Now().Add(time.Hour),
						Shards: []meta.ShardInfo{
							{
								ID:       uint64(1),
								OwnerIDs: []uint64{1},
							},
						},
					},
				},
			},
		},
	}, nil
}

func (t *testMetastore) Databases() ([]meta.DatabaseInfo, error) {
	db, _ := t.Database("foo")
	return []meta.DatabaseInfo{*db}, nil
}

func (t *testMetastore) User(name string) (*meta.UserInfo, error) { return nil, nil }

func (t *testMetastore) AdminUserExists() (bool, error) { return false, nil }

func (t *testMetastore) Authenticate(username, password string) (*meta.UserInfo, error) {
	return nil, nil
}

func (t *testMetastore) RetentionPolicy(database, name string) (rpi *meta.RetentionPolicyInfo, err error) {
	return &meta.RetentionPolicyInfo{
		Name: "bar",
		ShardGroups: []meta.ShardGroupInfo{
			{
				ID:        uint64(1),
				StartTime: time.Now().Add(-time.Hour),
				EndTime:   time.Now().Add(time.Hour),
				Shards: []meta.ShardInfo{
					{
						ID:       uint64(1),
						OwnerIDs: []uint64{1},
					},
				},
			},
		},
	}, nil
}

func (t *testMetastore) UserCount() (int, error) {
	return t.userCount, nil
}

// MustParseQuery parses an InfluxQL query. Panic on error.
func mustParseQuery(s string) *influxql.Query {
	q, err := influxql.NewParser(strings.NewReader(s)).ParseQuery()
	if err != nil {
		panic(err.Error())
	}
	return q
}
<|endoftext|>


<commit_before>/*
Copyright 2014 Rohith All rights reserved.

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 marathon

const (
	defaultEventsURL = "/event"

	/* --- api related constants --- */
	marathonAPIVersion      = "v2"
	marathonAPIEventStream  = marathonAPIVersion + "/events"
	marathonAPISubscription = marathonAPIVersion + "/eventSubscriptions"
	marathonAPIApps         = marathonAPIVersion + "/apps"
	marathonAPITasks        = marathonAPIVersion + "/tasks"
	marathonAPIDeployments  = marathonAPIVersion + "/deployments"
	marathonAPIGroups       = marathonAPIVersion + "/groups"
	marathonAPIQueue        = marathonAPIVersion + "/queue"
	marathonAPIInfo         = marathonAPIVersion + "/info"
	marathonAPILeader       = marathonAPIVersion + "/leader"
	marathonAPIPing         = "/ping"
)

const (
	// EventsTransportCallback activates callback events transport
	EventsTransportCallback EventsTransport = 1 << iota

	// EventsTransportSSE activates stream events transport
	EventsTransportSSE
)
<commit_msg>fixed ping path to make it work<commit_after>/*
Copyright 2014 Rohith All rights reserved.

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 marathon

const (
	defaultEventsURL = "/event"

	/* --- api related constants --- */
	marathonAPIVersion      = "v2"
	marathonAPIEventStream  = marathonAPIVersion + "/events"
	marathonAPISubscription = marathonAPIVersion + "/eventSubscriptions"
	marathonAPIApps         = marathonAPIVersion + "/apps"
	marathonAPITasks        = marathonAPIVersion + "/tasks"
	marathonAPIDeployments  = marathonAPIVersion + "/deployments"
	marathonAPIGroups       = marathonAPIVersion + "/groups"
	marathonAPIQueue        = marathonAPIVersion + "/queue"
	marathonAPIInfo         = marathonAPIVersion + "/info"
	marathonAPILeader       = marathonAPIVersion + "/leader"
	marathonAPIPing         = "ping"
)

const (
	// EventsTransportCallback activates callback events transport
	EventsTransportCallback EventsTransport = 1 << iota

	// EventsTransportSSE activates stream events transport
	EventsTransportSSE
)
<|endoftext|>


<commit_before>package main

const (
	// Name of this command.
	Name string = "logcatf"
	// Version of this command.
	Version string = "0.1.0"
	// DefaultFormat will be used if format wasn't specified.
	DefaultFormat string = "%time %tag %priority %message"
)

// Message has message strings.
var Message = map[string]string{
	"commandDescription": "A command to format Android Logcat",
	"helpFormat": `Format of output
 - Available keyword:
   %time, %tag, %priority, %pid, %tid, %message
 - Example:
   adb logcat -v time | "%time %message"
   adb logcat -v threadtime | "%time, %tag, %message" > logcat.csv
 - Default Format:
   "%time %tag %priority %message"`,
	"helpTrigger":           "Regex for triggering a command.",
	"helpCommand":           "Command to be executed.",
	"msgUnavailableKeyword": "error: %s is not available. Please check `Availavle Keyword:` on help.",
	"msgDuplicatedKeyword":  "error: %s or %s is duplicated.",
}
<commit_msg>update help message<commit_after>package main

const (
	// Name of this command.
	Name string = "logcatf"
	// Version of this command.
	Version string = "0.1.0"
	// DefaultFormat will be used if format wasn't specified.
	DefaultFormat string = "%time %tag %priority %message"
)

// Message has message strings.
var Message = map[string]string{
	"commandDescription": "A command to format Android Logcat",
	"helpFormat": `Format of output
 - Available keyword:
     %t (%time), %a (%tag), %p (%priority)
     %i (%pid), %I (%tid), %m (%message)
 - Example:
   adb logcat -v time | "%time %message"
   adb logcat -v threadtime | "%t, %a, %p, %i, %I, %m" > logcat.csv
 - Default Format:
   "%time %tag %priority %message"`,
	"helpTrigger":           "Regex for triggering a command.",
	"helpCommand":           "Command to be executed.",
	"msgUnavailableKeyword": "error: %s is not available. Please check `Availavle Keyword:` on help.",
	"msgDuplicatedKeyword":  "error: %s or %s is duplicated.",
}
<|endoftext|>


<commit_before>package hermes

import (
	"fmt"
	"net/http"
	"net/url"
	"runtime"
	"sync"
	"time"

	"strings"

	"github.com/PuerkitoBio/fetchbot"
	"github.com/PuerkitoBio/goquery"
)

var (
	mu sync.Mutex // Protect access to dup

	dup = map[string]bool{} // Duplicates table

	settingLinks = map[string]bool{} // Tracking link settings

	ingestionSet []Document // ingestion data TODO make non global

	badLinks []string // bad links TODO make non global
)

// Crawl function that will take a url string and start firing out some crawling functions
// it will return true/false based on the url root it starts with.
func Crawl(settings Settings, linkSettings CustomSettings, u url.URL) ([]Document, bool) {
	// Create the muxer
	mux := fetchbot.NewMux()

	// Handle all errors the same
	mux.HandleErrors(fetchbot.HandlerFunc(func(ctx *fetchbot.Context, res *http.Response, err error) {
		fmt.Printf("[ERR] %s %s - %s\n", ctx.Cmd.Method(), ctx.Cmd.URL(), err)
	}))

	// Handle GET requests for html responses, to parse the body and enqueue all links as HEAD
	// requests.
	mux.Response().Method("GET").ContentType("text/html").Handler(fetchbot.HandlerFunc(
		func(ctx *fetchbot.Context, res *http.Response, err error) {
			// Process the body to find the links
			doc, err := goquery.NewDocumentFromReader(res.Body)
			if err != nil {
				// find the bad links in the documents
				badLinks = append(badLinks, ctx.Cmd.URL().String())
				fmt.Printf("[ERR] %s %s - %s\n", ctx.Cmd.Method(), ctx.Cmd.URL(), err)
				return
			}
			// Enqueue all links as HEAD requests
			enqueueLinks(ctx, doc, u.Host, linkSettings)
		}))

	// Handle HEAD requests for html responses coming from the source host - we don't want
	// to crawl links from other hosts.
	mux.Response().Method("HEAD").Host(u.Host).ContentType("text/html").Handler(fetchbot.HandlerFunc(
		func(ctx *fetchbot.Context, res *http.Response, err error) {
			if _, err := ctx.Q.SendStringGet(ctx.Cmd.URL().String()); err != nil {
				fmt.Printf("[ERR] %s %s - %s\n", ctx.Cmd.Method(), ctx.Cmd.URL(), err)
			}
		}))

	// Create the Fetcher, handle the logging first, then dispatch to the Muxer
	h := scrapeHandler(mux, linkSettings)
	if settings.StopAtURL != "" || settings.CancelAtURL != "" {
		stopURL := settings.StopAtURL
		if settings.CancelAtURL != "" {
			stopURL = settings.CancelAtURL
		}
		h = stopHandler(stopURL, settings.CancelAtURL != "", scrapeHandler(mux, linkSettings))
	}
	f := fetchbot.New(h)

	// set the fetchbots settings from flag parameters
	f.UserAgent = settings.UserAgent
	f.CrawlDelay = settings.CrawlDelay * time.Second
	f.WorkerIdleTTL = settings.WorkerIdleTTL * time.Second
	f.AutoClose = settings.AutoClose

	// First mem stat print must be right after creating the fetchbot
	if settings.MemStatsInterval > 0 {
		// Print starting stats
		printMemStats(nil)
		// Run at regular intervals
		runMemStats(f, settings.MemStatsInterval)
		// On exit, print ending stats after a GC
		defer func() {
			runtime.GC()
			printMemStats(nil)
		}()
	}

	// Start processing
	q := f.Start()

	// if a stop or cancel is requested after some duration, launch the goroutine
	// that will stop or cancel.
	if settings.StopDuration*time.Minute > 0 || settings.CancelDuration*time.Minute > 0 {
		after := settings.StopDuration * time.Minute
		stopFunc := q.Close
		if settings.CancelDuration != 0 {
			after = settings.CancelDuration * time.Minute
			stopFunc = q.Cancel
		}

		go func() {
			c := time.After(after)
			<-c
			stopFunc()
		}()
	}

	// Enqueue the seed, which is the first entry in the dup map
	dup[linkSettings.RootLink] = true
	_, err := q.SendStringGet(linkSettings.RootLink)
	if err != nil {
		fmt.Printf("[ERR] GET %s - %s\n", linkSettings.RootLink, err)
	}
	q.Block()

	return ingestionSet, true
}

// stopHandler stops the fetcher if the stopurl is reached. Otherwise it dispatches
// the call to the wrapped Handler.
func stopHandler(stopurl string, cancel bool, wrapped fetchbot.Handler) fetchbot.Handler {
	return fetchbot.HandlerFunc(func(ctx *fetchbot.Context, res *http.Response, err error) {
		if ctx.Cmd.URL().String() == stopurl {
			fmt.Printf(">>>>> STOP URL %s\n", ctx.Cmd.URL())
			// generally not a good idea to stop/block from a handler goroutine
			// so do it in a separate goroutine
			go func() {
				if cancel {
					ctx.Q.Cancel()
				} else {
					ctx.Q.Close()
				}
			}()
			return
		}
		wrapped.Handle(ctx, res, err)
	})
}

// scrapeHandler will fire a scraper function on the page if successful response,
// append the scraped document stored for index ingestion
// and dispatches the call to the wrapped Handler.
func scrapeHandler(wrapped fetchbot.Handler, linkSettings CustomSettings) fetchbot.Handler {
	return fetchbot.HandlerFunc(func(ctx *fetchbot.Context, res *http.Response, err error) {
		if err == nil {
			if res.StatusCode == 200 {
				url := ctx.Cmd.URL().String()
				responseDocument, err := Scrape(ctx, linkSettings)
				if err != nil {
					fmt.Printf("[ERR] SCRAPE URL: %s - %s", url, err)
				}
				ingestionSet = append(ingestionSet, responseDocument)
			}
			fmt.Printf("[%d] %s %s - %s\n", res.StatusCode, ctx.Cmd.Method(), ctx.Cmd.URL(), res.Header.Get("Content-Type"))
		}
		wrapped.Handle(ctx, res, err)
	})
}

// enqueueLinks will make sure we are adding links to the queue to be processed for crawling/scraping
// this will pull all the href attributes on pages, check for duplicates and add them to the queue
func enqueueLinks(ctx *fetchbot.Context, doc *goquery.Document, host string, settings CustomSettings) {
	mu.Lock()

	fmt.Printf("incoming url: -- %s\n   ", ctx.Cmd.URL().String())
	fmt.Println("Duplicate object")
	for key, value := range dup {
		fmt.Printf("Link %s - %t \n", key, value)
	}

	doc.Find("a[href]").Each(func(i int, s *goquery.Selection) {
		val, exists := s.Attr("href")
		if exists == false {
			fmt.Print("error: address within the document\n")
			return
		}

		// Resolve address
		u, err := ctx.Cmd.URL().Parse(val)
		if err != nil {
			fmt.Printf("error: resolve URL %s - %s\n", u, err)
			return
		}

		// remove the 'www' from the URL so that we have better duplicate detection
		normalizeLink(u)

		// catch the duplicate urls here before trying to add them to the queue
		if !dup[u.String()] {

			// tld & subdomain
			if settings.TopLevelDomain == true && settings.Subdomain == true {
				rootDomain := getDomain(host)
				current := getDomain(u.Host)

				if rootDomain == current {
					err := addLink(ctx, u)
					if err != nil {
						fmt.Printf("error: enqueue head %s - %s\n", u, err)
						return
					}
				} else {
					fmt.Printf("catch: out of domain scope -- %s != %s\n", u.Host, host)
					return
				}

				return
			}

			// tld check
			if settings.TopLevelDomain == true {
				rootTLD := getTLD(host)
				current := getTLD(u.Host)

				if rootTLD == current {
					err := addLink(ctx, u)
					if err != nil {
						fmt.Printf("error: enqueue head %s - %s\n", u, err)
						return
					}
				}

				return
			}

			// subdomain check
			if settings.Subdomain == true {
				rootDomain := getDomain(host)
				current := getDomain(u.Host)

				if rootDomain == current {
					err := addLink(ctx, u)
					if err != nil {
						fmt.Printf("error: enqueue head %s - %s\n", u, err)
						return
					}
				} else {
					fmt.Printf("catch: out of domain scope -- %s != %s\n", u.Host, host)
					return
				}
			}

			return
		}
	})
	mu.Unlock()
}

// remove the www from the URL host
func normalizeLink(u *url.URL) {
	s := strings.Split(u.Host, ".")
	if len(s) == 0 {
		fmt.Printf("[ERR] URL doesn't have a TLD: %s\n", u.Host)
	} else if s[0] == "www" {
		u.Host = strings.Join(s[1:], ".")
	}
}

// addLink will add a url to fetchbot's queue and to the global hashmap to audit for duplicates
func addLink(ctx *fetchbot.Context, u *url.URL) error {
	if _, err := ctx.Q.SendStringHead(u.String()); err != nil {
		return err
	}
	dup[u.String()] = true
	return nil
}

// getDomain will parse a url and return the domain with the tld on it (ie. example.com)
func getDomain(u string) (root string) {
	s := strings.Split(u, ".")
	if len(s) == 0 {
		root = u
		return
	}
	last := len(s) - 1
	runnerUp := last - 1
	root = s[runnerUp] + "." + s[last]
	return
}

// getTLD will parse a url type and return the top-level domain (.com, .edu, .gov, etc.)
func getTLD(u string) (tld string) {
	s := strings.Split(u, ".")
	if len(s) == 0 {
		tld = u
		return
	}
	last := len(s) - 1
	tld = s[last]
	return
}
<commit_msg>added error debug in scrapeHandler<commit_after>package hermes

import (
	"fmt"
	"net/http"
	"net/url"
	"runtime"
	"sync"
	"time"

	"strings"

	"github.com/PuerkitoBio/fetchbot"
	"github.com/PuerkitoBio/goquery"
)

var (
	mu sync.Mutex // Protect access to dup

	dup = map[string]bool{} // Duplicates table

	settingLinks = map[string]bool{} // Tracking link settings

	ingestionSet []Document // ingestion data TODO make non global

	badLinks []string // bad links TODO make non global
)

// Crawl function that will take a url string and start firing out some crawling functions
// it will return true/false based on the url root it starts with.
func Crawl(settings Settings, linkSettings CustomSettings, u url.URL) ([]Document, bool) {
	// Create the muxer
	mux := fetchbot.NewMux()

	// Handle all errors the same
	mux.HandleErrors(fetchbot.HandlerFunc(func(ctx *fetchbot.Context, res *http.Response, err error) {
		fmt.Printf("[ERR] %s %s - %s\n", ctx.Cmd.Method(), ctx.Cmd.URL(), err)
	}))

	// Handle GET requests for html responses, to parse the body and enqueue all links as HEAD
	// requests.
	mux.Response().Method("GET").ContentType("text/html").Handler(fetchbot.HandlerFunc(
		func(ctx *fetchbot.Context, res *http.Response, err error) {
			// Process the body to find the links
			doc, err := goquery.NewDocumentFromReader(res.Body)
			if err != nil {
				// find the bad links in the documents
				badLinks = append(badLinks, ctx.Cmd.URL().String())
				fmt.Printf("[ERR] %s %s - %s\n", ctx.Cmd.Method(), ctx.Cmd.URL(), err)
				return
			}
			// Enqueue all links as HEAD requests
			enqueueLinks(ctx, doc, u.Host, linkSettings)
		}))

	// Handle HEAD requests for html responses coming from the source host - we don't want
	// to crawl links from other hosts.
	mux.Response().Method("HEAD").Host(u.Host).ContentType("text/html").Handler(fetchbot.HandlerFunc(
		func(ctx *fetchbot.Context, res *http.Response, err error) {
			if _, err := ctx.Q.SendStringGet(ctx.Cmd.URL().String()); err != nil {
				fmt.Printf("[ERR] %s %s - %s\n", ctx.Cmd.Method(), ctx.Cmd.URL(), err)
			}
		}))

	// Create the Fetcher, handle the logging first, then dispatch to the Muxer
	h := scrapeHandler(mux, linkSettings)
	if settings.StopAtURL != "" || settings.CancelAtURL != "" {
		stopURL := settings.StopAtURL
		if settings.CancelAtURL != "" {
			stopURL = settings.CancelAtURL
		}
		h = stopHandler(stopURL, settings.CancelAtURL != "", scrapeHandler(mux, linkSettings))
	}
	f := fetchbot.New(h)

	// set the fetchbots settings from flag parameters
	f.UserAgent = settings.UserAgent
	f.CrawlDelay = settings.CrawlDelay * time.Second
	f.WorkerIdleTTL = settings.WorkerIdleTTL * time.Second
	f.AutoClose = settings.AutoClose

	// First mem stat print must be right after creating the fetchbot
	if settings.MemStatsInterval > 0 {
		// Print starting stats
		printMemStats(nil)
		// Run at regular intervals
		runMemStats(f, settings.MemStatsInterval)
		// On exit, print ending stats after a GC
		defer func() {
			runtime.GC()
			printMemStats(nil)
		}()
	}

	// Start processing
	q := f.Start()

	// if a stop or cancel is requested after some duration, launch the goroutine
	// that will stop or cancel.
	if settings.StopDuration*time.Minute > 0 || settings.CancelDuration*time.Minute > 0 {
		after := settings.StopDuration * time.Minute
		stopFunc := q.Close
		if settings.CancelDuration != 0 {
			after = settings.CancelDuration * time.Minute
			stopFunc = q.Cancel
		}

		go func() {
			c := time.After(after)
			<-c
			stopFunc()
		}()
	}

	// Enqueue the seed, which is the first entry in the dup map
	dup[linkSettings.RootLink] = true
	_, err := q.SendStringGet(linkSettings.RootLink)
	if err != nil {
		fmt.Printf("[ERR] GET %s - %s\n", linkSettings.RootLink, err)
	}
	q.Block()

	return ingestionSet, true
}

// stopHandler stops the fetcher if the stopurl is reached. Otherwise it dispatches
// the call to the wrapped Handler.
func stopHandler(stopurl string, cancel bool, wrapped fetchbot.Handler) fetchbot.Handler {
	return fetchbot.HandlerFunc(func(ctx *fetchbot.Context, res *http.Response, err error) {
		if ctx.Cmd.URL().String() == stopurl {
			fmt.Printf(">>>>> STOP URL %s\n", ctx.Cmd.URL())
			// generally not a good idea to stop/block from a handler goroutine
			// so do it in a separate goroutine
			go func() {
				if cancel {
					ctx.Q.Cancel()
				} else {
					ctx.Q.Close()
				}
			}()
			return
		}
		wrapped.Handle(ctx, res, err)
	})
}

// scrapeHandler will fire a scraper function on the page if successful response,
// append the scraped document stored for index ingestion
// and dispatches the call to the wrapped Handler.
func scrapeHandler(wrapped fetchbot.Handler, linkSettings CustomSettings) fetchbot.Handler {
	return fetchbot.HandlerFunc(func(ctx *fetchbot.Context, res *http.Response, err error) {
		if err == nil {
			if res.StatusCode == 200 {
				url := ctx.Cmd.URL().String()
				responseDocument, err := Scrape(ctx, linkSettings)
				if err != nil {
					fmt.Printf("[ERR] SCRAPE URL: %s - %s", url, err)
				}
				ingestionSet = append(ingestionSet, responseDocument)
			}
			fmt.Printf("[%d] %s %s - %s\n", res.StatusCode, ctx.Cmd.Method(), ctx.Cmd.URL(), res.Header.Get("Content-Type"))
		} else {
			fmt.Printf("ERR [%d] - %s -- %s\n", res.StatusCode, ctx.Cmd.URL(), err)
		}
		wrapped.Handle(ctx, res, err)
	})
}

// enqueueLinks will make sure we are adding links to the queue to be processed for crawling/scraping
// this will pull all the href attributes on pages, check for duplicates and add them to the queue
func enqueueLinks(ctx *fetchbot.Context, doc *goquery.Document, host string, settings CustomSettings) {
	mu.Lock()

	fmt.Printf("incoming url: -- %s\n   ", ctx.Cmd.URL().String())
	fmt.Println("Duplicate object")
	for key, value := range dup {
		fmt.Printf("Link %s - %t \n", key, value)
	}

	doc.Find("a[href]").Each(func(i int, s *goquery.Selection) {
		val, exists := s.Attr("href")
		if exists == false {
			fmt.Print("error: address within the document\n")
			return
		}

		// Resolve address
		u, err := ctx.Cmd.URL().Parse(val)
		if err != nil {
			fmt.Printf("error: resolve URL %s - %s\n", u, err)
			return
		}

		// remove the 'www' from the URL so that we have better duplicate detection
		normalizeLink(u)

		// catch the duplicate urls here before trying to add them to the queue
		if !dup[u.String()] {

			// tld & subdomain
			if settings.TopLevelDomain == true && settings.Subdomain == true {
				rootDomain := getDomain(host)
				current := getDomain(u.Host)

				if rootDomain == current {
					err := addLink(ctx, u)
					if err != nil {
						fmt.Printf("error: enqueue head %s - %s\n", u, err)
						return
					}
				} else {
					fmt.Printf("catch: out of domain scope -- %s != %s\n", u.Host, host)
					return
				}

				return
			}

			// tld check
			if settings.TopLevelDomain == true {
				rootTLD := getTLD(host)
				current := getTLD(u.Host)

				if rootTLD == current {
					err := addLink(ctx, u)
					if err != nil {
						fmt.Printf("error: enqueue head %s - %s\n", u, err)
						return
					}
				}

				return
			}

			// subdomain check
			if settings.Subdomain == true {
				rootDomain := getDomain(host)
				current := getDomain(u.Host)

				if rootDomain == current {
					err := addLink(ctx, u)
					if err != nil {
						fmt.Printf("error: enqueue head %s - %s\n", u, err)
						return
					}
				} else {
					fmt.Printf("catch: out of domain scope -- %s != %s\n", u.Host, host)
					return
				}
			}

			return
		}
	})
	mu.Unlock()
}

// remove the www from the URL host
func normalizeLink(u *url.URL) {
	s := strings.Split(u.Host, ".")
	if len(s) == 0 {
		fmt.Printf("[ERR] URL doesn't have a TLD: %s\n", u.Host)
	} else if s[0] == "www" {
		u.Host = strings.Join(s[1:], ".")
	}
}

// addLink will add a url to fetchbot's queue and to the global hashmap to audit for duplicates
func addLink(ctx *fetchbot.Context, u *url.URL) error {
	if _, err := ctx.Q.SendStringHead(u.String()); err != nil {
		return err
	}
	dup[u.String()] = true
	return nil
}

// getDomain will parse a url and return the domain with the tld on it (ie. example.com)
func getDomain(u string) (root string) {
	s := strings.Split(u, ".")
	if len(s) == 0 {
		root = u
		return
	}
	last := len(s) - 1
	runnerUp := last - 1
	root = s[runnerUp] + "." + s[last]
	return
}

// getTLD will parse a url type and return the top-level domain (.com, .edu, .gov, etc.)
func getTLD(u string) (tld string) {
	s := strings.Split(u, ".")
	if len(s) == 0 {
		tld = u
		return
	}
	last := len(s) - 1
	tld = s[last]
	return
}
<|endoftext|>