<commit_before>package router

import (
	"bufio"
	"errors"
	"io"
	"log"
	"os"
	"strings"
	"sync"
	"time"
	"regexp"
	"strconv"

	"github.com/fsouza/go-dockerclient"
)

var time_regexp_map = make(map[string]*regexp.Regexp)

func init() {
	//time_regexp = regexp.MustCompile(`^(\d{4})[-,/](\d{2})[-,/](\d{2})`)
	pump := &LogsPump{
		pumps:  make(map[string]*containerPump),
		routes: make(map[chan *update]struct{}),
	}
	LogRouters.Register(pump, "pump")
	Jobs.Register(pump, "pump")
}

func getopt(name, dfault string) string {
	value := os.Getenv(name)
	if value == "" {
		value = dfault
	}
	return value
}

func debug(v ...interface{}) {
	if os.Getenv("DEBUG") != "" {
		log.Println(v...)
	}
}

func assert(err error, context string) {
	if err != nil {
		log.Fatal(context+": ", err)
	}
}

func normalName(name string) string {
	return name[1:]
}

func normalID(id string) string {
	if len(id) > 12 {
		return id[:12]
	}
	return id
}

func logDriverSupported(container *docker.Container) bool {
	switch container.HostConfig.LogConfig.Type {
	case "json-file", "journald":
		return true
	default:
		return false
	}
}

func ignoreContainer(container *docker.Container) bool {
	for _, kv := range container.Config.Env {
		kvp := strings.SplitN(kv, "=", 2)
		if len(kvp) == 2 && kvp[0] == "LOGSPOUT" && strings.ToLower(kvp[1]) == "ignore" {
			return true
		}
	}
	excludeLabel := getopt("EXCLUDE_LABEL", "")
	if value, ok := container.Config.Labels[excludeLabel]; ok {
		return len(excludeLabel) > 0 && strings.ToLower(value) == "true"
	}
	return false
}

type update struct {
	*docker.APIEvents
	pump *containerPump
}

type LogsPump struct {
	mu     sync.Mutex
	pumps  map[string]*containerPump
	routes map[chan *update]struct{}
	client *docker.Client
}

func (p *LogsPump) Name() string {
	return "pump"
}

func (p *LogsPump) Setup() error {
	var err error
	p.client, err = docker.NewClientFromEnv()
	return err
}

func (p *LogsPump) rename(event *docker.APIEvents) {
	p.mu.Lock()
	defer p.mu.Unlock()
	container, err := p.client.InspectContainer(event.ID)
	assert(err, "pump")
	pump, ok := p.pumps[normalID(event.ID)]
	if !ok {
		debug("pump.rename(): ignore: pump not found, state:", container.State.StateString())
		return
	}
	pump.container.Name = container.Name
}

func (p *LogsPump) Run() error {
	containers, err := p.client.ListContainers(docker.ListContainersOptions{})
	if err != nil {
		return err
	}
	for _, listing := range containers {
		p.pumpLogs(&docker.APIEvents{
			ID:     normalID(listing.ID),
			Status: "start",
		}, false)
	}
	events := make(chan *docker.APIEvents)
	err = p.client.AddEventListener(events)
	if err != nil {
		return err
	}
	for event := range events {
		debug("pump.Run() event:", normalID(event.ID), event.Status)
		switch event.Status {
		case "start", "restart":
			go p.pumpLogs(event, true)
		case "rename":
			go p.rename(event)
		case "die":
			go p.update(event)
		}
	}
	return errors.New("docker event stream closed")
}

func (p *LogsPump) pumpLogs(event *docker.APIEvents, backlog bool) {
	id := normalID(event.ID)
	container, err := p.client.InspectContainer(id)
	assert(err, "pump")
	//add route

	filter := getopt("FILTER", "TOPIC")
	topic_env := getEnv(filter, container.Config.Env, "default")
	kafka_addr := getopt("KAFKA", "localhost:9092")
	if _, ok := time_regexp_map[topic_env]; !ok {
		regex_env := getEnv("LOG_TIME_REGEX", container.Config.Env, "^(\\d{4})[-,/](\\d{2})[-,/](\\d{2})")
		time_regexp_map[topic_env] = regexp.MustCompile(regex_env)
	}
	if kafka_addr != "" {
		route := &Route{
			ID:        topic_env,
			Adapter:   "kafka",
			FilterEnv: topic_env,
			Address:   kafka_addr + "/" + topic_env,
		}
		Routes.Add(route)
	  log.Println("add route kafka topic:", kafka_addr, topic_env)
	}
	if container.Config.Tty {
		debug("pump.pumpLogs():", id, "ignored: tty enabled")
		return
	}
	if ignoreContainer(container) {
		debug("pump.pumpLogs():", id, "ignored: environ ignore")
		return
	}
	if !logDriverSupported(container) {
		debug("pump.pumpLogs():", id, "ignored: log driver not supported")
		return
	}

	var sinceTime time.Time
	if backlog {
		sinceTime = time.Unix(0, 0)
	} else {
		str_sinceTime := os.Getenv("SINCE_TIME")//first read sincetime from file via env
		int_sinceTime, _ := strconv.ParseInt(str_sinceTime, 10, 64)
		sinceTime = time.Unix(int_sinceTime, 0)
	}

	p.mu.Lock()
	if _, exists := p.pumps[id]; exists {
		p.mu.Unlock()
		debug("pump.pumpLogs():", id, "pump exists")
		return
	}
	outrd, outwr := io.Pipe()
	errrd, errwr := io.Pipe()
	//containerPump use to recieve log from pump
	p.pumps[id] = newContainerPump(container, outrd, errrd)
	p.mu.Unlock()
	p.update(event)
	go func() {
		for {
			debug("pump.pumpLogs():", id, "started")
			err := p.client.Logs(docker.LogsOptions{
				Container:    id,
				OutputStream: outwr,
				ErrorStream:  errwr,
				Stdout:       true,
				Stderr:       true,
				Follow:       true,
				Tail:         "all",
				Since:        sinceTime.Unix(),
			})
			if err != nil {
				debug("pump.pumpLogs():", id, "stopped with error:", err)
			} else {
				debug("pump.pumpLogs():", id, "stopped")
			}

			sinceTime = time.Now()
			container, err := p.client.InspectContainer(id)
			if err != nil {
				_, four04 := err.(*docker.NoSuchContainer)
				if !four04 {
					assert(err, "pump")
				}
			} else if container.State.Running {
				continue
			}

			debug("pump.pumpLogs():", id, "dead")
			outwr.Close()
			errwr.Close()
			p.mu.Lock()
			delete(p.pumps, id)
			p.mu.Unlock()
			return
		}
	}()
}

func (p *LogsPump) update(event *docker.APIEvents) {
	p.mu.Lock()
	defer p.mu.Unlock()
	pump, pumping := p.pumps[normalID(event.ID)]
	if pumping {
		for r := range p.routes {
			select {
			case r <- &update{event, pump}:
			case <-time.After(time.Second * 1):
				debug("pump.update(): route timeout, dropping")
				defer delete(p.routes, r)
			}
		}
	}
}

func (p *LogsPump) RoutingFrom(id string) bool {
	p.mu.Lock()
	defer p.mu.Unlock()
	_, monitoring := p.pumps[normalID(id)]
	return monitoring
}

func (p *LogsPump) Route(route *Route, logstream chan *Message) {
	p.mu.Lock()
	for _, pump := range p.pumps {
		if route.MatchContainer(
			normalID(pump.container.ID),
			normalName(pump.container.Name),
			pump.container.Config.Env) {

			pump.add(logstream, route)
			defer pump.remove(logstream)
		}
		//add MatchEnv
	}
	updates := make(chan *update)
	p.routes[updates] = struct{}{}
	p.mu.Unlock()
	defer func() {
		p.mu.Lock()
		delete(p.routes, updates)
		p.mu.Unlock()
	}()
	for {
		select {
		case event := <-updates:
			switch event.Status {
			case "start", "restart":
				if route.MatchContainer(
					normalID(event.pump.container.ID),
					normalName(event.pump.container.Name),
					event.pump.container.Config.Env) {
					event.pump.add(logstream, route)
					defer event.pump.remove(logstream)
				}
			case "die":
				if strings.HasPrefix(route.FilterID, event.ID) {
					// If the route is just about a single container,
					// we can stop routing when it dies.
					return
				}
			}
		case <-route.Closer():
			return
		}
	}
}

type containerPump struct {
	sync.Mutex
	container  *docker.Container
	logstreams map[chan *Message]*Route
}

func newContainerPump(container *docker.Container, stdout, stderr io.Reader) *containerPump {
	cp := &containerPump{
		container:  container,
		logstreams: make(map[chan *Message]*Route),
	}
	filter := getopt("FILTER", "TOPIC")
	topic_env := getEnv(filter, container.Config.Env, "default")
	time_regexp, _ := time_regexp_map[topic_env]
	pump := func(source string, input io.Reader) {
		buf := bufio.NewReader(input)
		var line string
		line_count := 0
		line, err := buf.ReadString('\n')//add exception judge,read one line
		if err != nil {
			if err != io.EOF {
				debug("pump.newContainerPump():", normalID(container.ID), source+":", err)
			}
			return
		}
		for {
			secline, err := buf.ReadString('\n')//add exception judge,read one line
			if err != nil {
				if err != io.EOF {
					debug("pump.newContainerPump():", normalID(container.ID), source+":", err)
					return
				}
				cp.send(&Message{
					Data:      strings.TrimSuffix(line, "\n"),
					Container: container,
					Time:      time.Now(),
					Source:    source,
				})
				return
			}
			if !time_regexp.MatchString(secline) {
				line = line + secline
				line_count = line_count + 1
				if line_count < 100{
					continue
				}
			}
			cp.send(&Message{
				Data:      strings.TrimSuffix(line, "\n"),
				Container: container,
				Time:      time.Now(),
				Source:    source,
			})
			line = secline
		}
	}
	go pump("stdout", stdout)
	go pump("stderr", stderr)
	return cp
}

func (cp *containerPump) send(msg *Message) {
	cp.Lock()
	defer cp.Unlock()
	for logstream, route := range cp.logstreams {
		if !route.MatchMessage(msg) {
			continue
		}
		select {
		case logstream <- msg:
		case <-time.After(time.Second * 1):
			debug("pump.send(): send timeout, closing")
			// normal call to remove() triggered by
			// route.Closer() may not be able to grab
			// lock under heavy load, so we delete here
			defer delete(cp.logstreams, logstream)
		}
	}
}

func (cp *containerPump) add(logstream chan *Message, route *Route) {
	cp.Lock()
	defer cp.Unlock()
	cp.logstreams[logstream] = route
}

func (cp *containerPump) remove(logstream chan *Message) {
	cp.Lock()
	defer cp.Unlock()
	delete(cp.logstreams, logstream)
}
<commit_msg>change mustcompile to compile<commit_after>package router

import (
	"bufio"
	"errors"
	"io"
	"log"
	"os"
	"strings"
	"sync"
	"time"
	"regexp"
	"strconv"

	"github.com/fsouza/go-dockerclient"
)

var time_regexp_map = make(map[string]*regexp.Regexp)

func init() {
	//time_regexp = regexp.MustCompile(`^(\d{4})[-,/](\d{2})[-,/](\d{2})`)
	pump := &LogsPump{
		pumps:  make(map[string]*containerPump),
		routes: make(map[chan *update]struct{}),
	}
	LogRouters.Register(pump, "pump")
	Jobs.Register(pump, "pump")
}

func getopt(name, dfault string) string {
	value := os.Getenv(name)
	if value == "" {
		value = dfault
	}
	return value
}

func debug(v ...interface{}) {
	if os.Getenv("DEBUG") != "" {
		log.Println(v...)
	}
}

func assert(err error, context string) {
	if err != nil {
		log.Fatal(context+": ", err)
	}
}

func normalName(name string) string {
	return name[1:]
}

func normalID(id string) string {
	if len(id) > 12 {
		return id[:12]
	}
	return id
}

func logDriverSupported(container *docker.Container) bool {
	switch container.HostConfig.LogConfig.Type {
	case "json-file", "journald":
		return true
	default:
		return false
	}
}

func ignoreContainer(container *docker.Container) bool {
	for _, kv := range container.Config.Env {
		kvp := strings.SplitN(kv, "=", 2)
		if len(kvp) == 2 && kvp[0] == "LOGSPOUT" && strings.ToLower(kvp[1]) == "ignore" {
			return true
		}
	}
	excludeLabel := getopt("EXCLUDE_LABEL", "")
	if value, ok := container.Config.Labels[excludeLabel]; ok {
		return len(excludeLabel) > 0 && strings.ToLower(value) == "true"
	}
	return false
}

type update struct {
	*docker.APIEvents
	pump *containerPump
}

type LogsPump struct {
	mu     sync.Mutex
	pumps  map[string]*containerPump
	routes map[chan *update]struct{}
	client *docker.Client
}

func (p *LogsPump) Name() string {
	return "pump"
}

func (p *LogsPump) Setup() error {
	var err error
	p.client, err = docker.NewClientFromEnv()
	return err
}

func (p *LogsPump) rename(event *docker.APIEvents) {
	p.mu.Lock()
	defer p.mu.Unlock()
	container, err := p.client.InspectContainer(event.ID)
	assert(err, "pump")
	pump, ok := p.pumps[normalID(event.ID)]
	if !ok {
		debug("pump.rename(): ignore: pump not found, state:", container.State.StateString())
		return
	}
	pump.container.Name = container.Name
}

func (p *LogsPump) Run() error {
	containers, err := p.client.ListContainers(docker.ListContainersOptions{})
	if err != nil {
		return err
	}
	for _, listing := range containers {
		p.pumpLogs(&docker.APIEvents{
			ID:     normalID(listing.ID),
			Status: "start",
		}, false)
	}
	events := make(chan *docker.APIEvents)
	err = p.client.AddEventListener(events)
	if err != nil {
		return err
	}
	for event := range events {
		debug("pump.Run() event:", normalID(event.ID), event.Status)
		switch event.Status {
		case "start", "restart":
			go p.pumpLogs(event, true)
		case "rename":
			go p.rename(event)
		case "die":
			go p.update(event)
		}
	}
	return errors.New("docker event stream closed")
}

func (p *LogsPump) pumpLogs(event *docker.APIEvents, backlog bool) {
	id := normalID(event.ID)
	container, err := p.client.InspectContainer(id)
	assert(err, "pump")
	//add route

	filter := getopt("FILTER", "TOPIC")
	topic_env := getEnv(filter, container.Config.Env, "default")
	kafka_addr := getopt("KAFKA", "localhost:9092")
	if _, ok := time_regexp_map[topic_env]; !ok {
		regex_env := getEnv("LOG_TIME_REGEX", container.Config.Env, "^(\\d{4})[-,/](\\d{2})[-,/](\\d{2})")
		time_regexp_map[topic_env], err = regexp.Compile(regex_env)
		if err != nil {
			log.Fatal("LOG_TIME_REGEX set wrong:", err)
			return
		}
	}
	if kafka_addr != "" {
		route := &Route{
			ID:        topic_env,
			Adapter:   "kafka",
			FilterEnv: topic_env,
			Address:   kafka_addr + "/" + topic_env,
		}
		Routes.Add(route)
	  log.Println("add route kafka topic:", kafka_addr, topic_env)
	}
	if container.Config.Tty {
		debug("pump.pumpLogs():", id, "ignored: tty enabled")
		return
	}
	if ignoreContainer(container) {
		debug("pump.pumpLogs():", id, "ignored: environ ignore")
		return
	}
	if !logDriverSupported(container) {
		debug("pump.pumpLogs():", id, "ignored: log driver not supported")
		return
	}

	var sinceTime time.Time
	if backlog {
		sinceTime = time.Unix(0, 0)
	} else {
		str_sinceTime := os.Getenv("SINCE_TIME")//first read sincetime from file via env
		int_sinceTime, _ := strconv.ParseInt(str_sinceTime, 10, 64)
		sinceTime = time.Unix(int_sinceTime, 0)
	}

	p.mu.Lock()
	if _, exists := p.pumps[id]; exists {
		p.mu.Unlock()
		debug("pump.pumpLogs():", id, "pump exists")
		return
	}
	outrd, outwr := io.Pipe()
	errrd, errwr := io.Pipe()
	//containerPump use to recieve log from pump
	p.pumps[id] = newContainerPump(container, outrd, errrd)
	p.mu.Unlock()
	p.update(event)
	go func() {
		for {
			debug("pump.pumpLogs():", id, "started")
			err := p.client.Logs(docker.LogsOptions{
				Container:    id,
				OutputStream: outwr,
				ErrorStream:  errwr,
				Stdout:       true,
				Stderr:       true,
				Follow:       true,
				Tail:         "all",
				Since:        sinceTime.Unix(),
			})
			if err != nil {
				debug("pump.pumpLogs():", id, "stopped with error:", err)
			} else {
				debug("pump.pumpLogs():", id, "stopped")
			}

			sinceTime = time.Now()
			container, err := p.client.InspectContainer(id)
			if err != nil {
				_, four04 := err.(*docker.NoSuchContainer)
				if !four04 {
					assert(err, "pump")
				}
			} else if container.State.Running {
				continue
			}

			debug("pump.pumpLogs():", id, "dead")
			outwr.Close()
			errwr.Close()
			p.mu.Lock()
			delete(p.pumps, id)
			p.mu.Unlock()
			return
		}
	}()
}

func (p *LogsPump) update(event *docker.APIEvents) {
	p.mu.Lock()
	defer p.mu.Unlock()
	pump, pumping := p.pumps[normalID(event.ID)]
	if pumping {
		for r := range p.routes {
			select {
			case r <- &update{event, pump}:
			case <-time.After(time.Second * 1):
				debug("pump.update(): route timeout, dropping")
				defer delete(p.routes, r)
			}
		}
	}
}

func (p *LogsPump) RoutingFrom(id string) bool {
	p.mu.Lock()
	defer p.mu.Unlock()
	_, monitoring := p.pumps[normalID(id)]
	return monitoring
}

func (p *LogsPump) Route(route *Route, logstream chan *Message) {
	p.mu.Lock()
	for _, pump := range p.pumps {
		if route.MatchContainer(
			normalID(pump.container.ID),
			normalName(pump.container.Name),
			pump.container.Config.Env) {

			pump.add(logstream, route)
			defer pump.remove(logstream)
		}
		//add MatchEnv
	}
	updates := make(chan *update)
	p.routes[updates] = struct{}{}
	p.mu.Unlock()
	defer func() {
		p.mu.Lock()
		delete(p.routes, updates)
		p.mu.Unlock()
	}()
	for {
		select {
		case event := <-updates:
			switch event.Status {
			case "start", "restart":
				if route.MatchContainer(
					normalID(event.pump.container.ID),
					normalName(event.pump.container.Name),
					event.pump.container.Config.Env) {
					event.pump.add(logstream, route)
					defer event.pump.remove(logstream)
				}
			case "die":
				if strings.HasPrefix(route.FilterID, event.ID) {
					// If the route is just about a single container,
					// we can stop routing when it dies.
					return
				}
			}
		case <-route.Closer():
			return
		}
	}
}

type containerPump struct {
	sync.Mutex
	container  *docker.Container
	logstreams map[chan *Message]*Route
}

func newContainerPump(container *docker.Container, stdout, stderr io.Reader) *containerPump {
	cp := &containerPump{
		container:  container,
		logstreams: make(map[chan *Message]*Route),
	}
	filter := getopt("FILTER", "TOPIC")
	topic_env := getEnv(filter, container.Config.Env, "default")
	time_regexp, _ := time_regexp_map[topic_env]
	pump := func(source string, input io.Reader) {
		buf := bufio.NewReader(input)
		var line string
		line_count := 0
		line, err := buf.ReadString('\n')//add exception judge,read one line
		if err != nil {
			if err != io.EOF {
				debug("pump.newContainerPump():", normalID(container.ID), source+":", err)
			}
			return
		}
		for {
			secline, err := buf.ReadString('\n')//add exception judge,read one line
			if err != nil {
				if err != io.EOF {
					debug("pump.newContainerPump():", normalID(container.ID), source+":", err)
					return
				}
				cp.send(&Message{
					Data:      strings.TrimSuffix(line, "\n"),
					Container: container,
					Time:      time.Now(),
					Source:    source,
				})
				return
			}
			if !time_regexp.MatchString(secline) {
				line = line + secline
				line_count = line_count + 1
				if line_count < 100{
					continue
				}
			}
			cp.send(&Message{
				Data:      strings.TrimSuffix(line, "\n"),
				Container: container,
				Time:      time.Now(),
				Source:    source,
			})
			line = secline
		}
	}
	go pump("stdout", stdout)
	go pump("stderr", stderr)
	return cp
}

func (cp *containerPump) send(msg *Message) {
	cp.Lock()
	defer cp.Unlock()
	for logstream, route := range cp.logstreams {
		if !route.MatchMessage(msg) {
			continue
		}
		select {
		case logstream <- msg:
		case <-time.After(time.Second * 1):
			debug("pump.send(): send timeout, closing")
			// normal call to remove() triggered by
			// route.Closer() may not be able to grab
			// lock under heavy load, so we delete here
			defer delete(cp.logstreams, logstream)
		}
	}
}

func (cp *containerPump) add(logstream chan *Message, route *Route) {
	cp.Lock()
	defer cp.Unlock()
	cp.logstreams[logstream] = route
}

func (cp *containerPump) remove(logstream chan *Message) {
	cp.Lock()
	defer cp.Unlock()
	delete(cp.logstreams, logstream)
}
<|endoftext|>


<commit_before>package routes

import (
  "github.com/gillesdemey/npm-registry/auth"
  "github.com/gillesdemey/npm-registry/model"
  "github.com/gillesdemey/npm-registry/storage"
  "gopkg.in/gin-gonic/gin.v1"
  "net/http"
  "log"
  "regexp"
)

func Login(c *gin.Context) {
  var login model.Login
  var err error

  if err = c.BindJSON(&login); err != nil {
    c.JSON(http.StatusBadRequest, gin.H{"error": "bad request"})
    return
  }

  storage := c.Value("storage").(storage.StorageEngine)

  username := login.Username
  password := login.Password

  token, err := auth.Login(username, password)
  if err != nil {
    c.JSON(http.StatusUnauthorized, gin.H{"error": "invalid credentials"})
    return
  }

  err = storage.StoreUserToken(token, username)
  if err != nil {
    c.Status(http.StatusInternalServerError)
    return
  }

  c.JSON(http.StatusCreated, gin.H{"token": token})
}

// Return the username associated with the NPM token
func Whoami(c *gin.Context) {
  re := regexp.MustCompile("(?i)Bearer ")
  authHeader := c.Request.Header.Get("Authorization")
  token := re.ReplaceAllString(authHeader, "")
  log.Printf("Whoami request for token '%s'", token)

  storage := c.Value("storage").(storage.StorageEngine)
  username, err := storage.RetrieveUsernameFromToken(token)

  if err != nil {
    c.Status(http.StatusUnauthorized)
    return
  }

  c.JSON(http.StatusOK, gin.H{"username": username})
}
<commit_msg>✨ cleanup<commit_after>package routes

import (
  "github.com/gillesdemey/npm-registry/auth"
  "github.com/gillesdemey/npm-registry/model"
  "github.com/gillesdemey/npm-registry/storage"
  "gopkg.in/gin-gonic/gin.v1"
  "net/http"
  "log"
  "regexp"
)

// Create or verify a user named <username>
func Login(c *gin.Context) {
  var login model.Login
  var err error

  if err = c.BindJSON(&login); err != nil {
    c.JSON(http.StatusBadRequest, gin.H{"error": "bad request"})
    return
  }

  storage := c.Value("storage").(storage.StorageEngine)

  username := login.Username
  password := login.Password

  token, err := auth.Login(username, password)
  if err != nil {
    c.JSON(http.StatusUnauthorized, gin.H{"error": "invalid credentials"})
    return
  }

  err = storage.StoreUserToken(token, username)
  if err != nil {
    c.Status(http.StatusInternalServerError)
    return
  }

  c.JSON(http.StatusCreated, gin.H{"token": token})
}

// Return the username associated with the NPM token
func Whoami(c *gin.Context) {
  re := regexp.MustCompile("(?i)Bearer ")
  authHeader := c.Request.Header.Get("Authorization")
  token := re.ReplaceAllString(authHeader, "")

  log.Printf("Whoami request for token '%s'", token)

  storage := c.Value("storage").(storage.StorageEngine)
  username, err := storage.RetrieveUsernameFromToken(token)

  if err != nil {
    c.Status(http.StatusUnauthorized)
    return
  }

  c.JSON(http.StatusOK, gin.H{"username": username})
}
<|endoftext|>


<commit_before>package main

import (
        "fmt"
)

func main() {
        fmt.Printf("yo!!")
}
<commit_msg>make go version of 'yo' consistent with C & ooc versions<commit_after>package main

import (
  "fmt"
  "time"
)

func main() {
  fmt.Printf("yo!!\n")
  time.Sleep(time.Second)
}
<|endoftext|>


<commit_before>package stash

import (
	"fmt"

	"bazil.org/bazil/cas"
	"bazil.org/bazil/cas/chunks"
	"bazil.org/bazil/idpool"
)

// New creates a new Stash.
func New(bs chunks.Store) *Stash {
	s := &Stash{
		chunks: bs,
		local:  make(map[uint64]*chunks.Chunk),
	}
	return s
}

// Stash is a proxy for a chunks.Store, but it keeps Private Keys
// local, only saving them to the Store when Save is called.
type Stash struct {
	chunks chunks.Store
	ids    idpool.Pool
	local  map[uint64]*chunks.Chunk
}

// Get returns a chunk either from the local stash, or from the
// Store (for Private keys).
//
// For Private keys, modifying the returned chunk *will* cause the
// locally stored data to change. This is the intended usage of a
// stash.
func (s *Stash) Get(key cas.Key, typ string, level uint8) (*chunks.Chunk, error) {
	priv, ok := key.Private()
	if ok {
		chunk, ok := s.local[priv]
		if !ok {
			return nil, cas.NotFound{
				Type:  typ,
				Level: level,
				Key:   key,
			}
		}
		return chunk, nil
	}

	chunk, err := s.chunks.Get(key, typ, level)
	return chunk, err
}

func (s *Stash) drop(key cas.Key) {
	priv, ok := key.Private()
	if !ok {
		panic(fmt.Sprintf("Cannot drop non-private key: %s", key))
	}
	s.ids.Put(priv)
	delete(s.local, priv)
}

// Clone is like Get but clones the chunk if it's not already private.
// Chunks that are already private are returned as-is.
//
// A cloned chunk will have a buffer of size bytes. This is intended
// to use for re-inflating zero-trimmed chunks.
//
// Modifying the returned chunk *will* cause the locally stored data
// to change. This is the intended usage of a stash.
func (s *Stash) Clone(key cas.Key, typ string, level uint8, size uint32) (cas.Key, *chunks.Chunk, error) {
	priv, ok := key.Private()
	if ok {
		chunk, ok := s.local[priv]
		if !ok {
			return key, nil, cas.NotFound{
				Type:  typ,
				Level: level,
				Key:   key,
			}
		}
		return key, chunk, nil
	}

	chunk, err := s.chunks.Get(key, typ, level)
	if err != nil {
		return key, nil, err
	}

	// clone the byte slice
	tmp := make([]byte, size)
	copy(tmp, chunk.Buf)
	chunk.Buf = tmp

	priv = s.ids.Get()
	privkey := cas.NewKeyPrivateNum(priv)
	s.local[priv] = chunk
	return privkey, chunk, nil
}

// Save the local Chunk to the Store.
//
// On success, the old key becomes invalid.
func (s *Stash) Save(key cas.Key) (cas.Key, error) {
	priv, ok := key.Private()
	if !ok {
		return key, nil
	}

	chunk, ok := s.local[priv]
	if !ok {
		return key, cas.NotFound{
			Key: key,
		}
	}

	newkey, err := s.chunks.Add(chunk)
	if err != nil {
		return key, err
	}
	s.drop(key)
	return newkey, nil
}
<commit_msg>cas/chunks/stash: drop can work on priv number; avoids a potential panic<commit_after>package stash

import (
	"bazil.org/bazil/cas"
	"bazil.org/bazil/cas/chunks"
	"bazil.org/bazil/idpool"
)

// New creates a new Stash.
func New(bs chunks.Store) *Stash {
	s := &Stash{
		chunks: bs,
		local:  make(map[uint64]*chunks.Chunk),
	}
	return s
}

// Stash is a proxy for a chunks.Store, but it keeps Private Keys
// local, only saving them to the Store when Save is called.
type Stash struct {
	chunks chunks.Store
	ids    idpool.Pool
	local  map[uint64]*chunks.Chunk
}

// Get returns a chunk either from the local stash, or from the
// Store (for Private keys).
//
// For Private keys, modifying the returned chunk *will* cause the
// locally stored data to change. This is the intended usage of a
// stash.
func (s *Stash) Get(key cas.Key, typ string, level uint8) (*chunks.Chunk, error) {
	priv, ok := key.Private()
	if ok {
		chunk, ok := s.local[priv]
		if !ok {
			return nil, cas.NotFound{
				Type:  typ,
				Level: level,
				Key:   key,
			}
		}
		return chunk, nil
	}

	chunk, err := s.chunks.Get(key, typ, level)
	return chunk, err
}

func (s *Stash) drop(priv uint64) {
	s.ids.Put(priv)
	delete(s.local, priv)
}

// Clone is like Get but clones the chunk if it's not already private.
// Chunks that are already private are returned as-is.
//
// A cloned chunk will have a buffer of size bytes. This is intended
// to use for re-inflating zero-trimmed chunks.
//
// Modifying the returned chunk *will* cause the locally stored data
// to change. This is the intended usage of a stash.
func (s *Stash) Clone(key cas.Key, typ string, level uint8, size uint32) (cas.Key, *chunks.Chunk, error) {
	priv, ok := key.Private()
	if ok {
		chunk, ok := s.local[priv]
		if !ok {
			return key, nil, cas.NotFound{
				Type:  typ,
				Level: level,
				Key:   key,
			}
		}
		return key, chunk, nil
	}

	chunk, err := s.chunks.Get(key, typ, level)
	if err != nil {
		return key, nil, err
	}

	// clone the byte slice
	tmp := make([]byte, size)
	copy(tmp, chunk.Buf)
	chunk.Buf = tmp

	priv = s.ids.Get()
	privkey := cas.NewKeyPrivateNum(priv)
	s.local[priv] = chunk
	return privkey, chunk, nil
}

// Save the local Chunk to the Store.
//
// On success, the old key becomes invalid.
func (s *Stash) Save(key cas.Key) (cas.Key, error) {
	priv, ok := key.Private()
	if !ok {
		return key, nil
	}

	chunk, ok := s.local[priv]
	if !ok {
		return key, cas.NotFound{
			Key: key,
		}
	}

	newkey, err := s.chunks.Add(chunk)
	if err != nil {
		return key, err
	}
	s.drop(priv)
	return newkey, nil
}
<|endoftext|>


<commit_before>package api

import (
	"github.com/ethereum/go-ethereum/eth"
	"github.com/ethereum/go-ethereum/rpc/codec"
	"github.com/ethereum/go-ethereum/rpc/shared"
	"github.com/ethereum/go-ethereum/xeth"
)

const (
	NetApiVersion = "1.0"
)

var (
	// mapping between methods and handlers
	netMapping = map[string]nethandler{
		"net_version":   (*netApi).Version,
		"net_peerCount": (*netApi).PeerCount,
		"net_listening": (*netApi).IsListening,
		"net_peers":     (*netApi).Peers,
	}
)

// net callback handler
type nethandler func(*netApi, *shared.Request) (interface{}, error)

// net api provider
type netApi struct {
	xeth     *xeth.XEth
	ethereum *eth.Ethereum
	methods  map[string]nethandler
	codec    codec.ApiCoder
}

// create a new net api instance
func NewNetApi(xeth *xeth.XEth, eth *eth.Ethereum, coder codec.Codec) *netApi {
	return &netApi{
		xeth:     xeth,
		ethereum: eth,
		methods:  netMapping,
		codec:    coder.New(nil),
	}
}

// collection with supported methods
func (self *netApi) Methods() []string {
	methods := make([]string, len(self.methods))
	i := 0
	for k := range self.methods {
		methods[i] = k
		i++
	}
	return methods
}

// Execute given request
func (self *netApi) Execute(req *shared.Request) (interface{}, error) {
	if callback, ok := self.methods[req.Method]; ok {
		return callback(self, req)
	}

	return nil, shared.NewNotImplementedError(req.Method)
}

func (self *netApi) Name() string {
	return NetApiName
}

func (self *netApi) ApiVersion() string {
	return NetApiVersion
}

// Network version
func (self *netApi) Version(req *shared.Request) (interface{}, error) {
	return self.xeth.NetworkVersion(), nil
}

// Number of connected peers
func (self *netApi) PeerCount(req *shared.Request) (interface{}, error) {
	return self.xeth.PeerCount(), nil
}

func (self *netApi) IsListening(req *shared.Request) (interface{}, error) {
	return self.xeth.IsListening(), nil
}

func (self *netApi) Peers(req *shared.Request) (interface{}, error) {
	return self.ethereum.PeersInfo(), nil
}
<commit_msg>fixed rpc test failure in net_peerCount<commit_after>package api

import (
	"github.com/ethereum/go-ethereum/eth"
	"github.com/ethereum/go-ethereum/rpc/codec"
	"github.com/ethereum/go-ethereum/rpc/shared"
	"github.com/ethereum/go-ethereum/xeth"
)

const (
	NetApiVersion = "1.0"
)

var (
	// mapping between methods and handlers
	netMapping = map[string]nethandler{
		"net_version":   (*netApi).Version,
		"net_peerCount": (*netApi).PeerCount,
		"net_listening": (*netApi).IsListening,
		"net_peers":     (*netApi).Peers,
	}
)

// net callback handler
type nethandler func(*netApi, *shared.Request) (interface{}, error)

// net api provider
type netApi struct {
	xeth     *xeth.XEth
	ethereum *eth.Ethereum
	methods  map[string]nethandler
	codec    codec.ApiCoder
}

// create a new net api instance
func NewNetApi(xeth *xeth.XEth, eth *eth.Ethereum, coder codec.Codec) *netApi {
	return &netApi{
		xeth:     xeth,
		ethereum: eth,
		methods:  netMapping,
		codec:    coder.New(nil),
	}
}

// collection with supported methods
func (self *netApi) Methods() []string {
	methods := make([]string, len(self.methods))
	i := 0
	for k := range self.methods {
		methods[i] = k
		i++
	}
	return methods
}

// Execute given request
func (self *netApi) Execute(req *shared.Request) (interface{}, error) {
	if callback, ok := self.methods[req.Method]; ok {
		return callback(self, req)
	}

	return nil, shared.NewNotImplementedError(req.Method)
}

func (self *netApi) Name() string {
	return NetApiName
}

func (self *netApi) ApiVersion() string {
	return NetApiVersion
}

// Network version
func (self *netApi) Version(req *shared.Request) (interface{}, error) {
	return self.xeth.NetworkVersion(), nil
}

// Number of connected peers
func (self *netApi) PeerCount(req *shared.Request) (interface{}, error) {
	return newHexNum(self.xeth.PeerCount()), nil
}

func (self *netApi) IsListening(req *shared.Request) (interface{}, error) {
	return self.xeth.IsListening(), nil
}

func (self *netApi) Peers(req *shared.Request) (interface{}, error) {
	return self.ethereum.PeersInfo(), nil
}
<|endoftext|>


<commit_before>package system

import (
	"fmt"
	"strconv"
	"strings"
	"testing"
	"time"

	"github.com/quintilesims/layer0/common/config"
	"github.com/quintilesims/layer0/common/testutils"
	"github.com/quintilesims/layer0/tests/clients"
	"github.com/quintilesims/tftest"
)

type StressTestCase struct {
	NumDeploys        int
	NumDeployFamilies int
	NumEnvironments   int
	NumLoadBalancers  int
	NumServices       int
	NumTasks          int
}

func runTest(b *testing.B, c StressTestCase) {
	if c.NumTasks > 0 || c.NumServices > 0 {
		if c.NumEnvironments <= 0 {
			b.Fatalf("Cannot have Tasks and/or Services without Environments.")
		}

		if c.NumDeploys <= 0 {
			b.Fatalf("Cannot have Tasks and/or Services without Deploys.")
		}
	}

	vars := map[string]string{
		"endpoint":            config.APIEndpoint(),
		"token":               config.AuthToken(),
		"num_deploys":         strconv.Itoa(c.NumDeploys),
		"num_deploy_families": strconv.Itoa(c.NumDeployFamilies),
		"num_environments":    strconv.Itoa(c.NumEnvironments),
		"num_load_balancers":  strconv.Itoa(c.NumLoadBalancers),
		"num_services":        strconv.Itoa(c.NumServices),
	}

	terraform := tftest.NewTestContext(
		b,
		tftest.Dir("module"),
		tftest.Vars(vars),
		tftest.DryRun(*dry),
		tftest.Log(b),
	)

	layer0 := clients.NewLayer0TestClient(b, vars["endpoint"], vars["token"])

	terraform.Apply()
	defer terraform.Destroy()

	methodsToBenchmark := map[string]func(){}

	if c.NumDeploys > 0 {
		deployIDs := strings.Split(terraform.Output("deploy_ids"), ",\n")
		methodsToBenchmark["GetDeploy"] = func() { layer0.GetDeploy(deployIDs[0]) }
		methodsToBenchmark["ListDeploys"] = func() { layer0.ListDeploys() }
	}

	if c.NumEnvironments > 0 {
		environmentIDs := strings.Split(terraform.Output("environment_ids"), ",\n")
		methodsToBenchmark["GetEnvironment"] = func() { layer0.GetEnvironment(environmentIDs[0]) }
		methodsToBenchmark["ListEnvironments"] = func() { layer0.ListEnvironments() }
	}

	if c.NumLoadBalancers > 0 {
		loadBalancerIDs := strings.Split(terraform.Output("load_balancer_ids"), ",\n")
		methodsToBenchmark["GetLoadBalancer"] = func() { layer0.GetLoadBalancer(loadBalancerIDs[0]) }
		methodsToBenchmark["ListLoadBalancers"] = func() { layer0.ListLoadBalancers() }
	}

	if c.NumServices > 0 {
		serviceIDs := strings.Split(terraform.Output("service_ids"), ",\n")
		methodsToBenchmark["GetService"] = func() { layer0.GetService(serviceIDs[0]) }
		methodsToBenchmark["ListServices"] = func() { layer0.ListServices() }
	}

	if c.NumTasks > 0 {
		methodsToBenchmark["ListTasks"] = func() { layer0.ListTasks() }

		deployIDs := strings.Split(terraform.Output("deploy_ids"), ",\n")
		environmentIDs := strings.Split(terraform.Output("environment_ids"), ",\n")

		tasksCreated := 0
		for copies := c.NumTasks / 2; tasksCreated < c.NumTasks; copies = copies / 2 {
			taskName := fmt.Sprintf("Task%v", copies)
			go func() {
				log.Debugf("Creating task %v", taskName)
				layer0.CreateTask(taskName, environmentIDs[0], deployIDs[0], copies, nil)
			}()

			tasksCreated += copies
			if copies <= 1 {
				copies++
			}
		}

		testutils.WaitFor(b, time.Second*30, time.Minute*10, func() bool {
			log.Debug("Waiting for all tasks to run")

			var numTasks int
			for _, taskSummary := range layer0.ListTasks() {
				if taskSummary.EnvironmentID == environmentIDs[0] {
					numTasks++
				}
			}

			log.Debugf("%d/%d tasks have run", numTasks, c.NumTasks)
			return numTasks >= c.NumTasks
		})
	}

	benchmark(b, methodsToBenchmark)
}

func benchmark(b *testing.B, methods map[string]func()) {
	for name, fn := range methods {
		b.Run(name, func(b *testing.B) {
			for n := 0; n < b.N; n++ {
				fn()
			}
		})
	}
}

func Benchmark5Services(b *testing.B) {
	runTest(b, StressTestCase{
		NumDeploys:      1,
		NumEnvironments: 2,
		NumServices:     5,
	})
}

func Benchmark25Environments(b *testing.B) {
	runTest(b, StressTestCase{
		NumEnvironments: 25,
	})
}

func Benchmark10Environments10Deploys(b *testing.B) {
	runTest(b, StressTestCase{
		NumDeploys:      10,
		NumEnvironments: 10,
	})
}

func Benchmark20Environments20Deploys(b *testing.B) {
	runTest(b, StressTestCase{
		NumDeploys:      20,
		NumEnvironments: 20,
	})
}

func Benchmark5Environments100Deploys(b *testing.B) {
	runTest(b, StressTestCase{
		NumDeploys:      100,
		NumEnvironments: 5,
	})
}

func Benchmark10Environments10Deploys10Services(b *testing.B) {
	runTest(b, StressTestCase{
		NumDeploys:      10,
		NumEnvironments: 10,
		NumServices:     10,
	})
}

func Benchmark5Environments5Deploys50Services(b *testing.B) {
	runTest(b, StressTestCase{
		NumDeploys:      5,
		NumEnvironments: 5,
		NumServices:     50,
	})
}

func Benchmark15Environments15Deploys15Services15LoadBalancers(b *testing.B) {
	runTest(b, StressTestCase{
		NumDeploys:       15,
		NumEnvironments:  15,
		NumLoadBalancers: 15,
		NumServices:      15,
	})
}

func Benchmark25Environments25Deploys25Services25LoadBalancers(b *testing.B) {
	runTest(b, StressTestCase{
		NumDeploys:       25,
		NumEnvironments:  25,
		NumLoadBalancers: 25,
		NumServices:      25,
	})
}

func Benchmark100Tasks(b *testing.B) {
	runTest(b, StressTestCase{
		NumDeploys:      1,
		NumEnvironments: 1,
		NumTasks:        100,
	})
}
<commit_msg>Don't wait for tasks to finish<commit_after>package system

import (
	"fmt"
	"strconv"
	"strings"
	"testing"

	"github.com/quintilesims/layer0/common/config"
	"github.com/quintilesims/layer0/tests/clients"
	"github.com/quintilesims/tftest"
)

type StressTestCase struct {
	NumDeploys        int
	NumDeployFamilies int
	NumEnvironments   int
	NumLoadBalancers  int
	NumServices       int
	NumTasks          int
}

func runTest(b *testing.B, c StressTestCase) {
	if c.NumTasks > 0 || c.NumServices > 0 {
		if c.NumEnvironments <= 0 {
			b.Fatalf("Cannot have Tasks and/or Services without Environments.")
		}

		if c.NumDeploys <= 0 {
			b.Fatalf("Cannot have Tasks and/or Services without Deploys.")
		}
	}

	vars := map[string]string{
		"endpoint":            config.APIEndpoint(),
		"token":               config.AuthToken(),
		"num_deploys":         strconv.Itoa(c.NumDeploys),
		"num_deploy_families": strconv.Itoa(c.NumDeployFamilies),
		"num_environments":    strconv.Itoa(c.NumEnvironments),
		"num_load_balancers":  strconv.Itoa(c.NumLoadBalancers),
		"num_services":        strconv.Itoa(c.NumServices),
	}

	terraform := tftest.NewTestContext(
		b,
		tftest.Dir("module"),
		tftest.Vars(vars),
		tftest.DryRun(*dry),
		tftest.Log(b),
	)

	layer0 := clients.NewLayer0TestClient(b, vars["endpoint"], vars["token"])

	terraform.Apply()
	defer terraform.Destroy()

	methodsToBenchmark := map[string]func(){}

	if c.NumDeploys > 0 {
		deployIDs := strings.Split(terraform.Output("deploy_ids"), ",\n")
		methodsToBenchmark["GetDeploy"] = func() { layer0.GetDeploy(deployIDs[0]) }
		methodsToBenchmark["ListDeploys"] = func() { layer0.ListDeploys() }
	}

	if c.NumEnvironments > 0 {
		environmentIDs := strings.Split(terraform.Output("environment_ids"), ",\n")
		methodsToBenchmark["GetEnvironment"] = func() { layer0.GetEnvironment(environmentIDs[0]) }
		methodsToBenchmark["ListEnvironments"] = func() { layer0.ListEnvironments() }
	}

	if c.NumLoadBalancers > 0 {
		loadBalancerIDs := strings.Split(terraform.Output("load_balancer_ids"), ",\n")
		methodsToBenchmark["GetLoadBalancer"] = func() { layer0.GetLoadBalancer(loadBalancerIDs[0]) }
		methodsToBenchmark["ListLoadBalancers"] = func() { layer0.ListLoadBalancers() }
	}

	if c.NumServices > 0 {
		serviceIDs := strings.Split(terraform.Output("service_ids"), ",\n")
		methodsToBenchmark["GetService"] = func() { layer0.GetService(serviceIDs[0]) }
		methodsToBenchmark["ListServices"] = func() { layer0.ListServices() }
	}

	if c.NumTasks > 0 {
		methodsToBenchmark["ListTasks"] = func() { layer0.ListTasks() }

		deployIDs := strings.Split(terraform.Output("deploy_ids"), ",\n")
		environmentIDs := strings.Split(terraform.Output("environment_ids"), ",\n")

		tasksCreated := 0
		for copies := c.NumTasks / 2; tasksCreated < c.NumTasks; copies = copies / 2 {
			taskName := fmt.Sprintf("Task%v", copies)
			go func() {
				log.Debugf("Creating task %v", taskName)
				layer0.CreateTask(taskName, environmentIDs[0], deployIDs[0], copies, nil)
			}()

			tasksCreated += copies
			if copies <= 1 {
				copies++
			}
		}
	}

	benchmark(b, methodsToBenchmark)
}

func benchmark(b *testing.B, methods map[string]func()) {
	for name, fn := range methods {
		b.Run(name, func(b *testing.B) {
			for n := 0; n < b.N; n++ {
				fn()
			}
		})
	}
}

func Benchmark5Services(b *testing.B) {
	runTest(b, StressTestCase{
		NumDeploys:      1,
		NumEnvironments: 2,
		NumServices:     5,
	})
}

func Benchmark25Environments(b *testing.B) {
	runTest(b, StressTestCase{
		NumEnvironments: 25,
	})
}

func Benchmark10Environments10Deploys(b *testing.B) {
	runTest(b, StressTestCase{
		NumDeploys:      10,
		NumEnvironments: 10,
	})
}

func Benchmark20Environments20Deploys(b *testing.B) {
	runTest(b, StressTestCase{
		NumDeploys:      20,
		NumEnvironments: 20,
	})
}

func Benchmark5Environments100Deploys(b *testing.B) {
	runTest(b, StressTestCase{
		NumDeploys:      100,
		NumEnvironments: 5,
	})
}

func Benchmark10Environments10Deploys10Services(b *testing.B) {
	runTest(b, StressTestCase{
		NumDeploys:      10,
		NumEnvironments: 10,
		NumServices:     10,
	})
}

func Benchmark5Environments5Deploys50Services(b *testing.B) {
	runTest(b, StressTestCase{
		NumDeploys:      5,
		NumEnvironments: 5,
		NumServices:     50,
	})
}

func Benchmark15Environments15Deploys15Services15LoadBalancers(b *testing.B) {
	runTest(b, StressTestCase{
		NumDeploys:       15,
		NumEnvironments:  15,
		NumLoadBalancers: 15,
		NumServices:      15,
	})
}

func Benchmark25Environments25Deploys25Services25LoadBalancers(b *testing.B) {
	runTest(b, StressTestCase{
		NumDeploys:       25,
		NumEnvironments:  25,
		NumLoadBalancers: 25,
		NumServices:      25,
	})
}

func Benchmark100Tasks(b *testing.B) {
	runTest(b, StressTestCase{
		NumDeploys:      1,
		NumEnvironments: 1,
		NumTasks:        100,
	})
}
<|endoftext|>


<commit_before>// Copyright 2019 Istio Authors
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in conformance 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 conformance

import (
	"context"
	"fmt"
	"net/url"
	"regexp"
	"strconv"
	"strings"
	"testing"
	"time"

	envoyAdmin "github.com/envoyproxy/go-control-plane/envoy/admin/v2alpha"
	"gopkg.in/yaml.v2"

	"istio.io/istio/pkg/config/protocol"
	"istio.io/istio/pkg/test"
	"istio.io/istio/pkg/test/conformance"
	"istio.io/istio/pkg/test/conformance/constraint"
	epb "istio.io/istio/pkg/test/echo/proto"
	"istio.io/istio/pkg/test/framework"
	"istio.io/istio/pkg/test/framework/components/echo"
	"istio.io/istio/pkg/test/framework/components/echo/echoboot"
	"istio.io/istio/pkg/test/framework/components/environment"
	"istio.io/istio/pkg/test/framework/components/galley"
	"istio.io/istio/pkg/test/framework/components/namespace"
	"istio.io/istio/pkg/test/framework/components/pilot"
	"istio.io/istio/pkg/test/framework/label"
	"istio.io/istio/pkg/test/util/structpath"
	"istio.io/istio/pkg/test/util/tmpl"
)

func TestConformance(t *testing.T) {
	framework.Run(t, func(ctx framework.TestContext) {
		cases, err := loadCases()
		if err != nil {
			ctx.Fatalf("error loading test cases: %v", err)
		}

		gal := galley.NewOrFail(ctx, ctx, galley.Config{})
		p := pilot.NewOrFail(ctx, ctx, pilot.Config{Galley: gal})

		for _, ca := range cases {
			tst := ctx.NewSubTest(ca.Metadata.Name)

			for _, lname := range ca.Metadata.Labels {
				l, ok := label.Find(lname)
				if !ok {
					ctx.Fatalf("label not found: %v", lname)
				}
				tst = tst.Label(l)
			}

			if ca.Metadata.Exclusive {
				tst.Run(runCaseFn(p, gal, ca))
			} else {
				tst.RunParallel(runCaseFn(p, gal, ca))
			}
		}
	})
}

func runCaseFn(p pilot.Instance, gal galley.Instance, ca *conformance.Test) func(framework.TestContext) {
	return func(ctx framework.TestContext) {
		match := true
	mainloop:
		for _, ename := range ca.Metadata.Environments {
			match = false
			for _, n := range environment.Names() {
				if n.String() == ename && n == ctx.Environment().EnvironmentName() {
					match = true
					break mainloop
				}
			}
		}

		if !match {
			ctx.Skipf("None of the expected environment(s) not found: %v", ca.Metadata.Environments)
		}

		if ca.Metadata.Skip {
			ctx.Skipf("Test is marked as skip")
		}

		// If there are any changes to the mesh config, then capture the original and restore.
		for _, s := range ca.Stages {
			if s.MeshConfig != nil {
				// TODO: Add support to get/set old meshconfig to avoid cross-test interference.
				// originalMeshCfg := gal.GetMeshConfigOrFail(ctx)
				// defer gal.SetMeshConfigOrFail(ctx, originalMeshCfg)
				break
			}
		}

		ns := namespace.NewOrFail(ctx, ctx, namespace.Config{
			Prefix: "conf",
			Inject: true,
		})

		if len(ca.Stages) == 1 {
			runStage(ctx, p, gal, ns, ca.Stages[0])
		} else {
			for i, s := range ca.Stages {
				ctx.NewSubTest(fmt.Sprintf("%d", i)).Run(func(ctx framework.TestContext) {
					runStage(ctx, p, gal, ns, s)
				})
			}
		}
	}
}

func runStage(ctx framework.TestContext, pil pilot.Instance, gal galley.Instance, ns namespace.Instance, s *conformance.Stage) {
	if s.MeshConfig != nil {
		gal.SetMeshConfigOrFail(ctx, *s.MeshConfig)
	}

	i := s.Input
	gal.ApplyConfigOrFail(ctx, ns, i)
	defer func() {
		gal.DeleteConfigOrFail(ctx, ns, i)
	}()

	if s.MCP != nil {
		validateMCPState(ctx, gal, ns, s)
	}
	if s.Traffic != nil {
		validateTraffic(ctx, pil, gal, ns, s)
	}

	// More and different types of validations can go here
}

func validateMCPState(ctx test.Failer, gal galley.Instance, ns namespace.Instance, s *conformance.Stage) {
	p := constraint.Params{
		Namespace: ns.Name(),
	}
	for _, coll := range s.MCP.Constraints {
		gal.WaitForSnapshotOrFail(ctx, coll.Name, func(actuals []*galley.SnapshotObject) error {
			for _, rangeCheck := range coll.Check {
				a := make([]interface{}, len(actuals))
				for i, item := range actuals {
					// Deep copy the item so we can modify the metadata safely
					itemCopy := &galley.SnapshotObject{}
					metadata := *item.Metadata
					itemCopy.Metadata = &metadata
					itemCopy.Body = item.Body
					itemCopy.TypeURL = item.TypeURL

					a[i] = itemCopy

					// Clear out for stable comparison.
					itemCopy.Metadata.CreateTime = nil
					itemCopy.Metadata.Version = ""
					if itemCopy.Metadata.Annotations != nil {
						delete(itemCopy.Metadata.Annotations, "kubectl.kubernetes.io/last-applied-configuration")
						if len(itemCopy.Metadata.Annotations) == 0 {
							itemCopy.Metadata.Annotations = nil
						}
					}
				}

				if err := rangeCheck.ValidateItems(a, p); err != nil {
					return err
				}
			}
			return nil
		})
	}

}

// Returns success for Envoy configs containing routes for all of specified domains.
func domainAcceptFunc(domains []string) func(*envoyAdmin.ConfigDump) (bool, error) {
	return func(cfg *envoyAdmin.ConfigDump) (bool, error) {
		validator := structpath.ForProto(cfg)
		const q = "{.configs[*].dynamicRouteConfigs[*].routeConfig.virtualHosts[*].domains[?(@ == %q)]}"
		for _, domain := range domains {
			// TODO(qfel): Figure out how to get rid of the loop.
			if err := validator.Exists(q, domain).Check(); err != nil {
				return false, err
			}
		}
		return true, nil
	}
}

func virtualServiceHosts(ctx test.Failer, resources string) []string {
	// TODO(qfel): We should parse input early and work on structured data here. Or at least if some
	// other code path needs it.
	var inputs []map[interface{}]interface{}
	for _, inputYAML := range strings.Split(resources, "\n---\n") {
		var input map[interface{}]interface{}
		if err := yaml.Unmarshal([]byte(inputYAML), &input); err != nil {
			ctx.Fatalf("yaml.Unmarshal: %v", err)
		}
		inputs = append(inputs, input)
	}

	var vHosts []string
	for _, res := range inputs {
		if res["apiVersion"] != "networking.istio.io/v1alpha3" || res["kind"] != "VirtualService" {
			continue
		}
		spec := res["spec"].(map[interface{}]interface{})
		hosts := spec["hosts"].([]interface{})
		for _, h := range hosts {
			vHosts = append(vHosts, h.(string))
		}
	}
	return vHosts
}

func validateTraffic(t framework.TestContext, pil pilot.Instance, gal galley.Instance, ns namespace.Instance, stage *conformance.Stage) {
	echos := make([]echo.Instance, len(stage.Traffic.Services))
	b := echoboot.NewBuilderOrFail(t, t)
	for i, svc := range stage.Traffic.Services {
		var ports []echo.Port
		for _, p := range svc.Ports {
			ports = append(ports, echo.Port{
				Name:        p.Name,
				Protocol:    protocol.Instance(p.Protocol),
				ServicePort: int(p.ServicePort),
			})
		}
		b = b.With(&echos[i], echo.Config{
			Galley:    gal,
			Pilot:     pil,
			Service:   svc.Name,
			Namespace: ns,
			Ports:     ports,
		})
	}
	if err := b.Build(); err != nil {
		t.Fatal(err)
	}

	services := make(map[string]echo.Instance)
	for i, svc := range echos {
		services[stage.Traffic.Services[i].Name] = svc
		svc.WaitUntilCallableOrFail(t, echos...)
	}

	ready := make(map[string]bool)
	vHosts := virtualServiceHosts(t, stage.Input)
	for _, call := range stage.Traffic.Calls {
		call.URL = tmpl.EvaluateOrFail(t, call.URL, constraint.Params{
			Namespace: ns.Name(),
		})
		t.NewSubTest(call.URL).Run(func(t framework.TestContext) {
			hostname, err := wildcardToRegexp(call.Response.Callee)
			if err != nil {
				t.Fatalf("Internal error: regexp.Compile: %v", err)
			}
			caller := services[call.Caller]
			if !ready[call.Caller] && false {
				t.Logf("Waiting for sidecar(s) for %s to contain domains: %s", call.Caller, strings.Join(vHosts, ", "))
				for _, w := range caller.WorkloadsOrFail(t) {
					w.Sidecar().WaitForConfigOrFail(t, domainAcceptFunc(vHosts))
				}
				ready[call.Caller] = true
			}
			u, err := url.Parse(call.URL)
			if err != nil {
				t.Fatalf("Parse call URL: %v", err)
			}
			workload := caller.WorkloadsOrFail(t)[0]

			validateWithRedo(t, func(ctx context.Context) bool {
				responses, err := workload.ForwardEcho(ctx, &epb.ForwardEchoRequest{
					Url:   call.URL,
					Count: 1,
					Headers: []*epb.Header{{
						Key:   "Host",
						Value: u.Host, // TODO(qfel): Why cannot echo infer it from URL?
					}},
				})
				if err != nil {
					t.Logf("Initiating test call: %v", err)
					return false
				}

				if len(responses) != 1 {
					t.Logf("Received %d responses, want 1", len(responses))
					return false
				}
				resp := responses[0]
				var fail bool
				if resp.Code != strconv.Itoa(call.Response.Code) {
					t.Logf("Responded with %s, want %d", resp.Code, call.Response.Code)
					fail = true
				}
				if !hostname.MatchString(resp.Hostname) {
					t.Logf("Callee %q doesn't match %q", resp.Hostname, call.Response.Callee)
					fail = true
				}
				return !fail
			})
		})
	}
}

func validateWithRedo(t framework.TestContext, f func(context.Context) bool) {
	const (
		pollDelay           = time.Second // How much to wait after a failed attempt.
		firstSuccessTimeout = time.Minute
		// After the traffic flows successfully for the first time, repeat according to the parameters
		// below.
		redoAttempts = 10
		redoTimeout  = 20 * time.Second
		redoFraction = 0.9 // Fail if not enough attempts succeed.
	)

	ctx, cfn := context.WithTimeout(context.Background(), firstSuccessTimeout)
	defer cfn()
	for {
		if f(ctx) {
			break
		}
		if sleepCtx(ctx, pollDelay) != nil {
			t.Fatal("No successful attempt")
		}
	}

	var passed int
	ctx, cfn = context.WithTimeout(context.Background(), redoTimeout)
	defer cfn()
	for i := 0; i < redoAttempts; i++ {
		start := time.Now()
		res := f(ctx)
		t.Logf("Traffic attempt #%d: %v in %s", i, res, time.Since(start))
		if res {
			passed++
		} else if i+1 < redoAttempts {
			// No need to check errors. If ctx times out, further calls to f will time out.
			_ = sleepCtx(ctx, pollDelay)
		}
	}
	fr := float64(passed) / float64(redoAttempts)
	t.Logf("%d/%d (%f) attempts succeeded", passed, redoAttempts, fr)
	if fr < redoFraction {
		t.Errorf("Success rate is %f, need at least %f", fr, redoFraction)
	}
}

func sleepCtx(ctx context.Context, duration time.Duration) error {
	t := time.NewTimer(duration)
	select {
	case <-t.C:
		return nil
	case <-ctx.Done():
		t.Stop()
		return ctx.Err()
	}
}

func wildcardToRegexp(s string) (*regexp.Regexp, error) {
	s = fmt.Sprintf("^%s$", regexp.QuoteMeta(s))
	return regexp.Compile(strings.Replace(s, `\*`, ".*", -1))
}
<commit_msg>Use index instead of URL to name steps of traffic test (#16424)<commit_after>// Copyright 2019 Istio Authors
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in conformance 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 conformance

import (
	"context"
	"fmt"
	"net/url"
	"regexp"
	"strconv"
	"strings"
	"testing"
	"time"

	envoyAdmin "github.com/envoyproxy/go-control-plane/envoy/admin/v2alpha"
	"gopkg.in/yaml.v2"

	"istio.io/istio/pkg/config/protocol"
	"istio.io/istio/pkg/test"
	"istio.io/istio/pkg/test/conformance"
	"istio.io/istio/pkg/test/conformance/constraint"
	epb "istio.io/istio/pkg/test/echo/proto"
	"istio.io/istio/pkg/test/framework"
	"istio.io/istio/pkg/test/framework/components/echo"
	"istio.io/istio/pkg/test/framework/components/echo/echoboot"
	"istio.io/istio/pkg/test/framework/components/environment"
	"istio.io/istio/pkg/test/framework/components/galley"
	"istio.io/istio/pkg/test/framework/components/namespace"
	"istio.io/istio/pkg/test/framework/components/pilot"
	"istio.io/istio/pkg/test/framework/label"
	"istio.io/istio/pkg/test/util/structpath"
	"istio.io/istio/pkg/test/util/tmpl"
)

func TestConformance(t *testing.T) {
	framework.Run(t, func(ctx framework.TestContext) {
		cases, err := loadCases()
		if err != nil {
			ctx.Fatalf("error loading test cases: %v", err)
		}

		gal := galley.NewOrFail(ctx, ctx, galley.Config{})
		p := pilot.NewOrFail(ctx, ctx, pilot.Config{Galley: gal})

		for _, ca := range cases {
			tst := ctx.NewSubTest(ca.Metadata.Name)

			for _, lname := range ca.Metadata.Labels {
				l, ok := label.Find(lname)
				if !ok {
					ctx.Fatalf("label not found: %v", lname)
				}
				tst = tst.Label(l)
			}

			if ca.Metadata.Exclusive {
				tst.Run(runCaseFn(p, gal, ca))
			} else {
				tst.RunParallel(runCaseFn(p, gal, ca))
			}
		}
	})
}

func runCaseFn(p pilot.Instance, gal galley.Instance, ca *conformance.Test) func(framework.TestContext) {
	return func(ctx framework.TestContext) {
		match := true
	mainloop:
		for _, ename := range ca.Metadata.Environments {
			match = false
			for _, n := range environment.Names() {
				if n.String() == ename && n == ctx.Environment().EnvironmentName() {
					match = true
					break mainloop
				}
			}
		}

		if !match {
			ctx.Skipf("None of the expected environment(s) not found: %v", ca.Metadata.Environments)
		}

		if ca.Metadata.Skip {
			ctx.Skipf("Test is marked as skip")
		}

		// If there are any changes to the mesh config, then capture the original and restore.
		for _, s := range ca.Stages {
			if s.MeshConfig != nil {
				// TODO: Add support to get/set old meshconfig to avoid cross-test interference.
				// originalMeshCfg := gal.GetMeshConfigOrFail(ctx)
				// defer gal.SetMeshConfigOrFail(ctx, originalMeshCfg)
				break
			}
		}

		ns := namespace.NewOrFail(ctx, ctx, namespace.Config{
			Prefix: "conf",
			Inject: true,
		})

		if len(ca.Stages) == 1 {
			runStage(ctx, p, gal, ns, ca.Stages[0])
		} else {
			for i, s := range ca.Stages {
				ctx.NewSubTest(fmt.Sprintf("%d", i)).Run(func(ctx framework.TestContext) {
					runStage(ctx, p, gal, ns, s)
				})
			}
		}
	}
}

func runStage(ctx framework.TestContext, pil pilot.Instance, gal galley.Instance, ns namespace.Instance, s *conformance.Stage) {
	if s.MeshConfig != nil {
		gal.SetMeshConfigOrFail(ctx, *s.MeshConfig)
	}

	i := s.Input
	gal.ApplyConfigOrFail(ctx, ns, i)
	defer func() {
		gal.DeleteConfigOrFail(ctx, ns, i)
	}()

	if s.MCP != nil {
		validateMCPState(ctx, gal, ns, s)
	}
	if s.Traffic != nil {
		validateTraffic(ctx, pil, gal, ns, s)
	}

	// More and different types of validations can go here
}

func validateMCPState(ctx test.Failer, gal galley.Instance, ns namespace.Instance, s *conformance.Stage) {
	p := constraint.Params{
		Namespace: ns.Name(),
	}
	for _, coll := range s.MCP.Constraints {
		gal.WaitForSnapshotOrFail(ctx, coll.Name, func(actuals []*galley.SnapshotObject) error {
			for _, rangeCheck := range coll.Check {
				a := make([]interface{}, len(actuals))
				for i, item := range actuals {
					// Deep copy the item so we can modify the metadata safely
					itemCopy := &galley.SnapshotObject{}
					metadata := *item.Metadata
					itemCopy.Metadata = &metadata
					itemCopy.Body = item.Body
					itemCopy.TypeURL = item.TypeURL

					a[i] = itemCopy

					// Clear out for stable comparison.
					itemCopy.Metadata.CreateTime = nil
					itemCopy.Metadata.Version = ""
					if itemCopy.Metadata.Annotations != nil {
						delete(itemCopy.Metadata.Annotations, "kubectl.kubernetes.io/last-applied-configuration")
						if len(itemCopy.Metadata.Annotations) == 0 {
							itemCopy.Metadata.Annotations = nil
						}
					}
				}

				if err := rangeCheck.ValidateItems(a, p); err != nil {
					return err
				}
			}
			return nil
		})
	}

}

// Returns success for Envoy configs containing routes for all of specified domains.
func domainAcceptFunc(domains []string) func(*envoyAdmin.ConfigDump) (bool, error) {
	return func(cfg *envoyAdmin.ConfigDump) (bool, error) {
		validator := structpath.ForProto(cfg)
		const q = "{.configs[*].dynamicRouteConfigs[*].routeConfig.virtualHosts[*].domains[?(@ == %q)]}"
		for _, domain := range domains {
			// TODO(qfel): Figure out how to get rid of the loop.
			if err := validator.Exists(q, domain).Check(); err != nil {
				return false, err
			}
		}
		return true, nil
	}
}

func virtualServiceHosts(ctx test.Failer, resources string) []string {
	// TODO(qfel): We should parse input early and work on structured data here. Or at least if some
	// other code path needs it.
	var inputs []map[interface{}]interface{}
	for _, inputYAML := range strings.Split(resources, "\n---\n") {
		var input map[interface{}]interface{}
		if err := yaml.Unmarshal([]byte(inputYAML), &input); err != nil {
			ctx.Fatalf("yaml.Unmarshal: %v", err)
		}
		inputs = append(inputs, input)
	}

	var vHosts []string
	for _, res := range inputs {
		if res["apiVersion"] != "networking.istio.io/v1alpha3" || res["kind"] != "VirtualService" {
			continue
		}
		spec := res["spec"].(map[interface{}]interface{})
		hosts := spec["hosts"].([]interface{})
		for _, h := range hosts {
			vHosts = append(vHosts, h.(string))
		}
	}
	return vHosts
}

func validateTraffic(t framework.TestContext, pil pilot.Instance, gal galley.Instance, ns namespace.Instance, stage *conformance.Stage) {
	echos := make([]echo.Instance, len(stage.Traffic.Services))
	b := echoboot.NewBuilderOrFail(t, t)
	for i, svc := range stage.Traffic.Services {
		var ports []echo.Port
		for _, p := range svc.Ports {
			ports = append(ports, echo.Port{
				Name:        p.Name,
				Protocol:    protocol.Instance(p.Protocol),
				ServicePort: int(p.ServicePort),
			})
		}
		b = b.With(&echos[i], echo.Config{
			Galley:    gal,
			Pilot:     pil,
			Service:   svc.Name,
			Namespace: ns,
			Ports:     ports,
		})
	}
	if err := b.Build(); err != nil {
		t.Fatal(err)
	}

	services := make(map[string]echo.Instance)
	for i, svc := range echos {
		services[stage.Traffic.Services[i].Name] = svc
		svc.WaitUntilCallableOrFail(t, echos...)
	}

	ready := make(map[string]bool)
	vHosts := virtualServiceHosts(t, stage.Input)
	for i, call := range stage.Traffic.Calls {
		call.URL = tmpl.EvaluateOrFail(t, call.URL, constraint.Params{
			Namespace: ns.Name(),
		})
		t.NewSubTest(strconv.Itoa(i)).Run(func(t framework.TestContext) {
			hostname, err := wildcardToRegexp(call.Response.Callee)
			if err != nil {
				t.Fatalf("Internal error: regexp.Compile: %v", err)
			}
			caller := services[call.Caller]
			if !ready[call.Caller] && false {
				t.Logf("Waiting for sidecar(s) for %s to contain domains: %s", call.Caller, strings.Join(vHosts, ", "))
				for _, w := range caller.WorkloadsOrFail(t) {
					w.Sidecar().WaitForConfigOrFail(t, domainAcceptFunc(vHosts))
				}
				ready[call.Caller] = true
			}
			u, err := url.Parse(call.URL)
			if err != nil {
				t.Fatalf("Parse call URL: %v", err)
			}
			workload := caller.WorkloadsOrFail(t)[0]

			validateWithRedo(t, func(ctx context.Context) bool {
				responses, err := workload.ForwardEcho(ctx, &epb.ForwardEchoRequest{
					Url:   call.URL,
					Count: 1,
					Headers: []*epb.Header{{
						Key:   "Host",
						Value: u.Host, // TODO(qfel): Why cannot echo infer it from URL?
					}},
				})
				if err != nil {
					t.Logf("Initiating test call: %v", err)
					return false
				}

				if len(responses) != 1 {
					t.Logf("Received %d responses, want 1", len(responses))
					return false
				}
				resp := responses[0]
				var fail bool
				if resp.Code != strconv.Itoa(call.Response.Code) {
					t.Logf("Responded with %s, want %d", resp.Code, call.Response.Code)
					fail = true
				}
				if !hostname.MatchString(resp.Hostname) {
					t.Logf("Callee %q doesn't match %q", resp.Hostname, call.Response.Callee)
					fail = true
				}
				return !fail
			})
		})
	}
}

func validateWithRedo(t framework.TestContext, f func(context.Context) bool) {
	const (
		pollDelay           = time.Second // How much to wait after a failed attempt.
		firstSuccessTimeout = time.Minute
		// After the traffic flows successfully for the first time, repeat according to the parameters
		// below.
		redoAttempts = 10
		redoTimeout  = 20 * time.Second
		redoFraction = 0.9 // Fail if not enough attempts succeed.
	)

	ctx, cfn := context.WithTimeout(context.Background(), firstSuccessTimeout)
	defer cfn()
	for {
		if f(ctx) {
			break
		}
		if sleepCtx(ctx, pollDelay) != nil {
			t.Fatal("No successful attempt")
		}
	}

	var passed int
	ctx, cfn = context.WithTimeout(context.Background(), redoTimeout)
	defer cfn()
	for i := 0; i < redoAttempts; i++ {
		start := time.Now()
		res := f(ctx)
		t.Logf("Traffic attempt #%d: %v in %s", i, res, time.Since(start))
		if res {
			passed++
		} else if i+1 < redoAttempts {
			// No need to check errors. If ctx times out, further calls to f will time out.
			_ = sleepCtx(ctx, pollDelay)
		}
	}
	fr := float64(passed) / float64(redoAttempts)
	t.Logf("%d/%d (%f) attempts succeeded", passed, redoAttempts, fr)
	if fr < redoFraction {
		t.Errorf("Success rate is %f, need at least %f", fr, redoFraction)
	}
}

func sleepCtx(ctx context.Context, duration time.Duration) error {
	t := time.NewTimer(duration)
	select {
	case <-t.C:
		return nil
	case <-ctx.Done():
		t.Stop()
		return ctx.Err()
	}
}

func wildcardToRegexp(s string) (*regexp.Regexp, error) {
	s = fmt.Sprintf("^%s$", regexp.QuoteMeta(s))
	return regexp.Compile(strings.Replace(s, `\*`, ".*", -1))
}
<|endoftext|>


<commit_before>package rpcc

import (
	"context"
	"encoding/json"
	"errors"
	"sync"
)

var (
	// ErrStreamClosing indicates that the operation is illegal because
	// the stream is closing and there are no pending messages.
	ErrStreamClosing = errors.New("rpcc: the stream is closing")
)

type streamMsg struct {
	method string
	data   []byte
}

type messageBuffer struct {
	ch    chan *streamMsg
	mu    sync.Mutex
	queue []*streamMsg
}

func newMessageBuffer() *messageBuffer {
	return &messageBuffer{
		ch: make(chan *streamMsg, 1),
	}
}

func (b *messageBuffer) store(m *streamMsg) {
	b.mu.Lock()
	defer b.mu.Unlock()
	if len(b.queue) == 0 {
		select {
		case b.ch <- m:
			return
		default:
		}
	}
	b.queue = append(b.queue, m)
}

func (b *messageBuffer) load() {
	b.mu.Lock()
	defer b.mu.Unlock()
	if len(b.queue) > 0 {
		select {
		case b.ch <- b.queue[0]:
			b.queue[0] = nil // Remove reference from underlying array.
			b.queue = b.queue[1:]
		default:
		}
	}
}

func (b *messageBuffer) get() <-chan *streamMsg {
	return b.ch
}

// Stream represents a stream of notifications for a certain method.
type Stream interface {
	// RecvMsg unmarshals pending messages onto m. Blocks until the
	// next message is received, context is canceled or stream is
	// closed.
	//
	// When m is a *[]byte the message will not be decoded and the
	// raw bytes are copied into m.
	RecvMsg(m interface{}) error
	// Close closes the stream and no new messages will be received.
	// RecvMsg will return ErrStreamClosing once all pending messages
	// have been received.
	Close() error
}

// NewStream creates a new stream that listens to notifications from the
// RPC server. This function is called by generated code.
func NewStream(ctx context.Context, method string, conn *Conn) (Stream, error) {
	if ctx == nil {
		ctx = context.Background()
	}

	s := &streamClient{userCtx: ctx, done: make(chan struct{})}
	s.msgBuf = newMessageBuffer()
	s.ctx, s.cancel = context.WithCancel(context.Background())

	var err error
	s.remove, err = conn.listen(method, s)
	if err != nil {
		return nil, err
	}

	go func() {
		select {
		case <-s.ctx.Done():
		case <-conn.ctx.Done():
			s.close(ErrConnClosing)
		case <-ctx.Done():
			s.close(ctx.Err())
		}
	}()

	return s, nil
}

type streamClient struct {
	userCtx context.Context
	ctx     context.Context
	cancel  context.CancelFunc

	// msgBuf stores all incoming messages
	// until they are ready to be received.
	msgBuf *messageBuffer

	mu     sync.Mutex // Protects following.
	remove func()     // Unsubscribes from messages.

	done chan struct{} // Protects err.
	err  error
}

func (s *streamClient) RecvMsg(m interface{}) (err error) {
	msg, err := s.recv()
	if err != nil {
		return err
	}

	if m, ok := m.(*[]byte); ok {
		*m = append(*m, msg.data...)
		return nil
	}

	return json.Unmarshal(msg.data, m)
}

func (s *streamClient) recv() (m *streamMsg, err error) {
	userCancelled := func() bool {
		select {
		case <-s.userCtx.Done():
			return true
		default:
			return false
		}
	}

	select {
	case <-s.userCtx.Done():
		return m, s.userCtx.Err()
	case <-s.ctx.Done():
		// Give precedence for user cancellation.
		if userCancelled() {
			return m, s.userCtx.Err()
		}

		// Send all messages before returning error.
		select {
		case m = <-s.msgBuf.get():
		default:
			<-s.done
			return m, s.err
		}
	case m = <-s.msgBuf.get():
		// Give precedence for user cancellation.
		if userCancelled() {
			return m, s.userCtx.Err()
		}
	}

	// Preload the next message.
	s.msgBuf.load()

	return m, nil
}

// Close closes the stream client.
func (s *streamClient) close(err error) error {
	s.mu.Lock()
	remove := s.remove
	s.remove = nil
	s.mu.Unlock()

	if remove == nil {
		return errors.New("rpcc: the stream is already closed")
	}

	if err == nil {
		err = ErrStreamClosing
	}

	// Unsubscribe first to prevent incoming messages.
	remove()
	s.cancel()
	s.err = err
	close(s.done)

	return nil
}

// Close closes the stream client.
func (s *streamClient) Close() error {
	return s.close(nil)
}

// streamClients handles multiple instances of streamClient and
// enables sending of the same message to multiple clients.
type streamClients struct {
	mu      sync.Mutex
	seq     uint64
	clients map[uint64]*streamClient
}

func newStreamService() *streamClients {
	return &streamClients{
		clients: make(map[uint64]*streamClient),
	}
}

func (s *streamClients) add(client *streamClient) (seq uint64) {
	s.mu.Lock()
	seq = s.seq
	s.seq++
	s.clients[seq] = client
	s.mu.Unlock()
	return seq
}

func (s *streamClients) remove(seq uint64) {
	s.mu.Lock()
	delete(s.clients, seq)
	s.mu.Unlock()
}

func (s *streamClients) send(method string, args []byte) {
	s.mu.Lock()
	for _, client := range s.clients {
		client.msgBuf.store(&streamMsg{method: method, data: args})
	}
	s.mu.Unlock()
}
<commit_msg>rpcc: Prevent loss of stream msg on context cancellation<commit_after>package rpcc

import (
	"context"
	"encoding/json"
	"errors"
	"sync"
)

var (
	// ErrStreamClosing indicates that the operation is illegal because
	// the stream is closing and there are no pending messages.
	ErrStreamClosing = errors.New("rpcc: the stream is closing")
)

type streamMsg struct {
	method string
	data   []byte
}

type messageBuffer struct {
	ch    chan *streamMsg
	mu    sync.Mutex
	queue []*streamMsg
}

func newMessageBuffer() *messageBuffer {
	return &messageBuffer{
		ch: make(chan *streamMsg, 1),
	}
}

func (b *messageBuffer) store(m *streamMsg) {
	b.mu.Lock()
	defer b.mu.Unlock()
	if len(b.queue) == 0 {
		select {
		case b.ch <- m:
			return
		default:
		}
	}
	b.queue = append(b.queue, m)
}

func (b *messageBuffer) load() {
	b.mu.Lock()
	defer b.mu.Unlock()
	if len(b.queue) > 0 {
		select {
		case b.ch <- b.queue[0]:
			b.queue[0] = nil // Remove reference from underlying array.
			b.queue = b.queue[1:]
		default:
		}
	}
}

func (b *messageBuffer) get() <-chan *streamMsg {
	return b.ch
}

// Stream represents a stream of notifications for a certain method.
type Stream interface {
	// RecvMsg unmarshals pending messages onto m. Blocks until the
	// next message is received, context is canceled or stream is
	// closed.
	//
	// When m is a *[]byte the message will not be decoded and the
	// raw bytes are copied into m.
	RecvMsg(m interface{}) error
	// Close closes the stream and no new messages will be received.
	// RecvMsg will return ErrStreamClosing once all pending messages
	// have been received.
	Close() error
}

// NewStream creates a new stream that listens to notifications from the
// RPC server. This function is called by generated code.
func NewStream(ctx context.Context, method string, conn *Conn) (Stream, error) {
	if ctx == nil {
		ctx = context.Background()
	}

	s := &streamClient{userCtx: ctx, done: make(chan struct{})}
	s.msgBuf = newMessageBuffer()
	s.ctx, s.cancel = context.WithCancel(context.Background())

	var err error
	s.remove, err = conn.listen(method, s)
	if err != nil {
		return nil, err
	}

	go func() {
		select {
		case <-s.ctx.Done():
		case <-conn.ctx.Done():
			s.close(ErrConnClosing)
		case <-ctx.Done():
			s.close(ctx.Err())
		}
	}()

	return s, nil
}

type streamClient struct {
	userCtx context.Context
	ctx     context.Context
	cancel  context.CancelFunc

	// msgBuf stores all incoming messages
	// until they are ready to be received.
	msgBuf *messageBuffer

	mu     sync.Mutex // Protects following.
	remove func()     // Unsubscribes from messages.

	done chan struct{} // Protects err.
	err  error
}

func (s *streamClient) RecvMsg(m interface{}) (err error) {
	msg, err := s.recv()
	if err != nil {
		return err
	}

	if m, ok := m.(*[]byte); ok {
		*m = append(*m, msg.data...)
		return nil
	}

	return json.Unmarshal(msg.data, m)
}

func (s *streamClient) recv() (m *streamMsg, err error) {
	userCancelled := func() bool {
		select {
		case <-s.userCtx.Done():
			return true
		default:
			return false
		}
	}

	// Check cancellation once here to avoid race in select.
	if userCancelled() {
		return m, s.userCtx.Err()
	}

	select {
	case <-s.userCtx.Done():
		return m, s.userCtx.Err()
	case <-s.ctx.Done():
		// Give precedence for user cancellation.
		if userCancelled() {
			return m, s.userCtx.Err()
		}

		// Send all messages before returning error.
		select {
		case m = <-s.msgBuf.get():
		default:
			<-s.done
			return m, s.err
		}
	case m = <-s.msgBuf.get():
		// We could check for userCancelled here,
		// but this message would be lost.
	}

	// Preload the next message.
	s.msgBuf.load()

	return m, nil
}

// Close closes the stream client.
func (s *streamClient) close(err error) error {
	s.mu.Lock()
	remove := s.remove
	s.remove = nil
	s.mu.Unlock()

	if remove == nil {
		return errors.New("rpcc: the stream is already closed")
	}

	if err == nil {
		err = ErrStreamClosing
	}

	// Unsubscribe first to prevent incoming messages.
	remove()
	s.cancel()
	s.err = err
	close(s.done)

	return nil
}

// Close closes the stream client.
func (s *streamClient) Close() error {
	return s.close(nil)
}

// streamClients handles multiple instances of streamClient and
// enables sending of the same message to multiple clients.
type streamClients struct {
	mu      sync.Mutex
	seq     uint64
	clients map[uint64]*streamClient
}

func newStreamService() *streamClients {
	return &streamClients{
		clients: make(map[uint64]*streamClient),
	}
}

func (s *streamClients) add(client *streamClient) (seq uint64) {
	s.mu.Lock()
	seq = s.seq
	s.seq++
	s.clients[seq] = client
	s.mu.Unlock()
	return seq
}

func (s *streamClients) remove(seq uint64) {
	s.mu.Lock()
	delete(s.clients, seq)
	s.mu.Unlock()
}

func (s *streamClients) send(method string, args []byte) {
	s.mu.Lock()
	for _, client := range s.clients {
		client.msgBuf.store(&streamMsg{method: method, data: args})
	}
	s.mu.Unlock()
}
<|endoftext|>


<commit_before>package main

import (
	"database/sql"
	"encoding/json"
	"flag"
	"fmt"
	"html/template"
	"io/ioutil"
	"net/http"
	"os"
	"regexp"
	"time"

	_ "github.com/mattn/go-sqlite3"
)

var settings struct {
	DailyTarget int
	Database    string
}

func init() {
	flag.IntVar(&settings.DailyTarget, "target", 900, "The number of words to write daily")
	flag.StringVar(&settings.Database, "db", "diary.db", "The name of the database file")
}

func main() {
	flag.Parse()

	db, err := sql.Open("sqlite3", settings.Database)
	if err != nil {
		panic(err)
	}

	_, err = db.Exec("CREATE TABLE IF NOT EXISTS entries (date TEXT PRIMARY KEY, text TEXT, words INTEGER)")
	if err != nil {
		panic(err)
	}

	if flag.NArg() >= 1 {
		cmd := flag.Arg(0)
		if cmd == "import" {
			importEntry(db)
			return
		}
	}

	http.HandleFunc("/", func(w http.ResponseWriter, req *http.Request) {
		now := time.Now()
		days := daysOfMonth(now)
		annotatedDays, err := annotateDays(db, days)
		if err != nil {
			fmt.Fprintf(os.Stderr, "Error: %s\n", err)
		}

		text := ""
		words := 0
		rows, err := db.Query("SELECT text, words FROM entries WHERE date = ?", now.Format("2006-01-02"))
		if err != nil {
			fmt.Fprintf(os.Stderr, "Error: %s\n", err)
		}
		defer rows.Close()

		if rows.Next() {
			err = rows.Scan(&text, &words)
			if err != nil {
				fmt.Fprintf(os.Stderr, "Error: %s\n", err)
			}
		}

		err = indexTmpl.Execute(w, map[string]interface{}{
			"Title": fmt.Sprintf("%d words", settings.DailyTarget),

			"Now":  now,
			"Days": annotatedDays,

			"Text":  text,
			"Words": words,
		})
		if err != nil {
			fmt.Fprintf(os.Stderr, "Error: %s\n", err)
		}
	})

	http.HandleFunc("/save", func(w http.ResponseWriter, req *http.Request) {
		var entry map[string]string
		dec := json.NewDecoder(req.Body)
		err := dec.Decode(&entry)
		if err != nil {
			respondWithError(w, http.StatusBadRequest, err)
			return
		}

		text, ok := entry["text"]
		if !ok {
			respondWithError(w, http.StatusBadRequest, fmt.Errorf("Missing field 'text'"))
			return
		}

		now := time.Now()
		err = saveEntry(db, now, text)
		if err != nil {
			respondWithError(w, http.StatusInternalServerError, fmt.Errorf("%s", http.StatusText(http.StatusInternalServerError)))
			return
		}

		enc := json.NewEncoder(w)
		err = enc.Encode(map[string]string{"message": "saved post", "time": now.Format(time.RFC3339)})
		if err != nil {
			fmt.Fprintf(os.Stderr, "Error: %s\n", err)
		}
	})

	http.ListenAndServe("localhost:12345", nil)
}

func respondWithError(w http.ResponseWriter, status int, err error) {
	res := map[string]string{"error": err.Error()}
	out, err := json.Marshal(res)
	if err != nil {
		w.WriteHeader(http.StatusInternalServerError)
		fmt.Fprintf(w, "{\"error\": %q}", http.StatusText(http.StatusInternalServerError))
		return
	}
	w.WriteHeader(status)
	w.Write(out)
}

func saveEntry(db *sql.DB, date time.Time, text string) error {
	_, err := db.Exec("INSERT OR REPLACE INTO entries VALUES (?, ?, ?)", date.Format("2006-01-02"), text, countWords(text))
	return err
}

func importEntry(db *sql.DB) {
	if flag.NArg() < 2 {
		fmt.Fprintf(os.Stderr, "Usage: %s import <date>\n", os.Args[0])
		os.Exit(1)
	}

	rawDate := flag.Arg(1)
	date, err := time.Parse("2006-01-02", rawDate)
	if err != nil {
		fmt.Fprintf(os.Stderr, "Error: invalid date '%s': %s\n", rawDate, err)
		os.Exit(1)
	}

	f := os.Stdin
	if flag.NArg() >= 3 {
		f, err = os.Open(flag.Arg(2))
		if err != nil {
			fmt.Fprintf(os.Stderr, "Error: %s\n", err)
			os.Exit(1)
		}
	}
	defer f.Close()

	text, err := ioutil.ReadAll(f)
	if err != nil {
		fmt.Fprintf(os.Stderr, "Error: reading %s: %s\n", f.Name(), err)
		os.Exit(1)
	}

	err = saveEntry(db, date, string(text))
	if err != nil {
		fmt.Fprintf(os.Stderr, "Error: saving entry: %s\n", err)
		os.Exit(1)
	}
}

var indexTmpl = template.Must(template.New("index").Parse(`<!doctype html>

		<meta charset="utf-8" />
		<title>{{ .Title }}</title>

{{ .Title }}

			{{ $now := .Now }}
			{{ range $day := .Days -}}

{{ $day.Date.Day }}

			{{ end }}

{{ .Now.Format "Monday, January 2, 2006" }}

				<textarea id="editor">{{ .Text }}</textarea>

					0 words

				Made with <3 by strange adventures.  —  /about

`))

var wordSeperator = regexp.MustCompile(`\s+`)

func countWords(text string) int {
	ws := wordSeperator.Split(text, -1)
	c := 0
	for _, w := range ws {
		if w != "" {
			c += 1
		}
	}
	return c
}

func daysOfMonth(t time.Time) []time.Time {
	s := time.Date(t.Year(), t.Month(), 1, 0, 0, 0, 0, t.Location())
	e := s.AddDate(0, 1, 0)
	ts := make([]time.Time, 0, 31)
	for s.Before(e) {
		ts = append(ts, s)
		s = s.AddDate(0, 0, 1)
	}
	return ts
}

type Day struct {
	Date  time.Time
	Words int
}

func (d Day) Classes(now time.Time) string {
	var classes string
	if d.Date.Before(now) {
		classes = "past"
	} else {
		classes = "future"
	}

	if d.Words >= 1 {
		classes += " written"
	}

	if d.Words >= settings.DailyTarget {
		classes += " yay"
	}

	return classes
}

func annotateDays(db *sql.DB, days []time.Time) ([]Day, error) {
	rows, err := db.Query("SELECT date, words FROM entries WHERE date >= ? AND date <= ?", days[0].Format("2006-01-02"), days[len(days)-1].Format("2006-01-02"))
	if err != nil {
		return nil, err
	}
	defer rows.Close()

	var day string
	var words int
	if rows.Next() {
		err = rows.Scan(&day, &words)
		if err != nil {
			return nil, err
		}
	}

	annotatedDays := make([]Day, len(days))
	for i, t := range days {
		d := Day{Date: t, Words: 0}

		if t.Format("2006-01-02") == day {
			d.Words = words

			if rows.Next() {
				err = rows.Scan(&day, &words)
				if err != nil {
					return nil, err
				}
			}
		}

		annotatedDays[i] = d
	}

	return annotatedDays, nil
}
<commit_msg>Move stats display to the right<commit_after>package main

import (
	"database/sql"
	"encoding/json"
	"flag"
	"fmt"
	"html/template"
	"io/ioutil"
	"net/http"
	"os"
	"regexp"
	"time"

	_ "github.com/mattn/go-sqlite3"
)

var settings struct {
	DailyTarget int
	Database    string
}

func init() {
	flag.IntVar(&settings.DailyTarget, "target", 900, "The number of words to write daily")
	flag.StringVar(&settings.Database, "db", "diary.db", "The name of the database file")
}

func main() {
	flag.Parse()

	db, err := sql.Open("sqlite3", settings.Database)
	if err != nil {
		panic(err)
	}

	_, err = db.Exec("CREATE TABLE IF NOT EXISTS entries (date TEXT PRIMARY KEY, text TEXT, words INTEGER)")
	if err != nil {
		panic(err)
	}

	if flag.NArg() >= 1 {
		cmd := flag.Arg(0)
		if cmd == "import" {
			importEntry(db)
			return
		}
	}

	http.HandleFunc("/", func(w http.ResponseWriter, req *http.Request) {
		now := time.Now()
		days := daysOfMonth(now)
		annotatedDays, err := annotateDays(db, days)
		if err != nil {
			fmt.Fprintf(os.Stderr, "Error: %s\n", err)
		}

		text := ""
		words := 0
		rows, err := db.Query("SELECT text, words FROM entries WHERE date = ?", now.Format("2006-01-02"))
		if err != nil {
			fmt.Fprintf(os.Stderr, "Error: %s\n", err)
		}
		defer rows.Close()

		if rows.Next() {
			err = rows.Scan(&text, &words)
			if err != nil {
				fmt.Fprintf(os.Stderr, "Error: %s\n", err)
			}
		}

		err = indexTmpl.Execute(w, map[string]interface{}{
			"Title": fmt.Sprintf("%d words", settings.DailyTarget),

			"Now":  now,
			"Days": annotatedDays,

			"Text":  text,
			"Words": words,
		})
		if err != nil {
			fmt.Fprintf(os.Stderr, "Error: %s\n", err)
		}
	})

	http.HandleFunc("/save", func(w http.ResponseWriter, req *http.Request) {
		var entry map[string]string
		dec := json.NewDecoder(req.Body)
		err := dec.Decode(&entry)
		if err != nil {
			respondWithError(w, http.StatusBadRequest, err)
			return
		}

		text, ok := entry["text"]
		if !ok {
			respondWithError(w, http.StatusBadRequest, fmt.Errorf("Missing field 'text'"))
			return
		}

		now := time.Now()
		err = saveEntry(db, now, text)
		if err != nil {
			respondWithError(w, http.StatusInternalServerError, fmt.Errorf("%s", http.StatusText(http.StatusInternalServerError)))
			return
		}

		enc := json.NewEncoder(w)
		err = enc.Encode(map[string]string{"message": "saved post", "time": now.Format(time.RFC3339)})
		if err != nil {
			fmt.Fprintf(os.Stderr, "Error: %s\n", err)
		}
	})

	http.ListenAndServe("localhost:12345", nil)
}

func respondWithError(w http.ResponseWriter, status int, err error) {
	res := map[string]string{"error": err.Error()}
	out, err := json.Marshal(res)
	if err != nil {
		w.WriteHeader(http.StatusInternalServerError)
		fmt.Fprintf(w, "{\"error\": %q}", http.StatusText(http.StatusInternalServerError))
		return
	}
	w.WriteHeader(status)
	w.Write(out)
}

func saveEntry(db *sql.DB, date time.Time, text string) error {
	_, err := db.Exec("INSERT OR REPLACE INTO entries VALUES (?, ?, ?)", date.Format("2006-01-02"), text, countWords(text))
	return err
}

func importEntry(db *sql.DB) {
	if flag.NArg() < 2 {
		fmt.Fprintf(os.Stderr, "Usage: %s import <date>\n", os.Args[0])
		os.Exit(1)
	}

	rawDate := flag.Arg(1)
	date, err := time.Parse("2006-01-02", rawDate)
	if err != nil {
		fmt.Fprintf(os.Stderr, "Error: invalid date '%s': %s\n", rawDate, err)
		os.Exit(1)
	}

	f := os.Stdin
	if flag.NArg() >= 3 {
		f, err = os.Open(flag.Arg(2))
		if err != nil {
			fmt.Fprintf(os.Stderr, "Error: %s\n", err)
			os.Exit(1)
		}
	}
	defer f.Close()

	text, err := ioutil.ReadAll(f)
	if err != nil {
		fmt.Fprintf(os.Stderr, "Error: reading %s: %s\n", f.Name(), err)
		os.Exit(1)
	}

	err = saveEntry(db, date, string(text))
	if err != nil {
		fmt.Fprintf(os.Stderr, "Error: saving entry: %s\n", err)
		os.Exit(1)
	}
}

var indexTmpl = template.Must(template.New("index").Parse(`<!doctype html>

		<meta charset="utf-8" />
		<title>{{ .Title }}</title>

{{ .Title }}

			{{ $now := .Now }}
			{{ range $day := .Days -}}

{{ $day.Date.Day }}

			{{ end }}

{{ .Now.Format "Monday, January 2, 2006" }}

				<textarea id="editor">{{ .Text }}</textarea>

					0 words

				Made with <3 by strange adventures.  —  /about

`))

var wordSeperator = regexp.MustCompile(`\s+`)

func countWords(text string) int {
	ws := wordSeperator.Split(text, -1)
	c := 0
	for _, w := range ws {
		if w != "" {
			c += 1
		}
	}
	return c
}

func daysOfMonth(t time.Time) []time.Time {
	s := time.Date(t.Year(), t.Month(), 1, 0, 0, 0, 0, t.Location())
	e := s.AddDate(0, 1, 0)
	ts := make([]time.Time, 0, 31)
	for s.Before(e) {
		ts = append(ts, s)
		s = s.AddDate(0, 0, 1)
	}
	return ts
}

type Day struct {
	Date  time.Time
	Words int
}

func (d Day) Classes(now time.Time) string {
	var classes string
	if d.Date.Before(now) {
		classes = "past"
	} else {
		classes = "future"
	}

	if d.Words >= 1 {
		classes += " written"
	}

	if d.Words >= settings.DailyTarget {
		classes += " yay"
	}

	return classes
}

func annotateDays(db *sql.DB, days []time.Time) ([]Day, error) {
	rows, err := db.Query("SELECT date, words FROM entries WHERE date >= ? AND date <= ?", days[0].Format("2006-01-02"), days[len(days)-1].Format("2006-01-02"))
	if err != nil {
		return nil, err
	}
	defer rows.Close()

	var day string
	var words int
	if rows.Next() {
		err = rows.Scan(&day, &words)
		if err != nil {
			return nil, err
		}
	}

	annotatedDays := make([]Day, len(days))
	for i, t := range days {
		d := Day{Date: t, Words: 0}

		if t.Format("2006-01-02") == day {
			d.Words = words

			if rows.Next() {
				err = rows.Scan(&day, &words)
				if err != nil {
					return nil, err
				}
			}
		}

		annotatedDays[i] = d
	}

	return annotatedDays, nil
}
<|endoftext|>


<commit_before>package fullerene

import (
	"fmt"
	libage "github.com/bearbin/go-age"
	"github.com/stretchr/testify/assert"
	"strconv"
	"testing"
	"time"
)

func TestGetCurrentTime(t *testing.T) {
	mt, tt := Now(), time.Now()
	mty, mtm, mtd := mt.Date()
	tty, ttm, ttd := tt.Date()
	assert.Equal(t, mty, tty)
	assert.Equal(t, mtm, ttm)
	assert.Equal(t, mtd, ttd)

	assert.True(t, mt.Equal(mt))
}

func TestFullerene_After(t *testing.T) {
	mt, tt := Now(), time.Now()
	assert.False(t, mt.After(Fullerene{t: tt}))
}

func TestFullerene_Before(t *testing.T) {
	mt, tt := Now(), time.Now()
	assert.True(t, mt.Before(Fullerene{t: tt}))
}

func TestFullerene_IsZero(t *testing.T) {
	assert.True(t, Fullerene{}.IsZero())
	assert.False(t, Now().IsZero())
}

func TestFullerene_Equal(t *testing.T) {
	mt := Now()
	assert.True(t, mt.Equal(mt))
}

func TestFullerene_IsLeapYear(t *testing.T) {
	fr := Date(2016, 1, 1, 0, 0, 0, 0, &time.Location{})
	assert.True(t, fr.IsLeapYear())
	fr2 := Date(2000, 1, 1, 0, 0, 0, 0, &time.Location{})
	assert.True(t, fr2.IsLeapYear())
	fr3 := Date(1999, 1, 1, 0, 0, 0, 0, &time.Location{})
	assert.False(t, fr3.IsLeapYear())
	fr4 := Date(2100, 1, 1, 0, 0, 0, 0, &time.Location{})
	assert.False(t, fr4.IsLeapYear())
}

func TestFullerene_IsLeapDay(t *testing.T) {
	fr := Date(2016, 2, 29, 0, 0, 0, 0, &time.Location{})
	assert.True(t, fr.IsLeapDay())

	fr2 := Date(2016, 2, 28, 0, 0, 0, 0, &time.Location{})
	assert.False(t, fr2.IsLeapDay())
}

func TestFullerene_IsBirthday(t *testing.T) {
	emptyLocation := time.Location{}
	fr := Date(2014, 11, 18, 0, 0, 0, 0, &emptyLocation)
	assert.True(t, fr.IsBirthday(Date(1981, 11, 18, 0, 0, 0, 0, &emptyLocation), false))
	assert.False(t, fr.IsBirthday(Date(1981, 11, 19, 0, 0, 0, 0, &emptyLocation), false))

	frLeap := Date(1980, 2, 29, 0, 0, 0, 0, &emptyLocation)
	assert.False(t, frLeap.IsBirthday(Date(2015, 2, 28, 0, 0, 0, 0, &emptyLocation), false))
	assert.True(t, frLeap.IsBirthday(Date(2015, 3, 1, 0, 0, 0, 0, &emptyLocation), false))
	assert.True(t, frLeap.IsBirthday(Date(2015, 2, 28, 0, 0, 0, 0, &emptyLocation), true))
	assert.False(t, frLeap.IsBirthday(Date(2015, 3, 1, 0, 0, 0, 0, &emptyLocation), true))
	assert.False(t, frLeap.IsBirthday(Date(2016, 3, 1, 0, 0, 0, 0, &emptyLocation), true))
	assert.False(t, frLeap.IsBirthday(Date(2016, 3, 1, 0, 0, 0, 0, &emptyLocation), false))
	tm := time.Now()
	tm.Weekday()
	fr.isBirthdayEx(frLeap, false)
}

func TestFullerene_Age(t *testing.T) {
	type Case struct {
		Birthday    Fullerene
		ExpectedAge int
	}
	targetDate := Date(2016, 10, 2, 0, 0, 0, 0, &time.Location{})
	tests := []Case{
		{Date(2015, 1, 1, 0, 0, 0, 0, &time.Location{}), 1},
		{Date(2014, 12, 31, 0, 0, 0, 0, &time.Location{}), 1},
		{Date(1999, 7, 16, 0, 0, 0, 0, &time.Location{}), 17},
		{Date(1988, 2, 29, 0, 0, 0, 0, &time.Location{}), 28},
		{Date(1988, 3, 3, 0, 0, 0, 0, &time.Location{}), 28},
		{Date(1988, 10, 1, 0, 0, 0, 0, &time.Location{}), 27},
		{Date(1988, 10, 2, 0, 0, 0, 0, &time.Location{}), 28},
		{Date(1988, 10, 3, 0, 0, 0, 0, &time.Location{}), 28},
	}

	for _, d := range tests {
		assert.Equal(t, d.ExpectedAge, d.Birthday.Age(targetDate), d.Birthday.String())
	}
}

func BenchmarkFullerene_Date(b *testing.B) {
	location := &time.Location{}
	for i := 0; i < b.N; i++ {
		Date(2016, 1, 1, 0, 0, 0, 0, location)
	}
}

func BenchmarkFullerene_Date2(b *testing.B) {
	for i := 0; i < b.N; i++ {
		Date(2016, 1, 1, 0, 0, 0, 0, &time.Location{})
	}
}

func BenchmarkFullerene_Date3(b *testing.B) {
	b.StopTimer()
	fr := Date(2016, 1, 1, 0, 0, 0, 0, &time.Location{})
	b.StartTimer()
	for i := 0; i < b.N; i++ {
		_ = fr.Year()
	}
}

func BenchmarkFullerene_Date4(b *testing.B) {
	b.StopTimer()
	fr := Date(2016, 1, 1, 0, 0, 0, 0, &time.Location{})
	b.StartTimer()
	for i := 0; i < b.N; i++ {
		_, _ = fr.Month(), fr.Day()
	}
}

func BenchmarkFullerene_Date5(b *testing.B) {
	b.StopTimer()
	fr := Date(2016, 1, 1, 0, 0, 0, 0, &time.Location{})
	b.StartTimer()
	for i := 0; i < b.N; i++ {
		_, _, _ = fr.Date()
	}
}

func Benchmark_Date(b *testing.B) {
	b.StopTimer()
	location := &time.Location{}
	b.StartTimer()
	for i := 0; i < b.N; i++ {
		time.Date(2016, 1, 1, 0, 0, 0, 0, location)
	}
}

func Benchmark_Fullerene_Age(b *testing.B) {
	b.StopTimer()
	f := Date(2015, 1, 1, 0, 0, 0, 0, &time.Location{})
	b.StartTimer()
	now := Now()
	for i := 0; i < b.N; i++ {
		f.Age(now)
	}
}

func AgeFromString(birthday, targetDay string) int {
	y1, _ := strconv.Atoi(birthday[0:4])
	y2, _ := strconv.Atoi(targetDay[0:4])
	age := y2 - y1
	m1, _ := strconv.Atoi(birthday[4:6])
	m2, _ := strconv.Atoi(targetDay[4:6])
	if m1 < m2 {
		return age
	}
	if m1 > m2 {
		return age - 1
	}
	d1, _ := strconv.Atoi(birthday[6:8])
	d2, _ := strconv.Atoi(targetDay[6:8])
	if d1 < d2 {
		return age - 1
	}
	return age
}

func TestSimpleAge(t *testing.T) {
	assert.Equal(t, 6, AgeFromString("20100101", "20161020"))
	assert.Equal(t, 6, AgeFromString("20100101", "20160101"))
	assert.Equal(t, 5, AgeFromString("20100101", "20151231"))
}

func BenchmarkSimpleAge(b *testing.B) {
	b.StopTimer()
	f1 := Date(2010, 10, 21, 0, 0, 0, 0, &time.Location{})
	f2 := Date(2010, 10, 22, 0, 0, 0, 0, &time.Location{})
	d1 := (f1.String())[0:12]
	d2 := (f2.String())[0:12]
	fmt.Println(d1)
	fmt.Println(d2)
	b.StartTimer()
	for i := 0; i < b.N; i++ {
		AgeFromString(d1, d2)
	}
}

func BenchmarkFullerene_GoAge(b *testing.B) {
	b.StopTimer()
	t := time.Now().AddDate(0, 0, -1)
	b.StartTimer()
	for i := 0; i < b.N; i++ {
		libage.Age(t)
	}
}
<commit_msg>Add Holiday tests.<commit_after>package fullerene

import (
	"fmt"
	libage "github.com/bearbin/go-age"
	"github.com/stretchr/testify/assert"
	"strconv"
	"testing"
	"time"
)

func TestGetCurrentTime(t *testing.T) {
	mt, tt := Now(), time.Now()
	mty, mtm, mtd := mt.Date()
	tty, ttm, ttd := tt.Date()
	assert.Equal(t, mty, tty)
	assert.Equal(t, mtm, ttm)
	assert.Equal(t, mtd, ttd)

	assert.True(t, mt.Equal(mt))
}

func TestFullerene_After(t *testing.T) {
	mt, tt := Now(), time.Now()
	assert.False(t, mt.After(Fullerene{t: tt}))
}

func TestFullerene_Before(t *testing.T) {
	mt, tt := Now(), time.Now()
	assert.True(t, mt.Before(Fullerene{t: tt}))
}

func TestFullerene_IsZero(t *testing.T) {
	assert.True(t, Fullerene{}.IsZero())
	assert.False(t, Now().IsZero())
}

func TestFullerene_Equal(t *testing.T) {
	mt := Now()
	assert.True(t, mt.Equal(mt))
}

func TestFullerene_IsLeapYear(t *testing.T) {
	fr := Date(2016, 1, 1, 0, 0, 0, 0, &time.Location{})
	assert.True(t, fr.IsLeapYear())
	fr2 := Date(2000, 1, 1, 0, 0, 0, 0, &time.Location{})
	assert.True(t, fr2.IsLeapYear())
	fr3 := Date(1999, 1, 1, 0, 0, 0, 0, &time.Location{})
	assert.False(t, fr3.IsLeapYear())
	fr4 := Date(2100, 1, 1, 0, 0, 0, 0, &time.Location{})
	assert.False(t, fr4.IsLeapYear())
}

func TestFullerene_IsLeapDay(t *testing.T) {
	fr := Date(2016, 2, 29, 0, 0, 0, 0, &time.Location{})
	assert.True(t, fr.IsLeapDay())

	fr2 := Date(2016, 2, 28, 0, 0, 0, 0, &time.Location{})
	assert.False(t, fr2.IsLeapDay())
}

func TestFullerene_IsBirthday(t *testing.T) {
	emptyLocation := time.Location{}
	fr := Date(2014, 11, 18, 0, 0, 0, 0, &emptyLocation)
	assert.True(t, fr.IsBirthday(Date(1981, 11, 18, 0, 0, 0, 0, &emptyLocation), false))
	assert.False(t, fr.IsBirthday(Date(1981, 11, 19, 0, 0, 0, 0, &emptyLocation), false))

	frLeap := Date(1980, 2, 29, 0, 0, 0, 0, &emptyLocation)
	assert.False(t, frLeap.IsBirthday(Date(2015, 2, 28, 0, 0, 0, 0, &emptyLocation), false))
	assert.True(t, frLeap.IsBirthday(Date(2015, 3, 1, 0, 0, 0, 0, &emptyLocation), false))
	assert.True(t, frLeap.IsBirthday(Date(2015, 2, 28, 0, 0, 0, 0, &emptyLocation), true))
	assert.False(t, frLeap.IsBirthday(Date(2015, 3, 1, 0, 0, 0, 0, &emptyLocation), true))
	assert.False(t, frLeap.IsBirthday(Date(2016, 3, 1, 0, 0, 0, 0, &emptyLocation), true))
	assert.False(t, frLeap.IsBirthday(Date(2016, 3, 1, 0, 0, 0, 0, &emptyLocation), false))
	tm := time.Now()
	tm.Weekday()
	fr.isBirthdayEx(frLeap, false)
}

func TestFullerene_Age(t *testing.T) {
	type Case struct {
		Birthday    Fullerene
		ExpectedAge int
	}
	targetDate := Date(2016, 10, 2, 0, 0, 0, 0, &time.Location{})
	tests := []Case{
		{Date(2015, 1, 1, 0, 0, 0, 0, &time.Location{}), 1},
		{Date(2014, 12, 31, 0, 0, 0, 0, &time.Location{}), 1},
		{Date(1999, 7, 16, 0, 0, 0, 0, &time.Location{}), 17},
		{Date(1988, 2, 29, 0, 0, 0, 0, &time.Location{}), 28},
		{Date(1988, 3, 3, 0, 0, 0, 0, &time.Location{}), 28},
		{Date(1988, 10, 1, 0, 0, 0, 0, &time.Location{}), 27},
		{Date(1988, 10, 2, 0, 0, 0, 0, &time.Location{}), 28},
		{Date(1988, 10, 3, 0, 0, 0, 0, &time.Location{}), 28},
	}

	for _, d := range tests {
		assert.Equal(t, d.ExpectedAge, d.Birthday.Age(targetDate), d.Birthday.String())
	}
}

func BenchmarkFullerene_Date(b *testing.B) {
	location := &time.Location{}
	for i := 0; i < b.N; i++ {
		Date(2016, 1, 1, 0, 0, 0, 0, location)
	}
}

func BenchmarkFullerene_Date2(b *testing.B) {
	for i := 0; i < b.N; i++ {
		Date(2016, 1, 1, 0, 0, 0, 0, &time.Location{})
	}
}

func BenchmarkFullerene_Date3(b *testing.B) {
	b.StopTimer()
	fr := Date(2016, 1, 1, 0, 0, 0, 0, &time.Location{})
	b.StartTimer()
	for i := 0; i < b.N; i++ {
		_ = fr.Year()
	}
}

func BenchmarkFullerene_Date4(b *testing.B) {
	b.StopTimer()
	fr := Date(2016, 1, 1, 0, 0, 0, 0, &time.Location{})
	b.StartTimer()
	for i := 0; i < b.N; i++ {
		_, _ = fr.Month(), fr.Day()
	}
}

func BenchmarkFullerene_Date5(b *testing.B) {
	b.StopTimer()
	fr := Date(2016, 1, 1, 0, 0, 0, 0, &time.Location{})
	b.StartTimer()
	for i := 0; i < b.N; i++ {
		_, _, _ = fr.Date()
	}
}

func Benchmark_Date(b *testing.B) {
	b.StopTimer()
	location := &time.Location{}
	b.StartTimer()
	for i := 0; i < b.N; i++ {
		time.Date(2016, 1, 1, 0, 0, 0, 0, location)
	}
}

func Benchmark_Fullerene_Age(b *testing.B) {
	b.StopTimer()
	f := Date(2015, 1, 1, 0, 0, 0, 0, &time.Location{})
	b.StartTimer()
	now := Now()
	for i := 0; i < b.N; i++ {
		f.Age(now)
	}
}

func AgeFromString(birthday, targetDay string) int {
	y1, _ := strconv.Atoi(birthday[0:4])
	y2, _ := strconv.Atoi(targetDay[0:4])
	age := y2 - y1
	m1, _ := strconv.Atoi(birthday[4:6])
	m2, _ := strconv.Atoi(targetDay[4:6])
	if m1 < m2 {
		return age
	}
	if m1 > m2 {
		return age - 1
	}
	d1, _ := strconv.Atoi(birthday[6:8])
	d2, _ := strconv.Atoi(targetDay[6:8])
	if d1 < d2 {
		return age - 1
	}
	return age
}

func TestSimpleAge(t *testing.T) {
	assert.Equal(t, 6, AgeFromString("20100101", "20161020"))
	assert.Equal(t, 6, AgeFromString("20100101", "20160101"))
	assert.Equal(t, 5, AgeFromString("20100101", "20151231"))
}

func TestFullerene_IsHoliday(t *testing.T) {
	f1 := Date(2017, 1, 22, 0, 0, 0, 0, &time.Location{})
	assert.True(t, f1.IsHoliday())
	f2 := Date(2017, 1, 23, 0, 0, 0, 0, &time.Location{})
	assert.False(t, f2.IsHoliday())

}

func TestFullerene_IsJapanesePublicHoliday(t *testing.T) {
	for _, d := range JapanesePublicHolidays {
		assert.True(t, d.IsJapanesePublicHoliday())
	}
}

func BenchmarkSimpleAge(b *testing.B) {
	b.StopTimer()
	f1 := Date(2010, 10, 21, 0, 0, 0, 0, &time.Location{})
	f2 := Date(2010, 10, 22, 0, 0, 0, 0, &time.Location{})
	d1 := (f1.String())[0:12]
	d2 := (f2.String())[0:12]
	fmt.Println(d1)
	fmt.Println(d2)
	b.StartTimer()
	for i := 0; i < b.N; i++ {
		AgeFromString(d1, d2)
	}
}

func BenchmarkFullerene_GoAge(b *testing.B) {
	b.StopTimer()
	t := time.Now().AddDate(0, 0, -1)
	b.StartTimer()
	for i := 0; i < b.N; i++ {
		libage.Age(t)
	}
}
<|endoftext|>


<commit_before>package aspect

import (
	"testing"
)

func TestFunctions(t *testing.T) {
	// COUNT()
	count := Count(views.C["id"])
	expectedSQL(t, count, `COUNT("views"."id")`, 0)

	// DATE()
	date := DateOf(views.C["timestamp"])
	expectedSQL(t, date, `DATE("views"."timestamp")`, 0)

	// DATE_PART()
	datePart := DatePart(views.C["timestamp"], "quarter")
	expectedSQL(
		t,
		datePart,
		`DATE_PART('quarter', "views"."timestamp")`,
		0,
	)
}
<commit_msg>Testing of all column functions<commit_after>package aspect

import (
	"testing"
)

func TestFunctions(t *testing.T) {
	expectedSQL(t, Avg(views.C["id"]), `AVG("views"."id")`, 0)
	expectedSQL(t, Sum(views.C["id"]), `SUM("views"."id")`, 0)
	expectedSQL(t, Count(views.C["id"]), `COUNT("views"."id")`, 0)
	expectedSQL(
		t,
		DateOf(views.C["timestamp"]),
		`DATE("views"."timestamp")`,
		0,
	)
	expectedSQL(t, Max(views.C["id"]), `MAX("views"."id")`, 0)
	expectedSQL(
		t,
		DatePart(views.C["timestamp"], "quarter"),
		`DATE_PART('quarter', "views"."timestamp")`,
		0,
	)
}
<|endoftext|>