<commit_before>package scheduler_test

import (
	"context"
	"errors"
	"os"
	"sync"
	"testing"
	"time"

	"github.com/cozy/cozy-stack/pkg/config"
	"github.com/cozy/cozy-stack/pkg/couchdb"
	"github.com/cozy/cozy-stack/pkg/jobs"
	"github.com/cozy/cozy-stack/pkg/realtime"
	"github.com/cozy/cozy-stack/pkg/scheduler"
	"github.com/cozy/cozy-stack/pkg/stack"
	"github.com/cozy/cozy-stack/tests/testutils"
	"github.com/go-redis/redis"
	"github.com/stretchr/testify/assert"
)

const redisURL = "redis://localhost:6379/15"

var instanceName string

type testDoc struct {
	id      string
	rev     string
	doctype string
}

func (t *testDoc) ID() string      { return t.id }
func (t *testDoc) Rev() string     { return t.rev }
func (t *testDoc) DocType() string { return t.doctype }

type mockBroker struct {
	jobs []*jobs.JobRequest
}

func (b *mockBroker) PushJob(request *jobs.JobRequest) (*jobs.JobInfos, error) {
	b.jobs = append(b.jobs, request)
	return nil, nil
}

func (b *mockBroker) QueueLen(workerType string) (int, error) {
	count := 0
	for _, job := range b.jobs {
		if job.WorkerType == workerType {
			count++
		}
	}
	return count, nil
}

func (b *mockBroker) GetJobInfos(domain, id string) (*jobs.JobInfos, error) {
	return nil, errors.New("Not implemented")
}

func TestRedisSchedulerWithTimeTriggers(t *testing.T) {
	var wAt sync.WaitGroup
	var wIn sync.WaitGroup
	bro := jobs.NewMemBroker(1, jobs.WorkersList{
		"worker": {
			Concurrency:  1,
			MaxExecCount: 1,
			Timeout:      1 * time.Millisecond,
			WorkerFunc: func(ctx context.Context, m *jobs.Message) error {
				var msg string
				if err := m.Unmarshal(&msg); err != nil {
					return err
				}
				switch msg {
				case "@at":
					wAt.Done()
				case "@in":
					wIn.Done()
				}
				return nil
			},
		},
	})

	msg1, _ := jobs.NewMessage("json", "@at")
	msg2, _ := jobs.NewMessage("json", "@in")

	wAt.Add(1) // 1 time in @at
	wIn.Add(1) // 1 time in @in

	at := &scheduler.TriggerInfos{
		Type:       "@at",
		Domain:     instanceName,
		Arguments:  time.Now().Add(2 * time.Second).Format(time.RFC3339),
		WorkerType: "worker",
		Message:    msg1,
	}
	in := &scheduler.TriggerInfos{
		Domain:     instanceName,
		Type:       "@in",
		Arguments:  "1s",
		WorkerType: "worker",
		Message:    msg2,
	}

	sch := stack.GetScheduler().(*scheduler.RedisScheduler)
	sch.Stop()
	sch.Start(bro)

	tat, err := scheduler.NewTrigger(at)
	assert.NoError(t, err)
	err = sch.Add(tat)
	assert.NoError(t, err)
	atID := tat.Infos().TID

	tin, err := scheduler.NewTrigger(in)
	assert.NoError(t, err)
	err = sch.Add(tin)
	assert.NoError(t, err)
	inID := tin.Infos().TID

	ts, err := sch.GetAll(instanceName)
	assert.NoError(t, err)
	assert.Len(t, ts, 3) // 1 @event for thumbnails + 1 @at + 1 @in

	for _, trigger := range ts {
		switch trigger.Infos().TID {
		case atID:
			assert.Equal(t, at, trigger.Infos())
		case inID:
			assert.Equal(t, in, trigger.Infos())
		default:
			// Just ignore the @event trigger for generating thumbnails
			infos := trigger.Infos()
			if infos.Type != "@event" || infos.WorkerType != "thumbnail" {
				t.Fatalf("unknown trigger ID %s", trigger.Infos().TID)
			}
		}
	}

	done := make(chan bool)
	go func() {
		wAt.Wait()
		done <- true
	}()

	go func() {
		wIn.Wait()
		done <- true
	}()

	for i := 0; i < 2; i++ {
		select {
		case <-done:
		case <-time.After(2 * time.Second):
			t.Fatalf("Timeout")
		}
	}

	time.Sleep(50 * time.Millisecond)

	_, err = sch.Get(instanceName, atID)
	assert.Error(t, err)
	assert.Equal(t, scheduler.ErrNotFoundTrigger, err)

	_, err = sch.Get(instanceName, inID)
	assert.Error(t, err)
	assert.Equal(t, scheduler.ErrNotFoundTrigger, err)
}

func TestRedisSchedulerWithCronTriggers(t *testing.T) {
	opts, _ := redis.ParseURL(redisURL)
	client := redis.NewClient(opts)
	err := client.Del(scheduler.TriggersKey, scheduler.SchedKey).Err()
	assert.NoError(t, err)

	bro := &mockBroker{}
	sch := stack.GetScheduler().(*scheduler.RedisScheduler)
	sch.Stop()
	sch.Start(bro)
	sch.Stop()
	defer sch.Start(bro)

	msg, _ := jobs.NewMessage("json", "@cron")

	infos := &scheduler.TriggerInfos{
		Type:       "@cron",
		Domain:     instanceName,
		Arguments:  "*/3 * * * * *",
		WorkerType: "incr",
		Message:    msg,
	}
	trigger, err := scheduler.NewTrigger(infos)
	assert.NoError(t, err)
	err = sch.Add(trigger)
	assert.NoError(t, err)

	now := time.Now().UTC().Unix()
	for i := int64(0); i < 9; i++ {
		err = sch.Poll(now + i + 4)
		assert.NoError(t, err)
	}
	count, _ := bro.QueueLen("incr")
	assert.Equal(t, 4, count)
}

func TestRedisPollFromSchedKey(t *testing.T) {
	opts, _ := redis.ParseURL(redisURL)
	client := redis.NewClient(opts)
	err := client.Del(scheduler.TriggersKey, scheduler.SchedKey).Err()
	assert.NoError(t, err)

	bro := &mockBroker{}
	sch := stack.GetScheduler().(*scheduler.RedisScheduler)
	sch.Stop()
	sch.Start(bro)
	sch.Stop()
	defer sch.Start(bro)

	now := time.Now()
	msg, _ := jobs.NewMessage("json", "@at")

	at := &scheduler.TriggerInfos{
		Type:       "@at",
		Domain:     instanceName,
		Arguments:  now.Format(time.RFC3339),
		WorkerType: "incr",
		Message:    msg,
	}
	db := couchdb.SimpleDatabasePrefix(instanceName)
	err = couchdb.CreateDoc(db, at)
	assert.NoError(t, err)

	ts := now.UTC().Unix()
	key := instanceName + "/" + at.TID
	err = client.ZAdd(scheduler.SchedKey, redis.Z{
		Score:  float64(ts + 1),
		Member: key,
	}).Err()
	assert.NoError(t, err)

	err = sch.Poll(ts + 2)
	assert.NoError(t, err)
	<-time.After(1 * time.Millisecond)
	count, _ := bro.QueueLen("incr")
	assert.Equal(t, 0, count)

	err = sch.Poll(ts + 13)
	assert.NoError(t, err)
	<-time.After(1 * time.Millisecond)
	count, _ = bro.QueueLen("incr")
	assert.Equal(t, 1, count)
}

func TestRedisTriggerEvent(t *testing.T) {
	opts, _ := redis.ParseURL(redisURL)
	client := redis.NewClient(opts)
	err := client.Del(scheduler.TriggersKey, scheduler.SchedKey).Err()
	assert.NoError(t, err)

	bro := &mockBroker{}
	sch := stack.GetScheduler().(*scheduler.RedisScheduler)
	sch.Stop()
	time.Sleep(1 * time.Second)
	sch.Start(bro)

	evTrigger := &scheduler.TriggerInfos{
		Type:       "@event",
		Domain:     instanceName,
		Arguments:  "io.cozy.event-test:CREATED",
		WorkerType: "incr",
	}
	tri, err := scheduler.NewTrigger(evTrigger)
	assert.NoError(t, err)
	sch.Add(tri)

	realtime.GetHub().Publish(&realtime.Event{
		Domain: instanceName,
		Doc: &testDoc{
			id:      "foo",
			doctype: "io.cozy.event-test",
		},
		Type: realtime.EventCreate,
	})

	time.Sleep(10 * time.Millisecond)

	count, _ := bro.QueueLen("incr")
	if !assert.Equal(t, 1, count) {
		return
	}

	type eventMessage struct {
		Message string
		Event   map[string]interface{}
	}
	var data eventMessage
	err = bro.jobs[0].Message.Unmarshal(&data)
	assert.NoError(t, err)
	assert.Equal(t, data.Event["Domain"].(string), instanceName)
	assert.Equal(t, data.Event["Type"].(string), "CREATED")

	realtime.GetHub().Publish(&realtime.Event{
		Domain: instanceName,
		Doc: &testDoc{
			id:      "foo",
			doctype: "io.cozy.event-test",
		},
		Type: realtime.EventUpdate,
	})

	realtime.GetHub().Publish(&realtime.Event{
		Domain: instanceName,
		Doc: &testDoc{
			id:      "foo",
			doctype: "io.cozy.event-test.bad",
		},
		Type: realtime.EventCreate,
	})

	time.Sleep(10 * time.Millisecond)

	count, _ = bro.QueueLen("incr")
	assert.Equal(t, 1, count)
}

func TestMain(m *testing.M) {
	// prefix = "test:"
	config.UseTestFile()
	cfg := config.GetConfig()
	was := cfg.Jobs.URL
	cfg.Jobs.URL = redisURL

	testutils.NeedCouchdb()
	setup := testutils.NewSetup(m, "test_redis_scheduler")
	instanceName = setup.GetTestInstance().Domain

	setup.AddCleanup(func() error {
		cfg.Jobs.URL = was
		opts, _ := redis.ParseURL(redisURL)
		client := redis.NewClient(opts)
		return client.Del(scheduler.TriggersKey, scheduler.SchedKey).Err()
	})

	os.Exit(setup.Run())
}
<commit_msg>Increase wait in redis_trigger_event test<commit_after>package scheduler_test

import (
	"context"
	"errors"
	"os"
	"sync"
	"testing"
	"time"

	"github.com/cozy/cozy-stack/pkg/config"
	"github.com/cozy/cozy-stack/pkg/couchdb"
	"github.com/cozy/cozy-stack/pkg/jobs"
	"github.com/cozy/cozy-stack/pkg/realtime"
	"github.com/cozy/cozy-stack/pkg/scheduler"
	"github.com/cozy/cozy-stack/pkg/stack"
	"github.com/cozy/cozy-stack/tests/testutils"
	"github.com/go-redis/redis"
	"github.com/stretchr/testify/assert"
)

const redisURL = "redis://localhost:6379/15"

var instanceName string

type testDoc struct {
	id      string
	rev     string
	doctype string
}

func (t *testDoc) ID() string      { return t.id }
func (t *testDoc) Rev() string     { return t.rev }
func (t *testDoc) DocType() string { return t.doctype }

type mockBroker struct {
	jobs []*jobs.JobRequest
}

func (b *mockBroker) PushJob(request *jobs.JobRequest) (*jobs.JobInfos, error) {
	b.jobs = append(b.jobs, request)
	return nil, nil
}

func (b *mockBroker) QueueLen(workerType string) (int, error) {
	count := 0
	for _, job := range b.jobs {
		if job.WorkerType == workerType {
			count++
		}
	}
	return count, nil
}

func (b *mockBroker) GetJobInfos(domain, id string) (*jobs.JobInfos, error) {
	return nil, errors.New("Not implemented")
}

func TestRedisSchedulerWithTimeTriggers(t *testing.T) {
	var wAt sync.WaitGroup
	var wIn sync.WaitGroup
	bro := jobs.NewMemBroker(1, jobs.WorkersList{
		"worker": {
			Concurrency:  1,
			MaxExecCount: 1,
			Timeout:      1 * time.Millisecond,
			WorkerFunc: func(ctx context.Context, m *jobs.Message) error {
				var msg string
				if err := m.Unmarshal(&msg); err != nil {
					return err
				}
				switch msg {
				case "@at":
					wAt.Done()
				case "@in":
					wIn.Done()
				}
				return nil
			},
		},
	})

	msg1, _ := jobs.NewMessage("json", "@at")
	msg2, _ := jobs.NewMessage("json", "@in")

	wAt.Add(1) // 1 time in @at
	wIn.Add(1) // 1 time in @in

	at := &scheduler.TriggerInfos{
		Type:       "@at",
		Domain:     instanceName,
		Arguments:  time.Now().Add(2 * time.Second).Format(time.RFC3339),
		WorkerType: "worker",
		Message:    msg1,
	}
	in := &scheduler.TriggerInfos{
		Domain:     instanceName,
		Type:       "@in",
		Arguments:  "1s",
		WorkerType: "worker",
		Message:    msg2,
	}

	sch := stack.GetScheduler().(*scheduler.RedisScheduler)
	sch.Stop()
	sch.Start(bro)

	tat, err := scheduler.NewTrigger(at)
	assert.NoError(t, err)
	err = sch.Add(tat)
	assert.NoError(t, err)
	atID := tat.Infos().TID

	tin, err := scheduler.NewTrigger(in)
	assert.NoError(t, err)
	err = sch.Add(tin)
	assert.NoError(t, err)
	inID := tin.Infos().TID

	ts, err := sch.GetAll(instanceName)
	assert.NoError(t, err)
	assert.Len(t, ts, 3) // 1 @event for thumbnails + 1 @at + 1 @in

	for _, trigger := range ts {
		switch trigger.Infos().TID {
		case atID:
			assert.Equal(t, at, trigger.Infos())
		case inID:
			assert.Equal(t, in, trigger.Infos())
		default:
			// Just ignore the @event trigger for generating thumbnails
			infos := trigger.Infos()
			if infos.Type != "@event" || infos.WorkerType != "thumbnail" {
				t.Fatalf("unknown trigger ID %s", trigger.Infos().TID)
			}
		}
	}

	done := make(chan bool)
	go func() {
		wAt.Wait()
		done <- true
	}()

	go func() {
		wIn.Wait()
		done <- true
	}()

	for i := 0; i < 2; i++ {
		select {
		case <-done:
		case <-time.After(2 * time.Second):
			t.Fatalf("Timeout")
		}
	}

	time.Sleep(50 * time.Millisecond)

	_, err = sch.Get(instanceName, atID)
	assert.Error(t, err)
	assert.Equal(t, scheduler.ErrNotFoundTrigger, err)

	_, err = sch.Get(instanceName, inID)
	assert.Error(t, err)
	assert.Equal(t, scheduler.ErrNotFoundTrigger, err)
}

func TestRedisSchedulerWithCronTriggers(t *testing.T) {
	opts, _ := redis.ParseURL(redisURL)
	client := redis.NewClient(opts)
	err := client.Del(scheduler.TriggersKey, scheduler.SchedKey).Err()
	assert.NoError(t, err)

	bro := &mockBroker{}
	sch := stack.GetScheduler().(*scheduler.RedisScheduler)
	sch.Stop()
	sch.Start(bro)
	sch.Stop()
	defer sch.Start(bro)

	msg, _ := jobs.NewMessage("json", "@cron")

	infos := &scheduler.TriggerInfos{
		Type:       "@cron",
		Domain:     instanceName,
		Arguments:  "*/3 * * * * *",
		WorkerType: "incr",
		Message:    msg,
	}
	trigger, err := scheduler.NewTrigger(infos)
	assert.NoError(t, err)
	err = sch.Add(trigger)
	assert.NoError(t, err)

	now := time.Now().UTC().Unix()
	for i := int64(0); i < 9; i++ {
		err = sch.Poll(now + i + 4)
		assert.NoError(t, err)
	}
	count, _ := bro.QueueLen("incr")
	assert.Equal(t, 4, count)
}

func TestRedisPollFromSchedKey(t *testing.T) {
	opts, _ := redis.ParseURL(redisURL)
	client := redis.NewClient(opts)
	err := client.Del(scheduler.TriggersKey, scheduler.SchedKey).Err()
	assert.NoError(t, err)

	bro := &mockBroker{}
	sch := stack.GetScheduler().(*scheduler.RedisScheduler)
	sch.Stop()
	sch.Start(bro)
	sch.Stop()
	defer sch.Start(bro)

	now := time.Now()
	msg, _ := jobs.NewMessage("json", "@at")

	at := &scheduler.TriggerInfos{
		Type:       "@at",
		Domain:     instanceName,
		Arguments:  now.Format(time.RFC3339),
		WorkerType: "incr",
		Message:    msg,
	}
	db := couchdb.SimpleDatabasePrefix(instanceName)
	err = couchdb.CreateDoc(db, at)
	assert.NoError(t, err)

	ts := now.UTC().Unix()
	key := instanceName + "/" + at.TID
	err = client.ZAdd(scheduler.SchedKey, redis.Z{
		Score:  float64(ts + 1),
		Member: key,
	}).Err()
	assert.NoError(t, err)

	err = sch.Poll(ts + 2)
	assert.NoError(t, err)
	<-time.After(1 * time.Millisecond)
	count, _ := bro.QueueLen("incr")
	assert.Equal(t, 0, count)

	err = sch.Poll(ts + 13)
	assert.NoError(t, err)
	<-time.After(1 * time.Millisecond)
	count, _ = bro.QueueLen("incr")
	assert.Equal(t, 1, count)
}

func TestRedisTriggerEvent(t *testing.T) {
	opts, _ := redis.ParseURL(redisURL)
	client := redis.NewClient(opts)
	err := client.Del(scheduler.TriggersKey, scheduler.SchedKey).Err()
	assert.NoError(t, err)

	bro := &mockBroker{}
	sch := stack.GetScheduler().(*scheduler.RedisScheduler)
	sch.Stop()
	time.Sleep(1 * time.Second)
	sch.Start(bro)

	evTrigger := &scheduler.TriggerInfos{
		Type:       "@event",
		Domain:     instanceName,
		Arguments:  "io.cozy.event-test:CREATED",
		WorkerType: "incr",
	}
	tri, err := scheduler.NewTrigger(evTrigger)
	assert.NoError(t, err)
	sch.Add(tri)

	realtime.GetHub().Publish(&realtime.Event{
		Domain: instanceName,
		Doc: &testDoc{
			id:      "foo",
			doctype: "io.cozy.event-test",
		},
		Type: realtime.EventCreate,
	})

	time.Sleep(1 * time.Second)

	count, _ := bro.QueueLen("incr")
	if !assert.Equal(t, 1, count) {
		return
	}

	type eventMessage struct {
		Message string
		Event   map[string]interface{}
	}
	var data eventMessage
	err = bro.jobs[0].Message.Unmarshal(&data)
	assert.NoError(t, err)
	assert.Equal(t, data.Event["Domain"].(string), instanceName)
	assert.Equal(t, data.Event["Type"].(string), "CREATED")

	realtime.GetHub().Publish(&realtime.Event{
		Domain: instanceName,
		Doc: &testDoc{
			id:      "foo",
			doctype: "io.cozy.event-test",
		},
		Type: realtime.EventUpdate,
	})

	realtime.GetHub().Publish(&realtime.Event{
		Domain: instanceName,
		Doc: &testDoc{
			id:      "foo",
			doctype: "io.cozy.event-test.bad",
		},
		Type: realtime.EventCreate,
	})

	time.Sleep(10 * time.Millisecond)

	count, _ = bro.QueueLen("incr")
	assert.Equal(t, 1, count)
}

func TestMain(m *testing.M) {
	// prefix = "test:"
	config.UseTestFile()
	cfg := config.GetConfig()
	was := cfg.Jobs.URL
	cfg.Jobs.URL = redisURL

	testutils.NeedCouchdb()
	setup := testutils.NewSetup(m, "test_redis_scheduler")
	instanceName = setup.GetTestInstance().Domain

	setup.AddCleanup(func() error {
		cfg.Jobs.URL = was
		opts, _ := redis.ParseURL(redisURL)
		client := redis.NewClient(opts)
		return client.Del(scheduler.TriggersKey, scheduler.SchedKey).Err()
	})

	os.Exit(setup.Run())
}
<|endoftext|>


<commit_before>package cache

import (
	"context"
	"crypto/tls"
	"flag"
	"fmt"
	"net"
	"net/http"
	"strings"
	"sync"
	"time"

	"golang.org/x/net/http2"

	"github.com/weaveworks/common/instrument"

	"github.com/prometheus/client_golang/prometheus/promauto"

	"github.com/grafana/groupcache_exporter"
	"github.com/mailgun/groupcache/v2"
	"github.com/pkg/errors"
	"github.com/prometheus/client_golang/prometheus"

	"github.com/go-kit/log"
	"github.com/go-kit/log/level"
	"github.com/grafana/dskit/ring"

	"github.com/grafana/loki/pkg/logqlmodel/stats"
	lokiutil "github.com/grafana/loki/pkg/util"
)

var (
	ErrGroupcacheMiss = errors.New("cache miss")

	http2Transport = &http2.Transport{
		DialTLS: func(network, addr string, cfg *tls.Config) (net.Conn, error) {
			return net.Dial(network, addr)
		},
		AllowHTTP: true,
	}
)

type GroupCache struct {
	peerRing             ring.ReadRing
	cache                *groupcache.Group
	pool                 *groupcache.HTTPPool
	stopChan             chan struct{}
	updateInterval       time.Duration
	logger               log.Logger
	wg                   sync.WaitGroup
	reg                  prometheus.Registerer
	startWaitingForClose context.CancelFunc
	cacheBytes           int64
}

// RingCfg is a wrapper for the Groupcache ring configuration plus the replication factor.
type RingCfg struct {
	lokiutil.RingConfig `yaml:",inline"`
}

type GroupCacheConfig struct {
	Enabled    bool    `yaml:"enabled,omitempty"`
	Ring       RingCfg `yaml:"ring,omitempty"`
	CapacityMB int64   `yaml:"capacity_per_cache_mb,omitempty"`
	ListenPort int     `yaml:"listen_port,omitempty"`

	Cache Cache `yaml:"-"`
}

// Groupconfig represents config per Group.
type GroupConfig struct {
	CapacityMB int64 `yaml:"capacity_mb,omitempty"`
}

// RegisterFlagsWithPrefix adds the flags required to config this to the given FlagSet
func (cfg *GroupCacheConfig) RegisterFlagsWithPrefix(prefix, _ string, f *flag.FlagSet) {
	cfg.Ring.RegisterFlagsWithPrefix(prefix, "", f)

	f.BoolVar(&cfg.Enabled, prefix+".enabled", false, "Whether or not groupcache is enabled")
	f.IntVar(&cfg.ListenPort, prefix+".listen_port", 4100, "The port to use for groupcache communication")
	f.Int64Var(&cfg.CapacityMB, prefix+".capacity-per-cache-mb", 100, "Capacity of each groupcache group in MB (default: 100). "+
		"NOTE: there are 3 caches (result, chunk, and index query), so the maximum used memory will be *triple* the value specified here.")
}

type ringManager interface {
	Addr() string
	Ring() ring.ReadRing
}

func NewGroupCache(rm ringManager, config GroupCacheConfig, logger log.Logger, reg prometheus.Registerer) (*GroupCache, error) {
	addr := fmt.Sprintf("http://%s", rm.Addr())
	level.Info(logger).Log("msg", "groupcache local address set to", "addr", addr)

	pool := groupcache.NewHTTPPoolOpts(
		addr,
		&groupcache.HTTPPoolOptions{
			Transport: func(_ context.Context) http.RoundTripper {
				return http2Transport
			},
		},
	)

	startCtx, cancel := context.WithCancel(context.Background())
	cache := &GroupCache{
		peerRing:             rm.Ring(),
		pool:                 pool,
		logger:               logger,
		stopChan:             make(chan struct{}),
		updateInterval:       5 * time.Minute,
		wg:                   sync.WaitGroup{},
		startWaitingForClose: cancel,
		reg:                  reg,
		cacheBytes:           config.CapacityMB * 1e6, // MB => B
	}

	go cache.serveGroupcache(config.ListenPort)

	go func() {
		// Avoid starting the cache and peer discovery until
		// a cache is being used
		<-startCtx.Done()
		go cache.updatePeers()

		cache.wg.Wait()
		close(cache.stopChan)
	}()

	return cache, nil
}

func (c *GroupCache) serveGroupcache(listenPort int) {
	addr := fmt.Sprintf(":%d", listenPort)
	l, err := net.Listen("tcp", addr)
	if err != nil {
		level.Error(c.logger).Log("msg", "unable to serve groupcache", "err", err)
		return
	}
	level.Info(c.logger).Log("msg", "groupcache listening", "addr", addr)

	server := http2.Server{}
	for {
		select {
		case <-c.stopChan:
			return
		default:
			conn, err := l.Accept()
			if err != nil {
				level.Error(c.logger).Log("msg", "groupcache connection failed", "err", err)
				continue
			}

			go server.ServeConn(conn, &http2.ServeConnOpts{Handler: c.pool})
		}
	}
}

func (c *GroupCache) updatePeers() {
	c.update()

	t := time.NewTicker(c.updateInterval)
	for {
		select {
		case <-t.C:
			c.update()
		case <-c.stopChan:
			return
		}
	}
}

func (c *GroupCache) update() {
	urls, err := c.peerUrls()
	if err != nil {
		level.Warn(c.logger).Log("msg", "unable to get groupcache peer urls", "err", err)
		return
	}

	level.Info(c.logger).Log("msg", "got groupcache peers", "peers", strings.Join(urls, ","))
	c.pool.Set(urls...)
}

func (c *GroupCache) peerUrls() ([]string, error) {
	replicationSet, err := c.peerRing.GetAllHealthy(ring.WriteNoExtend)
	if err != nil {
		return nil, err
	}

	var addrs []string
	for _, i := range replicationSet.Instances {
		addrs = append(addrs, fmt.Sprintf("http://%s", i.Addr))
	}
	return addrs, nil
}

func (c *GroupCache) Stats() *groupcache.Stats {
	if c.cache == nil {
		return nil
	}

	return &c.cache.Stats
}

type group struct {
	cache     *groupcache.Group
	logger    log.Logger
	wg        *sync.WaitGroup
	cacheType stats.CacheType

	// cacheBytes represents maxSize (in bytes) this group can grow.
	cacheBytes int64 // TODO(kavi): expose it as _info metrics later?

	fetchDuration prometheus.Observer
	storeDuration prometheus.Observer
}

func (c *GroupCache) NewGroup(name string, cfg *GroupConfig, ct stats.CacheType) Cache {
	// Return a known error on miss to track which keys need to be inserted
	missGetter := groupcache.GetterFunc(func(_ context.Context, _ string, _ groupcache.Sink) error {
		return ErrGroupcacheMiss
	})

	c.wg.Add(1)
	c.startWaitingForClose()

	cap := c.cacheBytes
	if cfg.CapacityMB != 0 {
		cap = cfg.CapacityMB * 1e6 // MB into bytes
	}

	requestDuration := promauto.With(c.reg).NewHistogramVec(prometheus.HistogramOpts{
		Namespace:   "loki",
		Name:        "groupcache_request_duration_seconds",
		Help:        "Total time spent in seconds doing groupcache requests.",
		Buckets:     instrument.DefBuckets,
		ConstLabels: prometheus.Labels{"cache_type": string(ct)},
	}, []string{"operation"})

	g := &group{
		cache:         groupcache.NewGroup(name, cap, missGetter),
		cacheBytes:    cap,
		logger:        c.logger,
		wg:            &c.wg,
		cacheType:     ct,
		fetchDuration: requestDuration.WithLabelValues("fetch"),
		storeDuration: requestDuration.WithLabelValues("store"),
	}

	exp := groupcache_exporter.NewExporter(map[string]string{"cache_type": string(ct)}, g)
	prometheus.WrapRegistererWithPrefix("loki_groupcache_", c.reg).MustRegister(exp)

	return g
}

func (c *group) Fetch(ctx context.Context, keys []string) ([]string, [][]byte, []string, error) {
	var (
		start  = time.Now()
		values = make([][]byte, 0, len(keys))
		missed = make([]string, 0, len(keys))
		found  = make([]string, 0, len(keys))
		data   = groupcache.ByteView{}
		sink   = groupcache.ByteViewSink(&data)
	)

	for _, key := range keys {
		if err := c.cache.Get(ctx, key, sink); err != nil {
			if errors.Is(err, ErrGroupcacheMiss) {
				missed = append(missed, key)
				continue
			}

			level.Error(c.logger).Log("msg", "unable to fetch from groupcache", "err", err)
			return found, values, missed, err
		}

		found = append(found, key)
		values = append(values, data.ByteSlice())
	}

	c.fetchDuration.Observe(time.Since(start).Seconds())
	return found, values, missed, nil
}

func (c *group) Store(ctx context.Context, keys []string, values [][]byte) error {
	start := time.Now()

	var lastErr error
	for i, key := range keys {
		if err := c.cache.Set(ctx, key, values[i], time.Time{}, c.GetCacheType() != stats.ChunkCache); err != nil {
			level.Warn(c.logger).Log("msg", "failed to put to groupcache", "err", err)
			lastErr = err
		}
	}

	c.storeDuration.Observe(time.Since(start).Seconds())
	return lastErr
}

func (c *group) Stop() {
	c.wg.Done()
}

func (c *group) GetCacheType() stats.CacheType {
	return c.cacheType
}
<commit_msg>Add groupcache timeouts (#6808)<commit_after>package cache

import (
	"context"
	"crypto/tls"
	"flag"
	"fmt"
	"net"
	"net/http"
	"strings"
	"sync"
	"time"

	"golang.org/x/net/http2"

	"github.com/weaveworks/common/instrument"

	"github.com/prometheus/client_golang/prometheus/promauto"

	"github.com/grafana/groupcache_exporter"
	"github.com/mailgun/groupcache/v2"
	"github.com/pkg/errors"
	"github.com/prometheus/client_golang/prometheus"

	"github.com/go-kit/log"
	"github.com/go-kit/log/level"
	"github.com/grafana/dskit/ring"

	"github.com/grafana/loki/pkg/logqlmodel/stats"
	lokiutil "github.com/grafana/loki/pkg/util"
)

var (
	ErrGroupcacheMiss = errors.New("cache miss")

	http2Transport = &http2.Transport{
		DialTLS: func(network, addr string, cfg *tls.Config) (net.Conn, error) {
			return net.Dial(network, addr)
		},
		AllowHTTP: true,
	}
)

type GroupCache struct {
	peerRing             ring.ReadRing
	cache                *groupcache.Group
	pool                 *groupcache.HTTPPool
	stopChan             chan struct{}
	updateInterval       time.Duration
	logger               log.Logger
	wg                   sync.WaitGroup
	reg                  prometheus.Registerer
	startWaitingForClose context.CancelFunc
	cacheBytes           int64
}

// RingCfg is a wrapper for the Groupcache ring configuration plus the replication factor.
type RingCfg struct {
	lokiutil.RingConfig `yaml:",inline"`
}

type GroupCacheConfig struct {
	Enabled           bool          `yaml:"enabled,omitempty"`
	Ring              RingCfg       `yaml:"ring,omitempty"`
	CapacityMB        int64         `yaml:"capacity_per_cache_mb,omitempty"`
	ListenPort        int           `yaml:"listen_port,omitempty"`
	HeartbeatInterval time.Duration `yaml:"heartbeat_interval,omitempty"`
	HeartbeatTimeout  time.Duration `yaml:"heartbeat_timeout,omitempty"`
	WriteByteTimeout  time.Duration `yaml:"write_timeout,omitempty"`

	Cache Cache `yaml:"-"`
}

// RegisterFlagsWithPrefix adds the flags required to config this to the given FlagSet
func (cfg *GroupCacheConfig) RegisterFlagsWithPrefix(prefix, _ string, f *flag.FlagSet) {
	cfg.Ring.RegisterFlagsWithPrefix(prefix, "", f)

	f.BoolVar(&cfg.Enabled, prefix+".enabled", false, "Whether or not groupcache is enabled")
	f.IntVar(&cfg.ListenPort, prefix+".listen_port", 4100, "The port to use for groupcache communication")
	f.Int64Var(&cfg.CapacityMB, prefix+".capacity-per-cache-mb", 100, "Capacity of each groupcache group in MB (default: 100). "+
		"NOTE: there are 3 caches (result, chunk, and index query), so the maximum used memory will be *triple* the value specified here.")
	f.DurationVar(&cfg.HeartbeatInterval, "prefix.heartbeat-interval", time.Second, "If the connection is idle, the interval the cache will send heartbeats")
	f.DurationVar(&cfg.HeartbeatTimeout, "prefix.heartbeat-timeout", time.Second, "Timeout for heartbeat responses")
	f.DurationVar(&cfg.WriteByteTimeout, "prefix.write-timeout", time.Second, "Maximum time for the cache to try writing")
}

type ringManager interface {
	Addr() string
	Ring() ring.ReadRing
}

func NewGroupCache(rm ringManager, config GroupCacheConfig, logger log.Logger, reg prometheus.Registerer) (*GroupCache, error) {
	addr := fmt.Sprintf("http://%s", rm.Addr())
	level.Info(logger).Log("msg", "groupcache local address set to", "addr", addr)

	http2Transport.ReadIdleTimeout = config.HeartbeatInterval
	http2Transport.PingTimeout = config.HeartbeatTimeout
	http2Transport.WriteByteTimeout = config.WriteByteTimeout

	pool := groupcache.NewHTTPPoolOpts(
		addr,
		&groupcache.HTTPPoolOptions{
			Transport: func(_ context.Context) http.RoundTripper {
				return http2Transport
			},
		},
	)

	startCtx, cancel := context.WithCancel(context.Background())
	cache := &GroupCache{
		peerRing:             rm.Ring(),
		pool:                 pool,
		logger:               logger,
		stopChan:             make(chan struct{}),
		updateInterval:       5 * time.Minute,
		wg:                   sync.WaitGroup{},
		startWaitingForClose: cancel,
		reg:                  reg,
		cacheBytes:           config.CapacityMB * 1e6, // MB => B
	}

	go cache.serveGroupcache(config.ListenPort)

	go func() {
		// Avoid starting the cache and peer discovery until
		// a cache is being used
		<-startCtx.Done()
		go cache.updatePeers()

		cache.wg.Wait()
		close(cache.stopChan)
	}()

	return cache, nil
}

func (c *GroupCache) serveGroupcache(listenPort int) {
	addr := fmt.Sprintf(":%d", listenPort)
	l, err := net.Listen("tcp", addr)
	if err != nil {
		level.Error(c.logger).Log("msg", "unable to serve groupcache", "err", err)
		return
	}
	level.Info(c.logger).Log("msg", "groupcache listening", "addr", addr)

	server := http2.Server{}
	for {
		select {
		case <-c.stopChan:
			return
		default:
			conn, err := l.Accept()
			if err != nil {
				level.Error(c.logger).Log("msg", "groupcache connection failed", "err", err)
				continue
			}

			go server.ServeConn(conn, &http2.ServeConnOpts{Handler: c.pool})
		}
	}
}

func (c *GroupCache) updatePeers() {
	c.update()

	t := time.NewTicker(c.updateInterval)
	for {
		select {
		case <-t.C:
			c.update()
		case <-c.stopChan:
			return
		}
	}
}

func (c *GroupCache) update() {
	urls, err := c.peerUrls()
	if err != nil {
		level.Warn(c.logger).Log("msg", "unable to get groupcache peer urls", "err", err)
		return
	}

	level.Info(c.logger).Log("msg", "got groupcache peers", "peers", strings.Join(urls, ","))
	c.pool.Set(urls...)
}

func (c *GroupCache) peerUrls() ([]string, error) {
	replicationSet, err := c.peerRing.GetAllHealthy(ring.WriteNoExtend)
	if err != nil {
		return nil, err
	}

	var addrs []string
	for _, i := range replicationSet.Instances {
		addrs = append(addrs, fmt.Sprintf("http://%s", i.Addr))
	}
	return addrs, nil
}

func (c *GroupCache) Stats() *groupcache.Stats {
	if c.cache == nil {
		return nil
	}

	return &c.cache.Stats
}

// Groupconfig represents config per Group.
type GroupConfig struct {
	CapacityMB int64 `yaml:"capacity_mb,omitempty"`
}

type group struct {
	cache     *groupcache.Group
	logger    log.Logger
	wg        *sync.WaitGroup
	cacheType stats.CacheType

	// cacheBytes represents maxSize (in bytes) this group can grow.
	cacheBytes int64 // TODO(kavi): expose it as _info metrics later?

	fetchDuration prometheus.Observer
	storeDuration prometheus.Observer
}

func (c *GroupCache) NewGroup(name string, cfg *GroupConfig, ct stats.CacheType) Cache {
	// Return a known error on miss to track which keys need to be inserted
	missGetter := groupcache.GetterFunc(func(_ context.Context, _ string, _ groupcache.Sink) error {
		return ErrGroupcacheMiss
	})

	c.wg.Add(1)
	c.startWaitingForClose()

	cap := c.cacheBytes
	if cfg.CapacityMB != 0 {
		cap = cfg.CapacityMB * 1e6 // MB into bytes
	}

	requestDuration := promauto.With(c.reg).NewHistogramVec(prometheus.HistogramOpts{
		Namespace:   "loki",
		Name:        "groupcache_request_duration_seconds",
		Help:        "Total time spent in seconds doing groupcache requests.",
		Buckets:     instrument.DefBuckets,
		ConstLabels: prometheus.Labels{"cache_type": string(ct)},
	}, []string{"operation"})

	g := &group{
		cache:         groupcache.NewGroup(name, cap, missGetter),
		cacheBytes:    cap,
		logger:        c.logger,
		wg:            &c.wg,
		cacheType:     ct,
		fetchDuration: requestDuration.WithLabelValues("fetch"),
		storeDuration: requestDuration.WithLabelValues("store"),
	}

	exp := groupcache_exporter.NewExporter(map[string]string{"cache_type": string(ct)}, g)
	prometheus.WrapRegistererWithPrefix("loki_groupcache_", c.reg).MustRegister(exp)

	return g
}

func (c *group) Fetch(ctx context.Context, keys []string) ([]string, [][]byte, []string, error) {
	var (
		start  = time.Now()
		values = make([][]byte, 0, len(keys))
		missed = make([]string, 0, len(keys))
		found  = make([]string, 0, len(keys))
		data   = groupcache.ByteView{}
		sink   = groupcache.ByteViewSink(&data)
	)

	for _, key := range keys {
		if err := c.cache.Get(ctx, key, sink); err != nil {
			if errors.Is(err, ErrGroupcacheMiss) {
				missed = append(missed, key)
				continue
			}

			level.Error(c.logger).Log("msg", "unable to fetch from groupcache", "err", err)
			return found, values, missed, err
		}

		found = append(found, key)
		values = append(values, data.ByteSlice())
	}

	c.fetchDuration.Observe(time.Since(start).Seconds())
	return found, values, missed, nil
}

func (c *group) Store(ctx context.Context, keys []string, values [][]byte) error {
	start := time.Now()

	var lastErr error
	for i, key := range keys {
		if err := c.cache.Set(ctx, key, values[i], time.Time{}, c.GetCacheType() != stats.ChunkCache); err != nil {
			level.Warn(c.logger).Log("msg", "failed to put to groupcache", "err", err)
			lastErr = err
		}
	}

	c.storeDuration.Observe(time.Since(start).Seconds())
	return lastErr
}

func (c *group) Stop() {
	c.wg.Done()
}

func (c *group) GetCacheType() stats.CacheType {
	return c.cacheType
}
<|endoftext|>


<commit_before>// stack2struct parses raw golang stack traces ([]byte) to a slice of well formated structs.

// As this package will need to evolve with the development of go's stack trace
// format, this is the stack format the package currently works with:
//
//   1: goroutine x [running]:                 <-- ignore this line
//   2: path/to/package.functionName()
//   3: path/to/responsible/file:lineNumber +0xdeadbeef
//   ... repeat 2 + 3 for each stack trace element
//
// To work with this, you need a type satisfying the interface
//
//	 type stackTrace interface {
// 	   AddEntry(lineNumber int, packageName string, fileName string, methodName string)
//   }
//
// and from there you can do whatever you like with the accumulated data.
package stack2struct

import (
	"fmt"
	"strconv"
	"strings"
)

// stackTrace is the interface a target stack has to satisfy.
type stackTrace interface {
	AddEntry(lineNumber int, packageName string, fileName string, methodName string)
}

// Parse loads the stack trace (given as trace) into the given stack.
// See Current() on how to obtain a stack trace.
func Parse(trace []byte, stack stackTrace) {
	lines := strings.Split(string(trace), "\n")

	var lineNumber int
	var fileName, packageName, methodName string

	for index, line := range lines[1:] {
		if len(line) == 0 {
			continue
		}
		if index%2 == 0 {
			packageName, methodName = extractPackageName(line)
		} else {
			lineNumber, fileName = extractLineNumberAndFile(line)
			stack.AddEntry(lineNumber, packageName, fileName, methodName)
		}
	}
}

// extractPageName receives a trace line and extracts packageName and
// methodName.
func extractPackageName(line string) (packageName, methodName string) {
	packagePath, packageNameAndFunction := splitAtLastSlash(line)
	parts := strings.Split(packageNameAndFunction, ".")
	packageName = fmt.Sprintf("%s/%s", packagePath, parts[0])
	methodName = parts[1]
	return
}

// extractLineNumberAndFile receives a trace line and extracts lineNumber and
// fileName.
func extractLineNumberAndFile(line string) (lineNumber int, fileName string) {
	_, fileAndLine := splitAtLastSlash(line)
	fileAndLine = removeSpaceAndSuffix(fileAndLine)
	parts := strings.Split(fileAndLine, ":")

	numberAsString := parts[1]
	number, _ := strconv.ParseUint(numberAsString, 10, 32)
	lineNumber = int(number)

	fileName = parts[0]
	return lineNumber, fileName
}

// splitAtLastSlash splits a string at the last found slash and returns the
// respective strings left and right of the slash.
func splitAtLastSlash(line string) (left, right string) {
	parts := strings.Split(line, "/")
	right = parts[len(parts)-1]
	left = strings.Join(parts[:len(parts)-1], "/")
	return
}

// removeSpaceAndSuffix splits the given string at ' ' and cuts off the part
// found after the last space.
func removeSpaceAndSuffix(line string) string {
	parts := strings.Split(line, " ")
	return strings.Join(parts[:len(parts)-1], " ")
}
<commit_msg>add Current()<commit_after>// stack2struct parses raw golang stack traces ([]byte) to a slice of well formated structs.

// As this package will need to evolve with the development of go's stack trace
// format, this is the stack format the package currently works with:
//
//   1: goroutine x [running]:                 <-- ignore this line
//   2: path/to/package.functionName()
//   3: path/to/responsible/file:lineNumber +0xdeadbeef
//   ... repeat 2 + 3 for each stack trace element
//
// To work with this, you need a type satisfying the interface
//
//	 type stackTrace interface {
// 	   AddEntry(lineNumber int, packageName string, fileName string, methodName string)
//   }
//
// and from there you can do whatever you like with the accumulated data.
package stack2struct

import (
	"fmt"
	"runtime"
	"strconv"
	"strings"
)

// stackTrace is the interface a target stack has to satisfy.
type stackTrace interface {
	AddEntry(lineNumber int, packageName string, fileName string, methodName string)
}

// Current loads the current stacktrace into a given stack
func Current(stack stackTrace) {
	rawStack := make([]byte, 1<<16)
	rawStack = rawStack[:runtime.Stack(rawStack, false)]
	Parse(rawStack, stack)
}

// Parse loads the stack trace (given as trace) into the given stack.
// See Current() on how to obtain a stack trace.
func Parse(trace []byte, stack stackTrace) {
	lines := strings.Split(string(trace), "\n")

	var lineNumber int
	var fileName, packageName, methodName string

	for index, line := range lines[1:] {
		if len(line) == 0 {
			continue
		}
		if index%2 == 0 {
			packageName, methodName = extractPackageName(line)
		} else {
			lineNumber, fileName = extractLineNumberAndFile(line)
			stack.AddEntry(lineNumber, packageName, fileName, methodName)
		}
	}
}

// extractPageName receives a trace line and extracts packageName and
// methodName.
func extractPackageName(line string) (packageName, methodName string) {
	packagePath, packageNameAndFunction := splitAtLastSlash(line)
	parts := strings.Split(packageNameAndFunction, ".")
	packageName = fmt.Sprintf("%s/%s", packagePath, parts[0])
	methodName = parts[1]
	return
}

// extractLineNumberAndFile receives a trace line and extracts lineNumber and
// fileName.
func extractLineNumberAndFile(line string) (lineNumber int, fileName string) {
	_, fileAndLine := splitAtLastSlash(line)
	fileAndLine = removeSpaceAndSuffix(fileAndLine)
	parts := strings.Split(fileAndLine, ":")

	numberAsString := parts[1]
	number, _ := strconv.ParseUint(numberAsString, 10, 32)
	lineNumber = int(number)

	fileName = parts[0]
	return lineNumber, fileName
}

// splitAtLastSlash splits a string at the last found slash and returns the
// respective strings left and right of the slash.
func splitAtLastSlash(line string) (left, right string) {
	parts := strings.Split(line, "/")
	right = parts[len(parts)-1]
	left = strings.Join(parts[:len(parts)-1], "/")
	return
}

// removeSpaceAndSuffix splits the given string at ' ' and cuts off the part
// found after the last space.
func removeSpaceAndSuffix(line string) string {
	parts := strings.Split(line, " ")
	return strings.Join(parts[:len(parts)-1], " ")
}
<|endoftext|>


<commit_before>// Copyright (c) 2015 Uber Technologies, Inc.

// Permission is hereby granted, free of charge, to any person obtaining a copy
// of this software and associated documentation files (the "Software"), to deal
// in the Software without restriction, including without limitation the rights
// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
// copies of the Software, and to permit persons to whom the Software is
// furnished to do so, subject to the following conditions:
//
// The above copyright notice and this permission notice shall be included in
// all copies or substantial portions of the Software.
//
// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
// THE SOFTWARE.

package main

import (
	"bufio"
	"fmt"
	"io"
	"io/ioutil"
	"os"
	"os/exec"
	"path"
	"path/filepath"
	"strings"
	"sync"
	"testing"

	"github.com/stretchr/testify/require"
	"github.com/uber/tchannel-go/testutils"
)

// These tests ensure that the code generator generates valid code that can be built
// in combination with Thrift's autogenerated code.

var (
	_testGoPath     string
	_testGoPathOnce sync.Once
)

func TestMain(m *testing.M) {
	exitCode := m.Run()

	// If we created a fake GOPATH, we should clean it up.
	if _testGoPath != "" {
		os.RemoveAll(_testGoPath)
	}

	os.Exit(exitCode)
}

func getTChannelDir(goPath string) string {
	return filepath.Join(goPath, "src", "github.com", "uber", "tchannel-go")
}

func createGoPath(t *testing.T) {
	goPath, err := ioutil.TempDir("", "thrift-gen")
	require.NoError(t, err, "TempDir failed")

	// Create $GOPATH/src/github.com/uber/tchannel-go and symlink everything.
	// And then create a dummy directory for all the test output.
	tchannelDir := getTChannelDir(goPath)
	require.NoError(t, os.MkdirAll(tchannelDir, 0755), "MkDirAll failed")

	// Symlink everything in the real GOPATH.
	realTChannelDir := getTChannelDir(os.Getenv("GOPATH"))
	realDirContents, err := ioutil.ReadDir(realTChannelDir)
	require.NoError(t, err, "Failed to read real tchannel-go dir")

	for _, f := range realDirContents {
		realPath := filepath.Join(realTChannelDir, f.Name())
		err := os.Symlink(realPath, filepath.Join(tchannelDir, filepath.Base(f.Name())))
		require.NoError(t, err, "Failed to symlink %v", f.Name())
	}

	_testGoPath = goPath
	os.Setenv("GOPATH", goPath)
}

func getOutputDir(t *testing.T) string {
	_testGoPathOnce.Do(func() { createGoPath(t) })

	// Create a random directory inside of the GOPATH in tmp
	randDir := filepath.Join(getTChannelDir(_testGoPath), testutils.RandString(10))
	// In case it's not empty.
	os.RemoveAll(randDir)

	return randDir
}

func TestAllThrift(t *testing.T) {
	files, err := ioutil.ReadDir("test_files")
	require.NoError(t, err, "Cannot read test_files directory: %v", err)

	for _, f := range files {
		fname := f.Name()
		if f.IsDir() || filepath.Ext(fname) != ".thrift" {
			continue
		}

		if err := runBuildTest(t, filepath.Join("test_files", fname)); err != nil {
			t.Errorf("Thrift file %v failed: %v", fname, err)
		}
	}
}

func TestIncludeThrift(t *testing.T) {
	dirs, err := ioutil.ReadDir("test_files/include_test")
	require.NoError(t, err, "Cannot read test_files/include_test directory: %v", err)

	for _, d := range dirs {
		dname := d.Name()
		if !d.IsDir() {
			continue
		}

		thriftFile := filepath.Join(dname, path.Base(dname)+".thrift")
		if err := runBuildTest(t, filepath.Join("test_files/include_test/", thriftFile)); err != nil {
			t.Errorf("Thrift test %v failed: %v", dname, err)
		}
	}
}

func TestMultipleFiles(t *testing.T) {
	if err := runBuildTest(t, filepath.Join("test_files", "multi_test", "file1.thrift")); err != nil {
		t.Errorf("Multiple file test failed: %v", err)
	}
}

func TestExternalTemplate(t *testing.T) {
	template1 := `package {{ .Package }}

{{ range .AST.Services }}
// Service {{ .Name }} has {{ len .Methods }} methods.
{{ range .Methods }}
// func {{ .Name | goPublicName }} ({{ range .Arguments }}{{ .Type | goType }}, {{ end }}) ({{ if .ReturnType }}{{ .ReturnType | goType }}{{ end }}){{ end }}
{{ end }}
	`
	templateFile := writeTempFile(t, template1)
	defer os.Remove(templateFile)

	expected := `package service_extend

// Service S1 has 1 methods.

// func M1 ([]byte, ) ([]byte)

// Service S2 has 1 methods.

// func M2 (*S, int32, ) (*S)

// Service S3 has 1 methods.

// func M3 () ()
`

	opts := processOptions{
		InputFile:     "test_files/service_extend.thrift",
		TemplateFiles: []string{templateFile},
	}
	checks := func(dir string) error {
		dir = filepath.Join(dir, "service_extend")
		if err := checkDirectoryFiles(dir, 6); err != nil {
			return err
		}

		// Verify the contents of the extra file.
		outFile := filepath.Join(dir, packageName(templateFile)+"-service_extend.go")
		return verifyFileContents(outFile, expected)
	}
	if err := runTest(t, opts, checks); err != nil {
		t.Errorf("Failed to run test: %v", err)
	}
}

func writeTempFile(t *testing.T, contents string) string {
	tempFile, err := ioutil.TempFile("", "temp")
	require.NoError(t, err, "Failed to create temp file")
	tempFile.Close()
	require.NoError(t, ioutil.WriteFile(tempFile.Name(), []byte(contents), 0666),
		"Write temp file failed")
	return tempFile.Name()
}

func verifyFileContents(filename, expected string) error {
	bytes, err := ioutil.ReadFile(filename)
	if err != nil {
		return err
	}

	bytesStr := string(bytes)
	if bytesStr != expected {
		return fmt.Errorf("file contents mismatch. got:\n%vexpected:\n%v", bytesStr, expected)
	}

	return nil
}

func copyFile(src, dst string) error {
	f, err := os.Open(src)
	if err != nil {
		return err
	}
	defer f.Close()

	writeF, err := os.OpenFile(dst, os.O_CREATE|os.O_TRUNC|os.O_WRONLY, 0666)
	if err != nil {
		return err
	}
	defer writeF.Close()

	_, err = io.Copy(writeF, f)
	return err
}

// setupDirectory creates a temporary directory.
func setupDirectory(thriftFile string) (string, error) {
	tempDir, err := ioutil.TempDir("", "thrift-gen")
	if err != nil {
		return "", err
	}

	return tempDir, nil
}

func createAdditionalTestFile(thriftFile, tempDir string) error {
	f, err := os.Open(thriftFile)
	if err != nil {
		return err
	}

	var writer io.Writer
	rdr := bufio.NewReader(f)
	for {
		line, err := rdr.ReadString('\n')
		if err != nil {
			if err == io.EOF {
				return nil
			}
		}

		if strings.HasPrefix(line, "//Go code:") {
			fileName := strings.TrimSpace(strings.TrimPrefix(line, "//Go code:"))
			outFile := filepath.Join(tempDir, fileName)
			f, err := os.OpenFile(outFile, os.O_CREATE|os.O_TRUNC|os.O_WRONLY, 0666)
			if err != nil {
				return err
			}
			defer f.Close()
			writer = f
		} else if writer != nil {
			if strings.HasPrefix(line, "//") {
				writer.Write([]byte(strings.TrimPrefix(line, "//")))
			} else {
				return nil
			}
		}
	}
}

func checkDirectoryFiles(dir string, n int) error {
	dirContents, err := ioutil.ReadDir(dir)
	if err != nil {
		return err
	}

	if len(dirContents) < n {
		return fmt.Errorf("expected to generate at least %v files, but found: %v", n, len(dirContents))
	}

	return nil
}

func runBuildTest(t *testing.T, thriftFile string) error {
	extraChecks := func(dir string) error {
		return checkDirectoryFiles(filepath.Join(dir, packageName(thriftFile)), 4)
	}

	opts := processOptions{InputFile: thriftFile}
	return runTest(t, opts, extraChecks)
}

func runTest(t *testing.T, opts processOptions, extraChecks func(string) error) error {
	tempDir := getOutputDir(t)

	// Generate code from the Thrift file.
	*packagePrefix = strings.TrimPrefix(tempDir, _testGoPath+"/src/") + "/"
	opts.GenerateThrift = true
	opts.OutputDir = tempDir
	if err := processFile(opts); err != nil {
		return fmt.Errorf("processFile(%s) in %q failed: %v", opts.InputFile, tempDir, err)
	}

	// Create any extra Go files as specified in the Thrift file.
	if err := createAdditionalTestFile(opts.InputFile, tempDir); err != nil {
		return fmt.Errorf("failed creating additional test files for %s in %q: %v", opts.InputFile, tempDir, err)
	}

	// Run go build to ensure that the generated code builds.
	cmd := exec.Command("go", "build", "-i", "./...")
	cmd.Dir = tempDir
	// NOTE: we check output, since go build ./... returns 0 status code on failure:
	// https://github.com/golang/go/issues/11407
	if output, err := cmd.CombinedOutput(); err != nil || len(output) > 0 {
		return fmt.Errorf("build in %q failed.\nError: %v Output:\n%v", tempDir, err, string(output))
	}

	// Run any extra checks the user may want.
	if err := extraChecks(tempDir); err != nil {
		return err
	}

	// Only delete the temp directory on success.
	os.RemoveAll(tempDir)
	return nil
}
<commit_msg>Follow up test cleanup from #558 (#560)<commit_after>// Copyright (c) 2015 Uber Technologies, Inc.

// Permission is hereby granted, free of charge, to any person obtaining a copy
// of this software and associated documentation files (the "Software"), to deal
// in the Software without restriction, including without limitation the rights
// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
// copies of the Software, and to permit persons to whom the Software is
// furnished to do so, subject to the following conditions:
//
// The above copyright notice and this permission notice shall be included in
// all copies or substantial portions of the Software.
//
// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
// THE SOFTWARE.

package main

import (
	"bufio"
	"fmt"
	"io"
	"io/ioutil"
	"os"
	"os/exec"
	"path"
	"path/filepath"
	"strings"
	"sync"
	"testing"

	"github.com/stretchr/testify/require"
	"github.com/uber/tchannel-go/testutils"
)

// These tests ensure that the code generator generates valid code that can be built
// in combination with Thrift's autogenerated code.

const _tchannelPackage = "github.com/uber/tchannel-go"

var (
	_testGoPath     string
	_testGoPathOnce sync.Once
)

func TestMain(m *testing.M) {
	exitCode := m.Run()

	// If we created a fake GOPATH, we should clean it up.
	if _testGoPath != "" {
		os.RemoveAll(_testGoPath)
	}

	os.Exit(exitCode)
}

func getTChannelDir(goPath string) string {
	return filepath.Join(goPath, "src", _tchannelPackage)
}

func getCurrentTChannelPath(t *testing.T) string {
	wd, err := os.Getwd()
	require.NoError(t, err, "Failed to get working directory")

	// Walk up "wd" till we find "tchannel-go".
	for filepath.Base(wd) != filepath.Base(_tchannelPackage) {
		wd = filepath.Dir(wd)
		if wd == "" {
			t.Fatalf("Failed to find tchannel-go in parents of current directory")
		}
	}

	return wd
}

func createGoPath(t *testing.T) {
	goPath, err := ioutil.TempDir("", "thrift-gen")
	require.NoError(t, err, "TempDir failed")

	// Create $GOPATH/src/github.com/uber/tchannel-go and symlink everything.
	// And then create a dummy directory for all the test output.
	tchannelDir := getTChannelDir(goPath)
	require.NoError(t, os.MkdirAll(tchannelDir, 0755), "MkDirAll failed")

	// Symlink the contents of tchannel-go into the temp directory.
	realTChannelDir := getCurrentTChannelPath(t)
	realDirContents, err := ioutil.ReadDir(realTChannelDir)
	require.NoError(t, err, "Failed to read real tchannel-go dir")

	for _, f := range realDirContents {
		realPath := filepath.Join(realTChannelDir, f.Name())
		err := os.Symlink(realPath, filepath.Join(tchannelDir, filepath.Base(f.Name())))
		require.NoError(t, err, "Failed to symlink %v", f.Name())
	}

	_testGoPath = goPath

	// None of the other tests in this package should use GOPATH, so we don't
	// restore this.
	os.Setenv("GOPATH", goPath)
}

func getOutputDir(t *testing.T) (dir, pkg string) {
	_testGoPathOnce.Do(func() { createGoPath(t) })

	// Create a random directory inside of the GOPATH in tmp
	randStr := testutils.RandString(10)
	randDir := filepath.Join(getTChannelDir(_testGoPath), randStr)
	// In case it's not empty.
	os.RemoveAll(randDir)

	return randDir, filepath.Join(_tchannelPackage, randStr)
}

func TestAllThrift(t *testing.T) {
	files, err := ioutil.ReadDir("test_files")
	require.NoError(t, err, "Cannot read test_files directory: %v", err)

	for _, f := range files {
		fname := f.Name()
		if f.IsDir() || filepath.Ext(fname) != ".thrift" {
			continue
		}

		if err := runBuildTest(t, filepath.Join("test_files", fname)); err != nil {
			t.Errorf("Thrift file %v failed: %v", fname, err)
		}
	}
}

func TestIncludeThrift(t *testing.T) {
	dirs, err := ioutil.ReadDir("test_files/include_test")
	require.NoError(t, err, "Cannot read test_files/include_test directory: %v", err)

	for _, d := range dirs {
		dname := d.Name()
		if !d.IsDir() {
			continue
		}

		thriftFile := filepath.Join(dname, path.Base(dname)+".thrift")
		if err := runBuildTest(t, filepath.Join("test_files/include_test/", thriftFile)); err != nil {
			t.Errorf("Thrift test %v failed: %v", dname, err)
		}
	}
}

func TestMultipleFiles(t *testing.T) {
	if err := runBuildTest(t, filepath.Join("test_files", "multi_test", "file1.thrift")); err != nil {
		t.Errorf("Multiple file test failed: %v", err)
	}
}

func TestExternalTemplate(t *testing.T) {
	template1 := `package {{ .Package }}

{{ range .AST.Services }}
// Service {{ .Name }} has {{ len .Methods }} methods.
{{ range .Methods }}
// func {{ .Name | goPublicName }} ({{ range .Arguments }}{{ .Type | goType }}, {{ end }}) ({{ if .ReturnType }}{{ .ReturnType | goType }}{{ end }}){{ end }}
{{ end }}
	`
	templateFile := writeTempFile(t, template1)
	defer os.Remove(templateFile)

	expected := `package service_extend

// Service S1 has 1 methods.

// func M1 ([]byte, ) ([]byte)

// Service S2 has 1 methods.

// func M2 (*S, int32, ) (*S)

// Service S3 has 1 methods.

// func M3 () ()
`

	opts := processOptions{
		InputFile:     "test_files/service_extend.thrift",
		TemplateFiles: []string{templateFile},
	}
	checks := func(dir string) error {
		dir = filepath.Join(dir, "service_extend")
		if err := checkDirectoryFiles(dir, 6); err != nil {
			return err
		}

		// Verify the contents of the extra file.
		outFile := filepath.Join(dir, packageName(templateFile)+"-service_extend.go")
		return verifyFileContents(outFile, expected)
	}
	if err := runTest(t, opts, checks); err != nil {
		t.Errorf("Failed to run test: %v", err)
	}
}

func writeTempFile(t *testing.T, contents string) string {
	tempFile, err := ioutil.TempFile("", "temp")
	require.NoError(t, err, "Failed to create temp file")
	tempFile.Close()
	require.NoError(t, ioutil.WriteFile(tempFile.Name(), []byte(contents), 0666),
		"Write temp file failed")
	return tempFile.Name()
}

func verifyFileContents(filename, expected string) error {
	bytes, err := ioutil.ReadFile(filename)
	if err != nil {
		return err
	}

	bytesStr := string(bytes)
	if bytesStr != expected {
		return fmt.Errorf("file contents mismatch. got:\n%vexpected:\n%v", bytesStr, expected)
	}

	return nil
}

func copyFile(src, dst string) error {
	f, err := os.Open(src)
	if err != nil {
		return err
	}
	defer f.Close()

	writeF, err := os.OpenFile(dst, os.O_CREATE|os.O_TRUNC|os.O_WRONLY, 0666)
	if err != nil {
		return err
	}
	defer writeF.Close()

	_, err = io.Copy(writeF, f)
	return err
}

// setupDirectory creates a temporary directory.
func setupDirectory(thriftFile string) (string, error) {
	tempDir, err := ioutil.TempDir("", "thrift-gen")
	if err != nil {
		return "", err
	}

	return tempDir, nil
}

func createAdditionalTestFile(thriftFile, tempDir string) error {
	f, err := os.Open(thriftFile)
	if err != nil {
		return err
	}

	var writer io.Writer
	rdr := bufio.NewReader(f)
	for {
		line, err := rdr.ReadString('\n')
		if err != nil {
			if err == io.EOF {
				return nil
			}
		}

		if strings.HasPrefix(line, "//Go code:") {
			fileName := strings.TrimSpace(strings.TrimPrefix(line, "//Go code:"))
			outFile := filepath.Join(tempDir, fileName)
			f, err := os.OpenFile(outFile, os.O_CREATE|os.O_TRUNC|os.O_WRONLY, 0666)
			if err != nil {
				return err
			}
			defer f.Close()
			writer = f
		} else if writer != nil {
			if strings.HasPrefix(line, "//") {
				writer.Write([]byte(strings.TrimPrefix(line, "//")))
			} else {
				return nil
			}
		}
	}
}

func checkDirectoryFiles(dir string, n int) error {
	dirContents, err := ioutil.ReadDir(dir)
	if err != nil {
		return err
	}

	if len(dirContents) < n {
		return fmt.Errorf("expected to generate at least %v files, but found: %v", n, len(dirContents))
	}

	return nil
}

func runBuildTest(t *testing.T, thriftFile string) error {
	extraChecks := func(dir string) error {
		return checkDirectoryFiles(filepath.Join(dir, packageName(thriftFile)), 4)
	}

	opts := processOptions{InputFile: thriftFile}
	return runTest(t, opts, extraChecks)
}

func runTest(t *testing.T, opts processOptions, extraChecks func(string) error) error {
	tempDir, outputPkg := getOutputDir(t)

	// Generate code from the Thrift file.
	*packagePrefix = outputPkg + "/"
	opts.GenerateThrift = true
	opts.OutputDir = tempDir
	if err := processFile(opts); err != nil {
		return fmt.Errorf("processFile(%s) in %q failed: %v", opts.InputFile, tempDir, err)
	}

	// Create any extra Go files as specified in the Thrift file.
	if err := createAdditionalTestFile(opts.InputFile, tempDir); err != nil {
		return fmt.Errorf("failed creating additional test files for %s in %q: %v", opts.InputFile, tempDir, err)
	}

	// Run go build to ensure that the generated code builds.
	cmd := exec.Command("go", "build", "-i", "./...")
	cmd.Dir = tempDir
	// NOTE: we check output, since go build ./... returns 0 status code on failure:
	// https://github.com/golang/go/issues/11407
	if output, err := cmd.CombinedOutput(); err != nil || len(output) > 0 {
		return fmt.Errorf("build in %q failed.\nError: %v Output:\n%v", tempDir, err, string(output))
	}

	// Run any extra checks the user may want.
	if err := extraChecks(tempDir); err != nil {
		return err
	}

	// Only delete the temp directory on success.
	os.RemoveAll(tempDir)
	return nil
}
<|endoftext|>


<commit_before>package main

import (
	"bufio"
	"fmt"
	"os"
	"strconv"
	"strings"
)

func main() {

	modeline := ""
	blocks := map[string]int{}
	scanner := bufio.NewScanner(os.Stdin)
	for scanner.Scan() {
		line := scanner.Text()
		if !strings.HasPrefix(line, "mode:") {
			lastSpace := strings.LastIndex(line, " ")
			prefix := line[0:lastSpace]
			suffix := line[lastSpace+1:]
			count, err := strconv.Atoi(suffix)
			if err != nil {
				fmt.Printf("error parsing count: %v", err)
				continue
			}
			existingCount, exists := blocks[prefix]
			if exists {
				blocks[prefix] = existingCount + count
			} else {
				blocks[prefix] = count
			}
		} else if modeline == "" {
			modeline = line
		}
	}
	if err := scanner.Err(); err != nil {
		fmt.Fprintln(os.Stderr, "reading standard input:", err)
	}

	fmt.Println(modeline)
	for k, v := range blocks {
		fmt.Printf("%s %d\n", k, v)
	}
}
<commit_msg>add build tag protecting merge-coverprofile<commit_after>// Copyright © 2016 Couchbase, Inc.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

// +build ignore

package main

import (
	"bufio"
	"fmt"
	"os"
	"strconv"
	"strings"
)

func main() {

	modeline := ""
	blocks := map[string]int{}
	scanner := bufio.NewScanner(os.Stdin)
	for scanner.Scan() {
		line := scanner.Text()
		if !strings.HasPrefix(line, "mode:") {
			lastSpace := strings.LastIndex(line, " ")
			prefix := line[0:lastSpace]
			suffix := line[lastSpace+1:]
			count, err := strconv.Atoi(suffix)
			if err != nil {
				fmt.Printf("error parsing count: %v", err)
				continue
			}
			existingCount, exists := blocks[prefix]
			if exists {
				blocks[prefix] = existingCount + count
			} else {
				blocks[prefix] = count
			}
		} else if modeline == "" {
			modeline = line
		}
	}
	if err := scanner.Err(); err != nil {
		fmt.Fprintln(os.Stderr, "reading standard input:", err)
	}

	fmt.Println(modeline)
	for k, v := range blocks {
		fmt.Printf("%s %d\n", k, v)
	}
}
<|endoftext|>


<commit_before>package utils

import (
	"io"
	"net/http"
	"net/textproto"
	"strconv"
	"strings"
)

// ServeContent replies to the request using the content in the provided
// reader. The Content-Length and Content-Type headers are added with the
// provided values.
func ServeContent(w http.ResponseWriter, r *http.Request, contentType string, size int64, content io.Reader) {
	h := w.Header()
	if size > 0 {
		h.Set("Content-Length", strconv.FormatInt(size, 10))
	}
	if contentType != "" {
		h.Set("Content-Type", contentType)
	}
	w.WriteHeader(http.StatusOK)
	if r.Method != "HEAD" {
		io.Copy(w, content)
	}
}

// CheckPreconditions evaluates request preconditions based only on the Etag
// values.
func CheckPreconditions(w http.ResponseWriter, r *http.Request, etag string) (done bool) {
	if etag != "" && (checkIfMatch(w, r, etag) || checkIfNoneMatch(w, r, etag)) {
		writeNotModified(w)
		return true
	}
	return false
}

func checkIfMatch(w http.ResponseWriter, r *http.Request, definedETag string) (match bool) {
	im := r.Header.Get("If-Match")
	if im == "" {
		return false
	}
	for {
		im = textproto.TrimString(im)
		if len(im) == 0 {
			break
		}
		if im[0] == ',' {
			im = im[1:]
			continue
		}
		if im[0] == '*' {
			return true
		}
		etag, remain := scanETag(im)
		if etag == "" {
			break
		}
		if etagStrongMatch(etag, definedETag) {
			return true
		}
		im = remain
	}
	return false
}

func checkIfNoneMatch(w http.ResponseWriter, r *http.Request, definedETag string) (match bool) {
	inm := r.Header.Get("If-None-Match")
	if inm == "" {
		return false
	}
	buf := inm
	for {
		buf = textproto.TrimString(buf)
		if len(buf) == 0 {
			break
		}
		if buf[0] == ',' {
			buf = buf[1:]
		}
		if buf[0] == '*' {
			return false
		}
		etag, remain := scanETag(buf)
		if etag == "" {
			break
		}
		if etagWeakMatch(etag, definedETag) {
			return false
		}
		buf = remain
	}
	return true
}

// etagWeakMatch reports whether a and b match using weak ETag comparison.
// Assumes a and b are valid ETags.
func etagWeakMatch(a, b string) bool {
	return strings.TrimPrefix(a, "W/") == strings.TrimPrefix(b, "W/")
}

// etagStrongMatch reports whether a and b match using strong ETag comparison.
// Assumes a and b are valid ETags.
func etagStrongMatch(a, b string) bool {
	return a == b && a != "" && a[0] == '"'
}

// scanETag determines if a syntactically valid ETag is present at s. If so,
// the ETag and remaining text after consuming ETag is returned. Otherwise,
// it returns "", "".
func scanETag(s string) (etag string, remain string) {
	start := 0
	if len(s) >= 2 && s[0] == 'W' && s[1] == '/' {
		start = 2
	}
	if len(s[start:]) < 2 || s[start] != '"' {
		return "", ""
	}
	// ETag is either W/"text" or "text".
	// See RFC 7232 2.3.
	for i := start + 1; i < len(s); i++ {
		c := s[i]
		switch {
		// Character values allowed in ETags.
		case c == 0x21 || c >= 0x23 && c <= 0x7E || c >= 0x80:
		case c == '"':
			return s[:i+1], s[i+1:]
		default:
			return "", ""
		}
	}
	return "", ""
}

func writeNotModified(w http.ResponseWriter) {
	h := w.Header()
	delete(h, "Content-Type")
	delete(h, "Content-Length")
	w.WriteHeader(http.StatusNotModified)
}
<commit_msg>Fix If-None-Match header handling<commit_after>package utils

import (
	"io"
	"net/http"
	"net/textproto"
	"strconv"
	"strings"
)

// ServeContent replies to the request using the content in the provided
// reader. The Content-Length and Content-Type headers are added with the
// provided values.
func ServeContent(w http.ResponseWriter, r *http.Request, contentType string, size int64, content io.Reader) {
	h := w.Header()
	if size > 0 {
		h.Set("Content-Length", strconv.FormatInt(size, 10))
	}
	if contentType != "" {
		h.Set("Content-Type", contentType)
	}
	w.WriteHeader(http.StatusOK)
	if r.Method != "HEAD" {
		io.Copy(w, content)
	}
}

// CheckPreconditions evaluates request preconditions based only on the Etag
// values.
func CheckPreconditions(w http.ResponseWriter, r *http.Request, etag string) (done bool) {
	if etag != "" && !checkIfNoneMatch(w, r, etag) {
		writeNotModified(w)
		return true
	}
	return false
}

func checkIfNoneMatch(w http.ResponseWriter, r *http.Request, definedETag string) (match bool) {
	inm := r.Header.Get("If-None-Match")
	if inm == "" {
		return false
	}
	buf := inm
	for {
		buf = textproto.TrimString(buf)
		if len(buf) == 0 {
			break
		}
		if buf[0] == ',' {
			buf = buf[1:]
		}
		if buf[0] == '*' {
			return false
		}
		etag, remain := scanETag(buf)
		if etag == "" {
			break
		}
		if etagWeakMatch(etag, definedETag) {
			return false
		}
		buf = remain
	}
	return true
}

// etagWeakMatch reports whether a and b match using weak ETag comparison.
// Assumes a and b are valid ETags.
func etagWeakMatch(a, b string) bool {
	return strings.TrimPrefix(a, "W/") == strings.TrimPrefix(b, "W/")
}

// etagStrongMatch reports whether a and b match using strong ETag comparison.
// Assumes a and b are valid ETags.
func etagStrongMatch(a, b string) bool {
	return a == b && a != "" && a[0] == '"'
}

// scanETag determines if a syntactically valid ETag is present at s. If so,
// the ETag and remaining text after consuming ETag is returned. Otherwise,
// it returns "", "".
func scanETag(s string) (etag string, remain string) {
	start := 0
	if len(s) >= 2 && s[0] == 'W' && s[1] == '/' {
		start = 2
	}
	if len(s[start:]) < 2 || s[start] != '"' {
		return "", ""
	}
	// ETag is either W/"text" or "text".
	// See RFC 7232 2.3.
	for i := start + 1; i < len(s); i++ {
		c := s[i]
		switch {
		// Character values allowed in ETags.
		case c == 0x21 || c >= 0x23 && c <= 0x7E || c >= 0x80:
		case c == '"':
			return s[:i+1], s[i+1:]
		default:
			return "", ""
		}
	}
	return "", ""
}

func writeNotModified(w http.ResponseWriter) {
	h := w.Header()
	delete(h, "Content-Type")
	delete(h, "Content-Length")
	w.WriteHeader(http.StatusNotModified)
}
<|endoftext|>


<commit_before>package models

import (
	"database/sql"
	"errors"
	"time"
)

// ReleaseStatus is a type alias which will be used to create an enum of acceptable release status states.
type ReleaseStatus string

// ReleaseStatus pseudo-enum values
const (
	RStatusReleased ReleaseStatus = "released"
	RStatusDraft    ReleaseStatus = "draft"
)

var (
	RStatuses = []ReleaseStatus{RStatusReleased, RStatusDraft}
)

// Errors pertaining to the data in a Release or operations on Releases.
var (
	ErrInvalidReleaseStatus = errors.New("Invalid release status.")
	ErrNoSuchRelease        = errors.New("Could not find release.")
)

// Database queries for operations on Releases.
const (
	QInitTableReleases string = `create table if not exists releases (
		id int not null primary key,
		chapter varchar(255),
		version int,
		status varchar(255),
		checksum varchar(255),
		released_on timestamp,
		project_id int,
		foreign key(project_id) references projects(id)
);`

	QSaveRelease string = `insert into releases (
		chapter, version, status, checksum, released_on
) values (
		$1, $2, $3, $4, $5
);`

	QUpdateRelease string = `update releases set
chapter = $2, version = $3, status = $4, checksum = $5, released_on = $6
where id = $1;`

	QDeleteRelease string = `delete from releases where id = $1;`

	QListReleases string = `select (
		id, chapter, version, status, checksum, released_on
) from releases
where project_id = $1;`

	QFindRelease string = `select (
		chapter, version, status, checksum, released_on
) from releases
where id = $1;`
)

// Release contains information about a release, which there are many of under a given Project.  It contains information
// about which chapter of manga the release was created for, which version of the release of said chapter it is for, and
// the status of the release of the chapter itself, which may not be final right away.
type Release struct {
	Id         int           `json:"id"`
	Chapter    string        `json:"chapter"`
	Version    int           `json:"version"`
	Status     ReleaseStatus `json:"status"`
	Checksum   string        `json:"checksum"`
	ReleasedOn time.Time     `json:"releasedOn:`
}

// NewRelease constructs a brand new Release instance, with a default state lacking information its (future) position in
// a database.
func NewRelease(version int, chapterName string) Release {
	return Release{
		0,
		chapterName,
		version,
		RStatusDraft,
		"",
		time.Now(), // TODO -Update the release time when the status changes
	}
}

// FindRelease attempts to lookup a release by ID.
func FindRelease(id int, db *sql.DB) (Release, error) {
	r := Release{}
	row := db.QueryRow(QFindRelease, id)
	if row == nil {
		return Release{}, ErrNoSuchRelease
	}
	err := row.Scan(&r.Chapter, &r.Version, &r.Status, &r.Checksum, &r.ReleasedOn)
	if err != nil {
		return Release{}, err
	}
	r.Id = id
	return r, nil
}

// ListReleases attempts to obtain a list of all of the releases in the database.
func ListReleases(projectId int, db *sql.DB) ([]Release, error) {
	releases := []Release{}
	rows, err := db.Query(QListReleases, projectId)
	if err != nil {
		return []Release{}, err
	}
	defer rows.Close()
	for rows.Next() {
		var id, version int
		var chapter, status, checksum string
		var released time.Time
		scanErr := rows.Scan(&id, &chapter, &version, &status, &checksum, &released)
		if scanErr != nil {
			err = scanErr
		}
		releases = append(releases, Release{id, chapter, version, ReleaseStatus(status), checksum, released})
	}
	return releases, err
}

// Validate checks that the "status" of the project is one of the accepted ReleaseStatus values.
func (r *Release) Validate() error {
	for _, status := range RStatuses {
		if r.Status == status {
			return nil
		}
	}
	return ErrInvalidReleaseStatus
}

// Save inserts the release into the database and updates its Id field.
func (r *Release) Save(db *sql.DB) error {
	// TODO - Where should we compute checksums?
	_, err := db.Exec(QSaveRelease, r.Chapter, r.Version, r.Status, r.Checksum, r.ReleasedOn)
	if err != nil {
		return err
	}
	row := db.QueryRow(QLastInsertID)
	if row == nil {
		return ErrCouldNotGetID
	}
	return row.Scan(&r.Id)
}

// Update modifies all of the fields of a Release in place with whatever is currently in the struct.
func (r *Release) Update(db *sql.DB) error {
	_, err := db.Exec(QUpdateRelease, r.Id, r.Chapter, r.Version, r.Status, r.Checksum, time.Now())
	return err
}

// Delete removes the Release and all associated pages from the database.
func (r *Release) Delete(db *sql.DB) error {
	// TODO - Delete all associated pages.
	_, err := db.Exec(QDeleteRelease, r.Id)
	return err
}
<commit_msg>Update the released_on field of a release when its updated<commit_after>package models

import (
	"database/sql"
	"errors"
	"time"
)

// ReleaseStatus is a type alias which will be used to create an enum of acceptable release status states.
type ReleaseStatus string

// ReleaseStatus pseudo-enum values
const (
	RStatusReleased ReleaseStatus = "released"
	RStatusDraft    ReleaseStatus = "draft"
)

var (
	RStatuses = []ReleaseStatus{RStatusReleased, RStatusDraft}
)

// Errors pertaining to the data in a Release or operations on Releases.
var (
	ErrInvalidReleaseStatus = errors.New("Invalid release status.")
	ErrNoSuchRelease        = errors.New("Could not find release.")
)

// Database queries for operations on Releases.
const (
	QInitTableReleases string = `create table if not exists releases (
		id int not null primary key,
		chapter varchar(255),
		version int,
		status varchar(255),
		checksum varchar(255),
		released_on timestamp,
		project_id int,
		foreign key(project_id) references projects(id)
);`

	QSaveRelease string = `insert into releases (
		chapter, version, status, checksum, released_on
) values (
		$1, $2, $3, $4, $5
);`

	QUpdateRelease string = `update releases set
chapter = $2, version = $3, status = $4, checksum = $5, released_on = $6
where id = $1;`

	QDeleteRelease string = `delete from releases where id = $1;`

	QListReleases string = `select (
		id, chapter, version, status, checksum, released_on
) from releases
where project_id = $1;`

	QFindRelease string = `select (
		chapter, version, status, checksum, released_on
) from releases
where id = $1;`
)

// Release contains information about a release, which there are many of under a given Project.  It contains information
// about which chapter of manga the release was created for, which version of the release of said chapter it is for, and
// the status of the release of the chapter itself, which may not be final right away.
type Release struct {
	Id         int           `json:"id"`
	Chapter    string        `json:"chapter"`
	Version    int           `json:"version"`
	Status     ReleaseStatus `json:"status"`
	Checksum   string        `json:"checksum"`
	ReleasedOn time.Time     `json:"releasedOn:`
}

// NewRelease constructs a brand new Release instance, with a default state lacking information its (future) position in
// a database.
func NewRelease(version int, chapterName string) Release {
	return Release{
		0,
		chapterName,
		version,
		RStatusDraft,
		"",
		time.Now(),
	}
}

// FindRelease attempts to lookup a release by ID.
func FindRelease(id int, db *sql.DB) (Release, error) {
	r := Release{}
	row := db.QueryRow(QFindRelease, id)
	if row == nil {
		return Release{}, ErrNoSuchRelease
	}
	err := row.Scan(&r.Chapter, &r.Version, &r.Status, &r.Checksum, &r.ReleasedOn)
	if err != nil {
		return Release{}, err
	}
	r.Id = id
	return r, nil
}

// ListReleases attempts to obtain a list of all of the releases in the database.
func ListReleases(projectId int, db *sql.DB) ([]Release, error) {
	releases := []Release{}
	rows, err := db.Query(QListReleases, projectId)
	if err != nil {
		return []Release{}, err
	}
	defer rows.Close()
	for rows.Next() {
		var id, version int
		var chapter, status, checksum string
		var released time.Time
		scanErr := rows.Scan(&id, &chapter, &version, &status, &checksum, &released)
		if scanErr != nil {
			err = scanErr
		}
		releases = append(releases, Release{id, chapter, version, ReleaseStatus(status), checksum, released})
	}
	return releases, err
}

// Validate checks that the "status" of the project is one of the accepted ReleaseStatus values.
func (r *Release) Validate() error {
	for _, status := range RStatuses {
		if r.Status == status {
			return nil
		}
	}
	return ErrInvalidReleaseStatus
}

// Save inserts the release into the database and updates its Id field.
func (r *Release) Save(db *sql.DB) error {
	// TODO - Where should we compute checksums?
	_, err := db.Exec(QSaveRelease, r.Chapter, r.Version, r.Status, r.Checksum, r.ReleasedOn)
	if err != nil {
		return err
	}
	row := db.QueryRow(QLastInsertID)
	if row == nil {
		return ErrCouldNotGetID
	}
	return row.Scan(&r.Id)
}

// Update modifies all of the fields of a Release in place with whatever is currently in the struct.
func (r *Release) Update(db *sql.DB) error {
	now := time.Now()
	_, err := db.Exec(QUpdateRelease, r.Id, r.Chapter, r.Version, r.Status, r.Checksum, now)
	r.ReleasedOn = now
	return err
}

// Delete removes the Release and all associated pages from the database.
func (r *Release) Delete(db *sql.DB) error {
	// TODO - Delete all associated pages.
	_, err := db.Exec(QDeleteRelease, r.Id)
	return err
}
<|endoftext|>


<commit_before>package device

import "github.com/tarm/serial"

type Manager struct {
	devices map[string]serial.Config
}
<commit_msg>Add *Conn.Open<commit_after>package device

import "github.com/tarm/serial"

type Manager struct {
	devices map[string]serial.Config
}

type Conn struct {
	device serial.Config
	port   *serial.Port
	isOpen bool
}

func (c *Conn) Open() error {
	p, err := serial.OpenPort(&c.device)
	if err != nil {
		return nil
	}
	c.port = p
	c.isOpen = true
	return nil
}
<|endoftext|>


<commit_before>package vm

import (
	"strconv"

	"github.com/dlclark/regexp2"
	"github.com/goby-lang/goby/vm/classes"
	"github.com/goby-lang/goby/vm/errors"
)

// RegexpObject represents regexp instances, which of the type is actually string.
// Regexp object holds regexp strings.
// Powered by: github.com/dlclark/regexp2
// The regexp2 package had been ported from .NET Framework's regexp library, which is PCRE-compatible
// and is almost all equivalent to Ruby's Onigmo regexp library.
//
// ```ruby
// a = Regexp.new("orl")
// a.match?("Hello World")   #=> true
// a.match?("Hello Regexp")  #=> false
//
// b = Regexp.new("😏")
// b.match?("🤡 😏 😐")   #=> true
// b.match?("😝 😍 😊")   #=> false
//
// c = Regexp.new("居(ら(?=れ)|さ(?=せ)|る|ろ|れ(?=[ばる])|よ|(?=な[いかくけそ]|ま[しすせ]|そう|た|て))")
// c.match?("居られればいいのに")  #=> true
// c.match?("居ずまいを正す")      #=> false
// ```
//
// **Note:**
//
// - Currently, manipulations are based upon Golang's Unicode manipulations.
// - Currently, UTF-8 encoding is assumed based upon Golang's string manipulation, but the encoding is not actually specified(TBD).
// - `Regexp.new` is exceptionally supported.
//
// **To Goby maintainers**: avoid using Go's standard regexp package (slow and not rich). Consider the faster `Trim` or `Split` etc in Go's "strings" package first, or just use the dlclark/regexp2 instead.
// ToDo: Regexp literals with '/.../' and match operator `=~`
type RegexpObject struct {
	*baseObj
	Regexp *regexp2.Regexp
}

// Class methods --------------------------------------------------------
func builtInRegexpClassMethods() []*BuiltinMethodObject {
	return []*BuiltinMethodObject{
		{
			Name: "new",
			Fn: func(receiver Object) builtinMethodBody {
				return func(t *thread, args []Object, blockFrame *callFrame) Object {
					r := t.vm.initRegexpObject(args[0].toString())
					if r == nil {
						return t.vm.initErrorObject(errors.ArgumentError, "Invalid regexp: %v", args[0].toString())
					}
					return r
				}
			},
		},
	}
}

// Instance methods -----------------------------------------------------
func builtinRegexpInstanceMethods() []*BuiltinMethodObject {
	return []*BuiltinMethodObject{

		{
			// Returns boolean value to indicate the result of regexp match with the string given.
			//
			// ```ruby
			// r = Regexp.new("o")
			// r.match("pow")  # => true
			// r.match("gee")  # => false
			// ```
			//
			// @param string [String]
			// @return [Boolean]
			Name: "match?",
			Fn: func(receiver Object) builtinMethodBody {
				return func(t *thread, args []Object, blockFrame *callFrame) Object {

					if len(args) != 1 {
						return t.vm.initErrorObject(errors.ArgumentError, "Expect 1 argument. got=%v", strconv.Itoa(len(args)))
					}

					re := receiver.(*RegexpObject).Regexp
					m, _ := re.MatchString(args[0].toString())
					if m {
						return TRUE
					}
					return FALSE
				}
			},
		},
	}
}

// Internal functions ===================================================

// Functions for initialization -----------------------------------------

func (vm *VM) initRegexpObject(regexp string) *RegexpObject {
	r, err := regexp2.Compile(regexp, 0)
	if err != nil {
		return nil
	}
	return &RegexpObject{
		baseObj: &baseObj{class: vm.topLevelClass(classes.RegexpClass)},
		Regexp:  r,
	}
}

func (vm *VM) initRegexpClass() *RClass {
	rc := vm.initializeClass(classes.RegexpClass, false)
	rc.setBuiltinMethods(builtinRegexpInstanceMethods(), false)
	rc.setBuiltinMethods(builtInRegexpClassMethods(), true)
	return rc
}

// Polymorphic helper functions -----------------------------------------

// Value returns the object
func (r *RegexpObject) Value() interface{} {
	return r.toString()
}

// toString returns the object's name as the string format
func (r *RegexpObject) toString() string {
	return r.Regexp.String()
}

// toJSON just delegates to toString
func (r *RegexpObject) toJSON() string {
	return "\"" + r.toString() + "\""
}

// equal checks if the string values between receiver and argument are equal
func (r *RegexpObject) equal(e *RegexpObject) bool {
	return r.toString() == r.toString()
}
<commit_msg>Add '==' method to Regexp<commit_after>package vm

import (
	"strconv"

	"github.com/dlclark/regexp2"
	"github.com/goby-lang/goby/vm/classes"
	"github.com/goby-lang/goby/vm/errors"
)

// RegexpObject represents regexp instances, which of the type is actually string.
// Regexp object holds regexp strings.
// Powered by: github.com/dlclark/regexp2
// The regexp2 package had been ported from .NET Framework's regexp library, which is PCRE-compatible
// and is almost all equivalent to Ruby's Onigmo regexp library.
//
// ```ruby
// a = Regexp.new("orl")
// a.match?("Hello World")   #=> true
// a.match?("Hello Regexp")  #=> false
//
// b = Regexp.new("😏")
// b.match?("🤡 😏 😐")   #=> true
// b.match?("😝 😍 😊")   #=> false
//
// c = Regexp.new("居(ら(?=れ)|さ(?=せ)|る|ろ|れ(?=[ばる])|よ|(?=な[いかくけそ]|ま[しすせ]|そう|た|て))")
// c.match?("居られればいいのに")  #=> true
// c.match?("居ずまいを正す")      #=> false
// ```
//
// **Note:**
//
// - Currently, manipulations are based upon Golang's Unicode manipulations.
// - Currently, UTF-8 encoding is assumed based upon Golang's string manipulation, but the encoding is not actually specified(TBD).
// - `Regexp.new` is exceptionally supported.
//
// **To Goby maintainers**: avoid using Go's standard regexp package (slow and not rich). Consider the faster `Trim` or `Split` etc in Go's "strings" package first, or just use the dlclark/regexp2 instead.
// ToDo: Regexp literals with '/.../' and match operator `=~`
type RegexpObject struct {
	*baseObj
	Regexp *regexp2.Regexp
}

// Class methods --------------------------------------------------------
func builtInRegexpClassMethods() []*BuiltinMethodObject {
	return []*BuiltinMethodObject{
		{
			Name: "new",
			Fn: func(receiver Object) builtinMethodBody {
				return func(t *thread, args []Object, blockFrame *callFrame) Object {
					r := t.vm.initRegexpObject(args[0].toString())
					if r == nil {
						return t.vm.initErrorObject(errors.ArgumentError, "Invalid regexp: %v", args[0].toString())
					}
					return r
				}
			},
		},
	}
}

// Instance methods -----------------------------------------------------
func builtinRegexpInstanceMethods() []*BuiltinMethodObject {
	return []*BuiltinMethodObject{

		{
			// Returns true if the two regexp patterns are exactly the same, or returns false if not.
			// If comparing with non Regexp class, just returns false.
			//
			// ```ruby
			// r1 = Regexp.new("goby[0-9]+")
			// r2 = Regexp.new("goby[0-9]+")
			// r3 = Regexp.new("Goby[0-9]+")
			//
			// r1 == r2   # => true
			// r1 == r2   # => false
			// ```
			//
			// @param regexp [Regexp]
			// @return [Boolean]
			Name: "==",
			Fn: func(receiver Object) builtinMethodBody {
				return func(t *thread, args []Object, blockFrame *callFrame) Object {

					var rv string
					if len(args) != 1 {
						return t.vm.initErrorObject(errors.ArgumentError, "Expect 1 argument. got=%v", strconv.Itoa(len(args)))
					}

					lv := receiver.(*RegexpObject).Value()
					arg := args[0]

					if n := arg.Class().Name; n == "Regexp" {
						rv = arg.toString()
					} else {
						if n != "String" {
							rv = arg.toString()
						} else {
							return FALSE
						}
					}

					if lv == rv {
						return TRUE
					}
					return FALSE
				}
			},
		},
		{
			// Returns boolean value to indicate the result of regexp match with the string given.
			//
			// ```ruby
			// r = Regexp.new("o")
			// r.match?("pow")  # => true
			// r.match?("gee")  # => false
			// ```
			//
			// @param string [String]
			// @return [Boolean]
			Name: "match?",
			Fn: func(receiver Object) builtinMethodBody {
				return func(t *thread, args []Object, blockFrame *callFrame) Object {

					if len(args) != 1 {
						return t.vm.initErrorObject(errors.ArgumentError, "Expect 1 argument. got=%v", strconv.Itoa(len(args)))
					}

					re := receiver.(*RegexpObject).Regexp
					m, _ := re.MatchString(args[0].toString())
					if m {
						return TRUE
					}
					return FALSE
				}
			},
		},
	}
}

// Internal functions ===================================================

// Functions for initialization -----------------------------------------

func (vm *VM) initRegexpObject(regexp string) *RegexpObject {
	r, err := regexp2.Compile(regexp, 0)
	if err != nil {
		return nil
	}
	return &RegexpObject{
		baseObj: &baseObj{class: vm.topLevelClass(classes.RegexpClass)},
		Regexp:  r,
	}
}

func (vm *VM) initRegexpClass() *RClass {
	rc := vm.initializeClass(classes.RegexpClass, false)
	rc.setBuiltinMethods(builtinRegexpInstanceMethods(), false)
	rc.setBuiltinMethods(builtInRegexpClassMethods(), true)
	return rc
}

// Polymorphic helper functions -----------------------------------------

// Value returns the object
func (r *RegexpObject) Value() interface{} {
	return r.toString()
}

// toString returns the object's name as the string format
func (r *RegexpObject) toString() string {
	return r.Regexp.String()
}

// toJSON just delegates to toString
func (r *RegexpObject) toJSON() string {
	return "\"" + r.toString() + "\""
}

// equal checks if the string values between receiver and argument are equal
func (r *RegexpObject) equal(e *RegexpObject) bool {
	return r.toString() == r.toString()
}
<|endoftext|>