<commit_before>package main

import (
	"context"
	"flag"
	"fmt"
	"math"
	"math/rand"
	"net/http"
	"os"
	"sync/atomic"
	"testing"

	"github.com/bmatsuo/lmdb-go/lmdb"
	"github.com/dgraph-io/badger"
	"github.com/dgraph-io/badger-bench/store"
	"github.com/dgraph-io/badger/table"
	"github.com/dgraph-io/badger/y"
)

var (
	ctx           = context.Background()
	numKeys       = flag.Float64("keys_mil", 10.0, "How many million keys to write.")
	flagDir       = flag.String("dir", "bench-tmp", "Where data is temporarily stored.")
	flagValueSize = flag.Int("valsz", 128, "Size of each value.")
)

const Mi int = 1000000
const Mf float64 = 1000000

func getBadger() (*badger.KV, error) {
	opt := badger.DefaultOptions
	opt.MapTablesTo = table.LoadToRAM
	opt.Dir = *flagDir + "/badger"
	opt.ValueDir = opt.Dir
	fmt.Println(opt.Dir)
	opt.DoNotCompact = true
	opt.ValueGCThreshold = 0.0
	return badger.NewKV(&opt)
}

func getRocks() *store.Store {
	rdb, err := store.NewReadOnlyStore(*flagDir + "/rocks")
	y.Check(err)
	return rdb
}

func getLmdb() *lmdb.Env {
	lmdbEnv, err := lmdb.NewEnv()
	y.Check(err)
	err = lmdbEnv.SetMaxReaders(math.MaxInt64)
	y.Check(err)
	err = lmdbEnv.SetMaxDBs(1)
	y.Check(err)
	err = lmdbEnv.SetMapSize(1 << 38) // ~273Gb
	y.Check(err)

	err = lmdbEnv.Open(*flagDir+"/lmdb", lmdb.Readonly, 0777)
	y.Check(err)
	return lmdbEnv
}

func newKey() []byte {
	k := rand.Int() % int(*numKeys*Mf)
	key := fmt.Sprintf("vsz=%05d-k=%010d", *flagValueSize, k) // 22 bytes.
	return []byte(key)
}

func print(count int) {
	if count%100000 == 0 {
		fmt.Printf(".")
	} else if count%Mi == 0 {
		fmt.Printf("-")
	}
}

func BenchmarkReadRandomBadger(b *testing.B) {
	bdb, err := getBadger()
	y.Check(err)
	defer bdb.Close()

	var totalFound uint64
	var totalErr uint64
	var totalNotFound uint64
	b.Run("read-random-badger", func(b *testing.B) {
		b.RunParallel(func(pb *testing.PB) {
			var found, errored, notFound uint64
			for pb.Next() {
				key := newKey()
				var val badger.KVItem
				if err := bdb.Get(key, &val); err == nil && val.Value() != nil {
					found++
				} else if err != nil {
					errored++
				} else {
					notFound++
				}
			}
			atomic.AddUint64(&totalFound, found)
			atomic.AddUint64(&totalErr, errored)
			atomic.AddUint64(&totalNotFound, notFound)
		})
	})
	b.Logf("badger %d keys had valid values.", totalFound)
	b.Logf("badger %d keys had no values", totalNotFound)
	b.Logf("badger %d keys had errors", totalErr)
	b.Logf("badger hit rate : %.2f", float64(totalFound)/float64(totalFound+totalNotFound+totalErr))
}

func BenchmarkReadRandomRocks(b *testing.B) {
	rdb := getRocks()
	defer rdb.Close()

	var totalCount uint64
	b.Run("read-random-rocks", func(b *testing.B) {
		b.RunParallel(func(pb *testing.PB) {
			var count uint64
			for pb.Next() {
				key := newKey()
				if _, err := rdb.Get(key); err == nil {
					count++
				}
			}
			atomic.AddUint64(&totalCount, count)
		})
	})
	b.Logf("rocks %d keys had valid values.", totalCount)
}

func BenchmarkReadRandomLmdb(b *testing.B) {
	lmdbEnv := getLmdb()
	defer lmdbEnv.Close()

	var lmdbDBI lmdb.DBI
	// Acquire handle
	err := lmdbEnv.View(func(txn *lmdb.Txn) error {
		var err error
		lmdbDBI, err = txn.OpenDBI("bench", 0)
		return err
	})
	y.Check(err)
	defer lmdbEnv.CloseDBI(lmdbDBI)

	var totalFound uint64
	var totalErr uint64
	var totalNotFound uint64
	b.Run("read-random-lmdb", func(b *testing.B) {
		b.RunParallel(func(pb *testing.PB) {
			var found, errored, notFound uint64

			for pb.Next() {
				key := newKey()
				_ = lmdbEnv.View(func(txn *lmdb.Txn) error {
					_, err := txn.Get(lmdbDBI, key)
					if lmdb.IsNotFound(err) {
						notFound++
						return nil

					} else if err != nil {
						errored++
						return err
					}
					found++
					return nil
				})
			}
			atomic.AddUint64(&totalFound, found)
			atomic.AddUint64(&totalErr, errored)
			atomic.AddUint64(&totalNotFound, notFound)
		})
	})
	b.Logf("lmdb %d keys had valid values.", totalFound)
	b.Logf("lmdb %d keys had no values", totalNotFound)
	b.Logf("lmdb %d keys had errors", totalErr)
	b.Logf("lmdb hit rate : %.2f", float64(totalFound)/float64(totalFound+totalNotFound+totalErr))
}

func safecopy(dst []byte, src []byte) []byte {
	if cap(dst) < len(src) {
		dst = make([]byte, len(src))
	}
	dst = dst[0:len(src)]
	copy(dst, src)
	return dst
}

func BenchmarkIterateRocks(b *testing.B) {
	rdb := getRocks()
	k := make([]byte, 1024)
	v := make([]byte, Mi)
	b.ResetTimer()

	b.Run("rocksdb-iterate", func(b *testing.B) {
		for j := 0; j < b.N; j++ {
			itr := rdb.NewIterator()
			var count int
			for itr.SeekToFirst(); itr.Valid(); itr.Next() {
				{
					// do some processing.
					k = safecopy(k, itr.Key().Data())
					v = safecopy(v, itr.Value().Data())
				}
				count++
				print(count)
				if count > 2*Mi {
					break
				}
			}
			b.Logf("[%d] Counted %d keys\n", j, count)
		}
	})
}

func BenchmarkIterateLmdb(b *testing.B) {
	lmdbEnv := getLmdb()

	var lmdbDBI lmdb.DBI
	// Acquire handle
	err := lmdbEnv.View(func(txn *lmdb.Txn) error {
		var err error
		lmdbDBI, err = txn.OpenDBI("bench", 0)
		return err
	})
	y.Check(err)
	defer lmdbEnv.CloseDBI(lmdbDBI)

	k := make([]byte, 1024)
	v := make([]byte, Mi)
	b.ResetTimer()

	b.Run("lmdb-iterate", func(b *testing.B) {
		for j := 0; j < b.N; j++ {
			var count int
			err = lmdbEnv.View(func(txn *lmdb.Txn) error {
				cur, err := txn.OpenCursor(lmdbDBI)
				if err != nil {
					return err
				}
				defer cur.Close()

				for {
					k1, v1, err := cur.Get(nil, nil, lmdb.Next)
					if lmdb.IsNotFound(err) {
						return nil
					}
					if err != nil {
						return err
					}

					//fmt.Printf("%s %s\n", k, v)

					// do some processing.
					k = safecopy(k, k1)
					v = safecopy(v, v1)

					count++
					print(count)
					if count > 2*Mi {
						break
					}
				}
				return nil
			})
			y.Check(err)
			b.Logf("[%d] Counted %d keys\n", j, count)
		}
	})
}

func BenchmarkIterateBadgerOnlyKeys(b *testing.B) {
	bdb, err := getBadger()
	y.Check(err)
	k := make([]byte, 1024)
	b.ResetTimer()

	b.Run("badger-iterate-onlykeys", func(b *testing.B) {
		for j := 0; j < b.N; j++ {
			var count int
			// 100 = size, 0 = num workers, false = fwd direction.
			opt := badger.IteratorOptions{}
			opt.PrefetchSize = 10000
			itr := bdb.NewIterator(opt)
			for itr.Rewind(); itr.Valid(); itr.Next() {
				item := itr.Item()
				{
					// do some processing.
					k = safecopy(k, item.Key())
				}
				count++
				print(count)
				if count > 2*Mi {
					break
				}
			}
			b.Logf("[%d] Counted %d keys\n", j, count)
		}
	})
}

func BenchmarkIterateBadgerWithValues(b *testing.B) {
	bdb, err := getBadger()
	y.Check(err)
	k := make([]byte, 1024)
	v := make([]byte, Mi)
	b.ResetTimer()

	b.Run("badger-iterate-withvals", func(b *testing.B) {
		for j := 0; j < b.N; j++ {
			var count int
			opt := badger.IteratorOptions{}
			opt.PrefetchSize = 10000
			opt.FetchValues = true
			itr := bdb.NewIterator(opt)
			for itr.Rewind(); itr.Valid(); itr.Next() {
				item := itr.Item()
				{
					// do some processing.
					k = safecopy(k, item.Key())
					v = safecopy(v, item.Value())
				}
				count++
				print(count)
				if count >= 2*Mi {
					break
				}
			}
			b.Logf("[%d] Counted %d keys\n", j, count)
		}
	})
}

func TestMain(m *testing.M) {
	flag.Parse()
	// call flag.Parse() here if TestMain uses flags
	go http.ListenAndServe(":8080", nil)
	os.Exit(m.Run())
}
<commit_msg>Check for error on lmdbEnv.View call<commit_after>package main

import (
	"context"
	"flag"
	"fmt"
	"math"
	"math/rand"
	"net/http"
	"os"
	"sync/atomic"
	"testing"

	"github.com/bmatsuo/lmdb-go/lmdb"
	"github.com/dgraph-io/badger"
	"github.com/dgraph-io/badger-bench/store"
	"github.com/dgraph-io/badger/table"
	"github.com/dgraph-io/badger/y"
)

var (
	ctx           = context.Background()
	numKeys       = flag.Float64("keys_mil", 10.0, "How many million keys to write.")
	flagDir       = flag.String("dir", "bench-tmp", "Where data is temporarily stored.")
	flagValueSize = flag.Int("valsz", 128, "Size of each value.")
)

const Mi int = 1000000
const Mf float64 = 1000000

func getBadger() (*badger.KV, error) {
	opt := badger.DefaultOptions
	opt.MapTablesTo = table.LoadToRAM
	opt.Dir = *flagDir + "/badger"
	opt.ValueDir = opt.Dir
	fmt.Println(opt.Dir)
	opt.DoNotCompact = true
	opt.ValueGCThreshold = 0.0
	return badger.NewKV(&opt)
}

func getRocks() *store.Store {
	rdb, err := store.NewReadOnlyStore(*flagDir + "/rocks")
	y.Check(err)
	return rdb
}

func getLmdb() *lmdb.Env {
	lmdbEnv, err := lmdb.NewEnv()
	y.Check(err)
	err = lmdbEnv.SetMaxReaders(math.MaxInt64)
	y.Check(err)
	err = lmdbEnv.SetMaxDBs(1)
	y.Check(err)
	err = lmdbEnv.SetMapSize(1 << 38) // ~273Gb
	y.Check(err)

	err = lmdbEnv.Open(*flagDir+"/lmdb", lmdb.Readonly, 0777)
	y.Check(err)
	return lmdbEnv
}

func newKey() []byte {
	k := rand.Int() % int(*numKeys*Mf)
	key := fmt.Sprintf("vsz=%05d-k=%010d", *flagValueSize, k) // 22 bytes.
	return []byte(key)
}

func print(count int) {
	if count%100000 == 0 {
		fmt.Printf(".")
	} else if count%Mi == 0 {
		fmt.Printf("-")
	}
}

func BenchmarkReadRandomBadger(b *testing.B) {
	bdb, err := getBadger()
	y.Check(err)
	defer bdb.Close()

	var totalFound uint64
	var totalErr uint64
	var totalNotFound uint64
	b.Run("read-random-badger", func(b *testing.B) {
		b.RunParallel(func(pb *testing.PB) {
			var found, errored, notFound uint64
			for pb.Next() {
				key := newKey()
				var val badger.KVItem
				if err := bdb.Get(key, &val); err == nil && val.Value() != nil {
					found++
				} else if err != nil {
					errored++
				} else {
					notFound++
				}
			}
			atomic.AddUint64(&totalFound, found)
			atomic.AddUint64(&totalErr, errored)
			atomic.AddUint64(&totalNotFound, notFound)
		})
	})
	b.Logf("badger %d keys had valid values.", totalFound)
	b.Logf("badger %d keys had no values", totalNotFound)
	b.Logf("badger %d keys had errors", totalErr)
	b.Logf("badger hit rate : %.2f", float64(totalFound)/float64(totalFound+totalNotFound+totalErr))
}

func BenchmarkReadRandomRocks(b *testing.B) {
	rdb := getRocks()
	defer rdb.Close()

	var totalCount uint64
	b.Run("read-random-rocks", func(b *testing.B) {
		b.RunParallel(func(pb *testing.PB) {
			var count uint64
			for pb.Next() {
				key := newKey()
				if _, err := rdb.Get(key); err == nil {
					count++
				}
			}
			atomic.AddUint64(&totalCount, count)
		})
	})
	b.Logf("rocks %d keys had valid values.", totalCount)
}

func BenchmarkReadRandomLmdb(b *testing.B) {
	lmdbEnv := getLmdb()
	defer lmdbEnv.Close()

	var lmdbDBI lmdb.DBI
	// Acquire handle
	err := lmdbEnv.View(func(txn *lmdb.Txn) error {
		var err error
		lmdbDBI, err = txn.OpenDBI("bench", 0)
		return err
	})
	y.Check(err)
	defer lmdbEnv.CloseDBI(lmdbDBI)

	var totalFound uint64
	var totalErr uint64
	var totalNotFound uint64
	b.Run("read-random-lmdb", func(b *testing.B) {
		b.RunParallel(func(pb *testing.PB) {
			var found, errored, notFound uint64

			for pb.Next() {
				key := newKey()
				err = lmdbEnv.View(func(txn *lmdb.Txn) error {
					_, err := txn.Get(lmdbDBI, key)
					if lmdb.IsNotFound(err) {
						notFound++
						return nil

					} else if err != nil {
						return err
					}
					found++
					return nil
				})
				if err != nil {
					errored++
				}
			}
			atomic.AddUint64(&totalFound, found)
			atomic.AddUint64(&totalErr, errored)
			atomic.AddUint64(&totalNotFound, notFound)
		})
	})
	b.Logf("lmdb %d keys had valid values.", totalFound)
	b.Logf("lmdb %d keys had no values", totalNotFound)
	b.Logf("lmdb %d keys had errors", totalErr)
	b.Logf("lmdb hit rate : %.2f", float64(totalFound)/float64(totalFound+totalNotFound+totalErr))
}

func safecopy(dst []byte, src []byte) []byte {
	if cap(dst) < len(src) {
		dst = make([]byte, len(src))
	}
	dst = dst[0:len(src)]
	copy(dst, src)
	return dst
}

func BenchmarkIterateRocks(b *testing.B) {
	rdb := getRocks()
	k := make([]byte, 1024)
	v := make([]byte, Mi)
	b.ResetTimer()

	b.Run("rocksdb-iterate", func(b *testing.B) {
		for j := 0; j < b.N; j++ {
			itr := rdb.NewIterator()
			var count int
			for itr.SeekToFirst(); itr.Valid(); itr.Next() {
				{
					// do some processing.
					k = safecopy(k, itr.Key().Data())
					v = safecopy(v, itr.Value().Data())
				}
				count++
				print(count)
				if count > 2*Mi {
					break
				}
			}
			b.Logf("[%d] Counted %d keys\n", j, count)
		}
	})
}

func BenchmarkIterateLmdb(b *testing.B) {
	lmdbEnv := getLmdb()

	var lmdbDBI lmdb.DBI
	// Acquire handle
	err := lmdbEnv.View(func(txn *lmdb.Txn) error {
		var err error
		lmdbDBI, err = txn.OpenDBI("bench", 0)
		return err
	})
	y.Check(err)
	defer lmdbEnv.CloseDBI(lmdbDBI)

	k := make([]byte, 1024)
	v := make([]byte, Mi)
	b.ResetTimer()

	b.Run("lmdb-iterate", func(b *testing.B) {
		for j := 0; j < b.N; j++ {
			var count int
			err = lmdbEnv.View(func(txn *lmdb.Txn) error {
				cur, err := txn.OpenCursor(lmdbDBI)
				if err != nil {
					return err
				}
				defer cur.Close()

				for {
					k1, v1, err := cur.Get(nil, nil, lmdb.Next)
					if lmdb.IsNotFound(err) {
						return nil
					}
					if err != nil {
						return err
					}

					//fmt.Printf("%s %s\n", k, v)

					// do some processing.
					k = safecopy(k, k1)
					v = safecopy(v, v1)

					count++
					print(count)
					if count > 2*Mi {
						break
					}
				}
				return nil
			})
			y.Check(err)
			b.Logf("[%d] Counted %d keys\n", j, count)
		}
	})
}

func BenchmarkIterateBadgerOnlyKeys(b *testing.B) {
	bdb, err := getBadger()
	y.Check(err)
	k := make([]byte, 1024)
	b.ResetTimer()

	b.Run("badger-iterate-onlykeys", func(b *testing.B) {
		for j := 0; j < b.N; j++ {
			var count int
			// 100 = size, 0 = num workers, false = fwd direction.
			opt := badger.IteratorOptions{}
			opt.PrefetchSize = 10000
			itr := bdb.NewIterator(opt)
			for itr.Rewind(); itr.Valid(); itr.Next() {
				item := itr.Item()
				{
					// do some processing.
					k = safecopy(k, item.Key())
				}
				count++
				print(count)
				if count > 2*Mi {
					break
				}
			}
			b.Logf("[%d] Counted %d keys\n", j, count)
		}
	})
}

func BenchmarkIterateBadgerWithValues(b *testing.B) {
	bdb, err := getBadger()
	y.Check(err)
	k := make([]byte, 1024)
	v := make([]byte, Mi)
	b.ResetTimer()

	b.Run("badger-iterate-withvals", func(b *testing.B) {
		for j := 0; j < b.N; j++ {
			var count int
			opt := badger.IteratorOptions{}
			opt.PrefetchSize = 10000
			opt.FetchValues = true
			itr := bdb.NewIterator(opt)
			for itr.Rewind(); itr.Valid(); itr.Next() {
				item := itr.Item()
				{
					// do some processing.
					k = safecopy(k, item.Key())
					v = safecopy(v, item.Value())
				}
				count++
				print(count)
				if count >= 2*Mi {
					break
				}
			}
			b.Logf("[%d] Counted %d keys\n", j, count)
		}
	})
}

func TestMain(m *testing.M) {
	flag.Parse()
	// call flag.Parse() here if TestMain uses flags
	go http.ListenAndServe(":8080", nil)
	os.Exit(m.Run())
}
<|endoftext|>


<commit_before>package main

import "fmt"
import "math/big"

// A BigIntPoly represents a polynomial with big.Int coefficients mod
// some (N, X^R - 1).
//
// The zero value for a BigIntPoly represents the zero polynomial.
type BigIntPoly struct {
	R int
	// k is the number of big.Words required to hold a coefficient
	// in calculations without overflowing.
	k int
	// If p(x) is the polynomial as a function, phi is
	// p(2^{k*bitsize(big.Word)}). Since a coefficient fits into k
	// big.Words, this is a lossless transformation; that is, one
	// can recover all coefficients of p(x) from phi.
	phi big.Int
}

// Only polynomials built with the same value of N and R may be used
// together in one of the functions below.

// Builds a new BigIntPoly representing the zero polynomial
// mod (N, X^R - 1). R must fit into an int.
func NewBigIntPoly(N, R big.Int) *BigIntPoly {
	// A coefficient can be up to R*(N - 1)^2 in intermediate
	// calculations.
	var maxCoefficient big.Int
	maxCoefficient.Sub(&N, big.NewInt(1))
	maxCoefficient.Mul(&maxCoefficient, &maxCoefficient)
	maxCoefficient.Mul(&maxCoefficient, &R)

	var phi big.Int
	rInt := int(R.Int64())
	k := len(maxCoefficient.Bits())
	// Up to 2*R coefficients may be needed in intermediate
	// calculations.
	maxWordCount := 2 * rInt * k
	phi.SetBits(make([]big.Word, maxWordCount))
	return &BigIntPoly{rInt, k, phi}
}

// Returns 1 + the degree of this polynomial, or 0 if the polynomial
// is the zero polynomial.
func (p *BigIntPoly) getCoefficientCount() int {
	l := len(p.phi.Bits())
	if l == 0 {
		return 0
	}
	coefficientCount := l / p.k
	if l%p.k != 0 {
		coefficientCount++
	}
	return coefficientCount
}

// Sets the coefficient count to the given number, which must be at
// most p.R. The unused bytes of the leading coefficient must be
// cleared (via commitCoefficient()) prior to this being called.
func (p *BigIntPoly) setCoefficientCount(coefficientCount int) {
	p.phi.SetBits(p.phi.Bits()[0 : coefficientCount*p.k])
}

// Returns the ith coefficient of this polynomial. i must be less than
// p.getCoefficientCount().
func (p *BigIntPoly) getCoefficient(i int) big.Int {
	var c big.Int
	start := i * p.k
	if i == p.getCoefficientCount()-1 {
		// If the last coefficient is small enough, phi might
		// have fewer than p.R * p.k words.
		c.SetBits(p.phi.Bits()[start:])
	} else {
		end := (i + 1) * p.k
		c.SetBits(p.phi.Bits()[start:end])
	}
	return c
}

// Must be called after all changes have been made to a coefficient
// via a big.Int returned from p.getCoefficient().
func (p *BigIntPoly) commitCoefficient(c big.Int) {
	cBits := c.Bits()
	unusedBits := cBits[len(cBits):p.k]
	for j := 0; j < len(unusedBits); j++ {
		unusedBits[j] = 0
	}
}

// Sets p to X^k + a mod (N, X^R - 1).
func (p *BigIntPoly) Set(a, k, N big.Int) {
	c0 := p.getCoefficient(0)
	c0.Mod(&a, &N)
	p.commitCoefficient(c0)

	R := big.NewInt(int64(p.R))
	var kModRBig big.Int
	kModRBig.Mod(&k, R)
	kModR := int(kModRBig.Int64())

	for i := 1; i <= kModR; i++ {
		c := p.getCoefficient(i)
		c.Set(&big.Int{})
		p.commitCoefficient(c)
	}

	cKModR := p.getCoefficient(kModR)
	cKModR.Set(big.NewInt(1))
	p.commitCoefficient(cKModR)

	p.setCoefficientCount(kModR + 1)
}

// Returns whether p has the same coefficients as q.
func (p *BigIntPoly) Eq(q *BigIntPoly) bool {
	return p.phi.Cmp(&q.phi) == 0
}

// Sets p to the product of p and q mod (N, X^R - 1). Assumes R >=
// 2. tmp must not alias p or q.
func (p *BigIntPoly) mul(q *BigIntPoly, N big.Int, tmp *BigIntPoly) {
	tmp.phi.Mul(&p.phi, &q.phi)
	p.phi, tmp.phi = tmp.phi, p.phi

	// Mod p by X^R - 1.
	mid := p.R * p.k
	pBits := p.phi.Bits()
	if len(pBits) > mid {
		var lo, hi big.Int
		lo.SetBits(pBits[:mid])
		hi.SetBits(pBits[mid:])
		p.phi.Add(&lo, &hi)
	}

	// Mod p by N.
	newCoefficientCount := 0
	tmp2 := tmp.getCoefficient(0)
	tmp3 := tmp.getCoefficient(1)
	for i := 0; i < p.getCoefficientCount(); i++ {
		c := p.getCoefficient(i)
		if c.Cmp(&N) >= 0 {
			// Mod c by N. Use big.Int.QuoRem() instead of
			// big.Int.Mod() since the latter allocates an
			// extra big.Int.
			tmp2.QuoRem(&c, &N, &tmp3)
			c.Set(&tmp3)
			p.commitCoefficient(c)
		}
		if c.Sign() != 0 {
			newCoefficientCount = i + 1
		}
	}
	if newCoefficientCount > 0 {
		p.commitCoefficient(p.getCoefficient(newCoefficientCount - 1))
	}
	p.setCoefficientCount(newCoefficientCount)
}

// Sets p to p^N mod (N, X^R - 1), where R is the size of p. tmp1 and
// tmp2 must not alias each other or p.
func (p *BigIntPoly) Pow(N big.Int, tmp1, tmp2 *BigIntPoly) {
	tmp1.phi.Set(&p.phi)

	for i := N.BitLen() - 2; i >= 0; i-- {
		tmp1.mul(tmp1, N, tmp2)
		if N.Bit(i) != 0 {
			tmp1.mul(p, N, tmp2)
		}
	}

	p.phi, tmp1.phi = tmp1.phi, p.phi
}

// fmt.Formatter implementation.
func (p *BigIntPoly) Format(f fmt.State, c rune) {
	if p.phi.Sign() == 0 {
		fmt.Fprint(f, "0")
		return
	}

	// Formats coeff*x^deg.
	formatNonZeroMonomial := func(
		f fmt.State, c rune,
		coeff big.Int, deg int) {
		if coeff.Cmp(big.NewInt(1)) != 0 || deg == 0 {
			fmt.Fprint(f, &coeff)
		}
		if deg != 0 {
			fmt.Fprint(f, "x")
			if deg > 1 {
				fmt.Fprint(f, "^", deg)
			}
		}
	}

	i := p.getCoefficientCount() - 1
	formatNonZeroMonomial(f, c, p.getCoefficient(i), i)

	for i--; i >= 0; i-- {
		coeff := p.getCoefficient(i)
		if coeff.Sign() != 0 {
			fmt.Fprint(f, " + ")
			formatNonZeroMonomial(f, c, coeff, i)
		}
	}
}
<commit_msg>Avoid calling BigIntPoly.getCoefficientCount() repeatedly<commit_after>package main

import "fmt"
import "math/big"

// A BigIntPoly represents a polynomial with big.Int coefficients mod
// some (N, X^R - 1).
//
// The zero value for a BigIntPoly represents the zero polynomial.
type BigIntPoly struct {
	R int
	// k is the number of big.Words required to hold a coefficient
	// in calculations without overflowing.
	k int
	// If p(x) is the polynomial as a function, phi is
	// p(2^{k*bitsize(big.Word)}). Since a coefficient fits into k
	// big.Words, this is a lossless transformation; that is, one
	// can recover all coefficients of p(x) from phi.
	//
	// phi is set to have capacity for the largest possible
	// (intermediate) polynomial. No assumptions can be made about
	// the bytes in the unused capacity except for that the unused
	// bytes for the leading coefficient (if any) is guaranteed to
	// be zeroed out.
	phi big.Int
}

// Only polynomials built with the same value of N and R may be used
// together in one of the functions below.

// Builds a new BigIntPoly representing the zero polynomial
// mod (N, X^R - 1). R must fit into an int.
func NewBigIntPoly(N, R big.Int) *BigIntPoly {
	// A coefficient can be up to R*(N - 1)^2 in intermediate
	// calculations.
	var maxCoefficient big.Int
	maxCoefficient.Sub(&N, big.NewInt(1))
	maxCoefficient.Mul(&maxCoefficient, &maxCoefficient)
	maxCoefficient.Mul(&maxCoefficient, &R)

	var phi big.Int
	rInt := int(R.Int64())
	k := len(maxCoefficient.Bits())
	// Up to 2*R coefficients may be needed in intermediate
	// calculations.
	maxWordCount := 2 * rInt * k
	phi.SetBits(make([]big.Word, maxWordCount))
	return &BigIntPoly{rInt, k, phi}
}

// Returns 1 + the degree of this polynomial, or 0 if the polynomial
// is the zero polynomial.
func (p *BigIntPoly) getCoefficientCount() int {
	l := len(p.phi.Bits())
	if l == 0 {
		return 0
	}
	coefficientCount := l / p.k
	if l%p.k != 0 {
		coefficientCount++
	}
	return coefficientCount
}

// Sets the coefficient count to the given number, which must be at
// most p.R. The unused bytes of the leading coefficient must be
// cleared (via commitCoefficient()) prior to this being called.
func (p *BigIntPoly) setCoefficientCount(coefficientCount int) {
	p.phi.SetBits(p.phi.Bits()[0 : coefficientCount*p.k])
}

// Returns the ith coefficient of this polynomial. i must be less than
// p.getCoefficientCount().
func (p *BigIntPoly) getCoefficient(i int) big.Int {
	var c big.Int
	start := i * p.k
	end := (i + 1) * p.k
	// Since the unused data for the leading coefficient is
	// guaranteed to be zeroed out, this is okay.
	c.SetBits(p.phi.Bits()[start:end])
	return c
}

// Clears the unused bytes of the given coefficient. Must be called
// after all changes have been made to a coefficient via a big.Int
// returned from p.getCoefficient(). Also must be called on the
// leading coefficient before p.setCoefficientCount() is called.
func (p *BigIntPoly) commitCoefficient(c big.Int) {
	cBits := c.Bits()
	unusedBits := cBits[len(cBits):p.k]
	for j := 0; j < len(unusedBits); j++ {
		unusedBits[j] = 0
	}
}

// Sets p to X^k + a mod (N, X^R - 1).
func (p *BigIntPoly) Set(a, k, N big.Int) {
	c0 := p.getCoefficient(0)
	c0.Mod(&a, &N)
	p.commitCoefficient(c0)

	R := big.NewInt(int64(p.R))
	var kModRBig big.Int
	kModRBig.Mod(&k, R)
	kModR := int(kModRBig.Int64())

	for i := 1; i <= kModR; i++ {
		c := p.getCoefficient(i)
		c.Set(&big.Int{})
		p.commitCoefficient(c)
	}

	cKModR := p.getCoefficient(kModR)
	cKModR.Set(big.NewInt(1))
	p.commitCoefficient(cKModR)

	p.setCoefficientCount(kModR + 1)
}

// Returns whether p has the same coefficients as q.
func (p *BigIntPoly) Eq(q *BigIntPoly) bool {
	return p.phi.Cmp(&q.phi) == 0
}

// Sets p to the product of p and q mod (N, X^R - 1). Assumes R >=
// 2. tmp must not alias p or q.
func (p *BigIntPoly) mul(q *BigIntPoly, N big.Int, tmp *BigIntPoly) {
	tmp.phi.Mul(&p.phi, &q.phi)
	p.phi, tmp.phi = tmp.phi, p.phi

	// Mod p by X^R - 1.
	mid := p.R * p.k
	pBits := p.phi.Bits()
	if len(pBits) > mid {
		var lo, hi big.Int
		lo.SetBits(pBits[:mid])
		hi.SetBits(pBits[mid:])
		p.phi.Add(&lo, &hi)
		pBits = p.phi.Bits()
	}

	// Clear the unused bits of the leading coefficient if
	// necessary.
	if len(pBits)%p.k != 0 {
		start := len(pBits)
		end := start + p.k - start%p.k
		unusedBits := pBits[start:end]
		for i := 0; i < len(unusedBits); i++ {
			unusedBits[i] = 0
		}
	}
	// Commit the leading coefficient before we access it.
	oldCoefficientCount := p.getCoefficientCount()
	if oldCoefficientCount > 0 {
		p.commitCoefficient(p.getCoefficient(oldCoefficientCount - 1))
	}

	// Mod p by N.
	newCoefficientCount := 0
	tmp2 := tmp.getCoefficient(0)
	tmp3 := tmp.getCoefficient(1)
	for i := 0; i < oldCoefficientCount; i++ {
		c := p.getCoefficient(i)
		if c.Cmp(&N) >= 0 {
			// Mod c by N. Use big.Int.QuoRem() instead of
			// big.Int.Mod() since the latter allocates an
			// extra big.Int.
			tmp2.QuoRem(&c, &N, &tmp3)
			c.Set(&tmp3)
			p.commitCoefficient(c)
		}
		if c.Sign() != 0 {
			newCoefficientCount = i + 1
		}
	}
	p.setCoefficientCount(newCoefficientCount)
}

// Sets p to p^N mod (N, X^R - 1), where R is the size of p. tmp1 and
// tmp2 must not alias each other or p.
func (p *BigIntPoly) Pow(N big.Int, tmp1, tmp2 *BigIntPoly) {
	tmp1.phi.Set(&p.phi)

	for i := N.BitLen() - 2; i >= 0; i-- {
		tmp1.mul(tmp1, N, tmp2)
		if N.Bit(i) != 0 {
			tmp1.mul(p, N, tmp2)
		}
	}

	p.phi, tmp1.phi = tmp1.phi, p.phi
}

// fmt.Formatter implementation.
func (p *BigIntPoly) Format(f fmt.State, c rune) {
	if p.phi.Sign() == 0 {
		fmt.Fprint(f, "0")
		return
	}

	// Formats coeff*x^deg.
	formatNonZeroMonomial := func(
		f fmt.State, c rune,
		coeff big.Int, deg int) {
		if coeff.Cmp(big.NewInt(1)) != 0 || deg == 0 {
			fmt.Fprint(f, &coeff)
		}
		if deg != 0 {
			fmt.Fprint(f, "x")
			if deg > 1 {
				fmt.Fprint(f, "^", deg)
			}
		}
	}

	i := p.getCoefficientCount() - 1
	formatNonZeroMonomial(f, c, p.getCoefficient(i), i)

	for i--; i >= 0; i-- {
		coeff := p.getCoefficient(i)
		if coeff.Sign() != 0 {
			fmt.Fprint(f, " + ")
			formatNonZeroMonomial(f, c, coeff, i)
		}
	}
}
<|endoftext|>


<commit_before>/*
Copyright 2018 The Knative Authors
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
*/

// Multitenant autoscaler executable.
package main

import (
	"flag"
	"fmt"
	"log"
	"time"

	"github.com/knative/pkg/configmap"
	"github.com/knative/pkg/controller"
	"github.com/knative/pkg/injection"
	"github.com/knative/pkg/injection/clients/kubeclient"
	endpointsinformer "github.com/knative/pkg/injection/informers/kubeinformers/corev1/endpoints"
	"github.com/knative/pkg/logging"
	"github.com/knative/pkg/metrics"
	pkgmetrics "github.com/knative/pkg/metrics"
	"github.com/knative/pkg/signals"
	"github.com/knative/pkg/system"
	"github.com/knative/serving/pkg/apis/serving"
	"github.com/knative/serving/pkg/autoscaler"
	"github.com/knative/serving/pkg/autoscaler/statserver"
	"github.com/knative/serving/pkg/reconciler/autoscaling/hpa"
	"github.com/knative/serving/pkg/reconciler/autoscaling/kpa"
	"github.com/knative/serving/pkg/resources"
	basecmd "github.com/kubernetes-incubator/custom-metrics-apiserver/pkg/cmd"
	"github.com/spf13/pflag"
	"go.uber.org/zap"
	"golang.org/x/sync/errgroup"
	corev1informers "k8s.io/client-go/informers/core/v1"
	corev1listers "k8s.io/client-go/listers/core/v1"
	"k8s.io/client-go/rest"
	"k8s.io/client-go/tools/clientcmd"
)

const (
	statsServerAddr = ":8080"
	statsBufferLen  = 1000
	component       = "autoscaler"
)

var (
	masterURL  = flag.String("master", "", "The address of the Kubernetes API server. Overrides any value in kubeconfig. Only required if out-of-cluster.")
	kubeconfig = flag.String("kubeconfig", "", "Path to a kubeconfig. Only required if out-of-cluster.")
)

func main() {
	// Initialize early to get access to flags and merge them with the autoscaler flags.
	customMetricsAdapter := &basecmd.AdapterBase{}
	customMetricsAdapter.Flags().AddGoFlagSet(flag.CommandLine)
	pflag.Parse()

	// Set up signals so we handle the first shutdown signal gracefully.
	ctx := signals.NewContext()

	// Set up our logger.
	loggingConfigMap, err := configmap.Load("/etc/config-logging")
	if err != nil {
		log.Fatal("Error loading logging configuration:", err)
	}
	loggingConfig, err := logging.NewConfigFromMap(loggingConfigMap)
	if err != nil {
		log.Fatal("Error parsing logging configuration:", err)
	}
	logger, atomicLevel := logging.NewLoggerFromConfig(loggingConfig, component)
	defer flush(logger)
	ctx = logging.WithLogger(ctx, logger)

	cfg, err := clientcmd.BuildConfigFromFlags(*masterURL, *kubeconfig)
	if err != nil {
		logger.Fatalw("Error building kubeconfig", zap.Error(err))
	}

	logger.Infof("Registering %d clients", len(injection.Default.GetClients()))
	logger.Infof("Registering %d informer factories", len(injection.Default.GetInformerFactories()))
	logger.Infof("Registering %d informers", len(injection.Default.GetInformers()))
	logger.Infof("Registering %d controllers", 2)

	// Adjust our client's rate limits based on the number of controller's we are running.
	cfg.QPS = 2 * rest.DefaultQPS
	cfg.Burst = 2 * rest.DefaultBurst

	ctx, informers := injection.Default.SetupInformers(ctx, cfg)

	// statsCh is the main communication channel between the stats channel and multiscaler.
	statsCh := make(chan *autoscaler.StatMessage, statsBufferLen)
	defer close(statsCh)

	cmw := configmap.NewInformedWatcher(kubeclient.Get(ctx), system.Namespace())
	// Watch the logging config map and dynamically update logging levels.
	cmw.Watch(logging.ConfigMapName(), logging.UpdateLevelFromConfigMap(logger, atomicLevel, component))
	// Watch the observability config map and dynamically update metrics exporter.
	cmw.Watch(metrics.ConfigMapName(), metrics.UpdateExporterFromConfigMap(component, logger))

	endpointsInformer := endpointsinformer.Get(ctx)

	collector := autoscaler.NewMetricCollector(statsScraperFactoryFunc(endpointsInformer.Lister()), logger)
	customMetricsAdapter.WithCustomMetrics(autoscaler.NewMetricProvider())

	// Set up scalers.
	// uniScalerFactory depends endpointsInformer to be set.
	multiScaler := autoscaler.NewMultiScaler(ctx.Done(), uniScalerFactoryFunc(endpointsInformer, collector), logger)

	controllers := []*controller.Impl{
		kpa.NewController(ctx, cmw, multiScaler, collector),
		hpa.NewController(ctx, cmw),
	}

	// Set up a statserver.
	statsServer := statserver.New(statsServerAddr, statsCh, logger)
	defer statsServer.Shutdown(time.Second * 5)

	// Start watching the configs.
	if err := cmw.Start(ctx.Done()); err != nil {
		logger.Fatalw("Failed to start watching configs", zap.Error(err))
	}

	// Start all of the informers and wait for them to sync.
	if err := controller.StartInformers(ctx.Done(), informers...); err != nil {
		logger.Fatalw("Failed to start informers", err)
	}

	go controller.StartAll(ctx.Done(), controllers...)

	// Run the controllers and the statserver in a group.
	var eg errgroup.Group
	eg.Go(func() error {
		return customMetricsAdapter.Run(ctx.Done())
	})

	eg.Go(func() error {
		return statsServer.ListenAndServe()
	})

	go func() {
		for {
			sm, ok := <-statsCh
			if !ok {
				break
			}
			collector.Record(sm.Key, sm.Stat)
			multiScaler.Poke(sm.Key, sm.Stat)
		}
	}()

	egCh := make(chan struct{})

	go func() {
		if err := eg.Wait(); err != nil {
			logger.Errorw("Group error.", zap.Error(err))
		}
		close(egCh)
	}()

	select {
	case <-egCh:
	case <-ctx.Done():
	}
}

func setupLogger() (*zap.SugaredLogger, zap.AtomicLevel) {
	loggingConfigMap, err := configmap.Load("/etc/config-logging")
	if err != nil {
		log.Fatal("Error loading logging configuration:", err)
	}
	loggingConfig, err := logging.NewConfigFromMap(loggingConfigMap)
	if err != nil {
		log.Fatal("Error parsing logging configuration:", err)
	}
	return logging.NewLoggerFromConfig(loggingConfig, component)
}

func uniScalerFactoryFunc(endpointsInformer corev1informers.EndpointsInformer, metricClient autoscaler.MetricClient) func(decider *autoscaler.Decider) (autoscaler.UniScaler, error) {
	return func(decider *autoscaler.Decider) (autoscaler.UniScaler, error) {
		if v, ok := decider.Labels[serving.ConfigurationLabelKey]; !ok || v == "" {
			return nil, fmt.Errorf("label %q not found or empty in Decider %s", serving.ConfigurationLabelKey, decider.Name)
		}
		if decider.Spec.ServiceName == "" {
			return nil, fmt.Errorf("%s decider has empty ServiceName", decider.Name)
		}

		serviceName := decider.Labels[serving.ServiceLabelKey] // This can be empty.
		configName := decider.Labels[serving.ConfigurationLabelKey]

		// Create a stats reporter which tags statistics by PA namespace, configuration name, and PA name.
		reporter, err := autoscaler.NewStatsReporter(decider.Namespace, serviceName, configName, decider.Name)
		if err != nil {
			return nil, err
		}

		podCounter := resources.NewScopedEndpointsCounter(endpointsInformer.Lister(), decider.Namespace, decider.Name)
		return autoscaler.New(decider.Namespace, decider.Name, metricClient, podCounter, decider.Spec, reporter)
	}
}

func statsScraperFactoryFunc(endpointsLister corev1listers.EndpointsLister) func(metric *autoscaler.Metric) (autoscaler.StatsScraper, error) {
	return func(metric *autoscaler.Metric) (autoscaler.StatsScraper, error) {
		podCounter := resources.NewScopedEndpointsCounter(endpointsLister, metric.Namespace, metric.Name)
		return autoscaler.NewServiceScraper(metric, podCounter)
	}
}

func flush(logger *zap.SugaredLogger) {
	logger.Sync()
	pkgmetrics.FlushExporter()
}
<commit_msg>Fix the pod counting target. (#4296)<commit_after>/*
Copyright 2018 The Knative Authors
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
*/

// Multitenant autoscaler executable.
package main

import (
	"flag"
	"fmt"
	"log"
	"time"

	"github.com/knative/pkg/configmap"
	"github.com/knative/pkg/controller"
	"github.com/knative/pkg/injection"
	"github.com/knative/pkg/injection/clients/kubeclient"
	endpointsinformer "github.com/knative/pkg/injection/informers/kubeinformers/corev1/endpoints"
	"github.com/knative/pkg/logging"
	"github.com/knative/pkg/metrics"
	pkgmetrics "github.com/knative/pkg/metrics"
	"github.com/knative/pkg/signals"
	"github.com/knative/pkg/system"
	"github.com/knative/serving/pkg/apis/serving"
	"github.com/knative/serving/pkg/autoscaler"
	"github.com/knative/serving/pkg/autoscaler/statserver"
	"github.com/knative/serving/pkg/reconciler/autoscaling/hpa"
	"github.com/knative/serving/pkg/reconciler/autoscaling/kpa"
	"github.com/knative/serving/pkg/resources"
	basecmd "github.com/kubernetes-incubator/custom-metrics-apiserver/pkg/cmd"
	"github.com/spf13/pflag"
	"go.uber.org/zap"
	"golang.org/x/sync/errgroup"
	corev1informers "k8s.io/client-go/informers/core/v1"
	corev1listers "k8s.io/client-go/listers/core/v1"
	"k8s.io/client-go/rest"
	"k8s.io/client-go/tools/clientcmd"
)

const (
	statsServerAddr = ":8080"
	statsBufferLen  = 1000
	component       = "autoscaler"
)

var (
	masterURL  = flag.String("master", "", "The address of the Kubernetes API server. Overrides any value in kubeconfig. Only required if out-of-cluster.")
	kubeconfig = flag.String("kubeconfig", "", "Path to a kubeconfig. Only required if out-of-cluster.")
)

func main() {
	// Initialize early to get access to flags and merge them with the autoscaler flags.
	customMetricsAdapter := &basecmd.AdapterBase{}
	customMetricsAdapter.Flags().AddGoFlagSet(flag.CommandLine)
	pflag.Parse()

	// Set up signals so we handle the first shutdown signal gracefully.
	ctx := signals.NewContext()

	// Set up our logger.
	loggingConfigMap, err := configmap.Load("/etc/config-logging")
	if err != nil {
		log.Fatal("Error loading logging configuration:", err)
	}
	loggingConfig, err := logging.NewConfigFromMap(loggingConfigMap)
	if err != nil {
		log.Fatal("Error parsing logging configuration:", err)
	}
	logger, atomicLevel := logging.NewLoggerFromConfig(loggingConfig, component)
	defer flush(logger)
	ctx = logging.WithLogger(ctx, logger)

	cfg, err := clientcmd.BuildConfigFromFlags(*masterURL, *kubeconfig)
	if err != nil {
		logger.Fatalw("Error building kubeconfig", zap.Error(err))
	}

	logger.Infof("Registering %d clients", len(injection.Default.GetClients()))
	logger.Infof("Registering %d informer factories", len(injection.Default.GetInformerFactories()))
	logger.Infof("Registering %d informers", len(injection.Default.GetInformers()))
	logger.Infof("Registering %d controllers", 2)

	// Adjust our client's rate limits based on the number of controller's we are running.
	cfg.QPS = 2 * rest.DefaultQPS
	cfg.Burst = 2 * rest.DefaultBurst

	ctx, informers := injection.Default.SetupInformers(ctx, cfg)

	// statsCh is the main communication channel between the stats channel and multiscaler.
	statsCh := make(chan *autoscaler.StatMessage, statsBufferLen)
	defer close(statsCh)

	cmw := configmap.NewInformedWatcher(kubeclient.Get(ctx), system.Namespace())
	// Watch the logging config map and dynamically update logging levels.
	cmw.Watch(logging.ConfigMapName(), logging.UpdateLevelFromConfigMap(logger, atomicLevel, component))
	// Watch the observability config map and dynamically update metrics exporter.
	cmw.Watch(metrics.ConfigMapName(), metrics.UpdateExporterFromConfigMap(component, logger))

	endpointsInformer := endpointsinformer.Get(ctx)

	collector := autoscaler.NewMetricCollector(statsScraperFactoryFunc(endpointsInformer.Lister()), logger)
	customMetricsAdapter.WithCustomMetrics(autoscaler.NewMetricProvider())

	// Set up scalers.
	// uniScalerFactory depends endpointsInformer to be set.
	multiScaler := autoscaler.NewMultiScaler(ctx.Done(), uniScalerFactoryFunc(endpointsInformer, collector), logger)

	controllers := []*controller.Impl{
		kpa.NewController(ctx, cmw, multiScaler, collector),
		hpa.NewController(ctx, cmw),
	}

	// Set up a statserver.
	statsServer := statserver.New(statsServerAddr, statsCh, logger)
	defer statsServer.Shutdown(time.Second * 5)

	// Start watching the configs.
	if err := cmw.Start(ctx.Done()); err != nil {
		logger.Fatalw("Failed to start watching configs", zap.Error(err))
	}

	// Start all of the informers and wait for them to sync.
	if err := controller.StartInformers(ctx.Done(), informers...); err != nil {
		logger.Fatalw("Failed to start informers", err)
	}

	go controller.StartAll(ctx.Done(), controllers...)

	// Run the controllers and the statserver in a group.
	var eg errgroup.Group
	eg.Go(func() error {
		return customMetricsAdapter.Run(ctx.Done())
	})

	eg.Go(func() error {
		return statsServer.ListenAndServe()
	})

	go func() {
		for {
			sm, ok := <-statsCh
			if !ok {
				break
			}
			collector.Record(sm.Key, sm.Stat)
			multiScaler.Poke(sm.Key, sm.Stat)
		}
	}()

	egCh := make(chan struct{})

	go func() {
		if err := eg.Wait(); err != nil {
			logger.Errorw("Group error.", zap.Error(err))
		}
		close(egCh)
	}()

	select {
	case <-egCh:
	case <-ctx.Done():
	}
}

func setupLogger() (*zap.SugaredLogger, zap.AtomicLevel) {
	loggingConfigMap, err := configmap.Load("/etc/config-logging")
	if err != nil {
		log.Fatal("Error loading logging configuration:", err)
	}
	loggingConfig, err := logging.NewConfigFromMap(loggingConfigMap)
	if err != nil {
		log.Fatal("Error parsing logging configuration:", err)
	}
	return logging.NewLoggerFromConfig(loggingConfig, component)
}

func uniScalerFactoryFunc(endpointsInformer corev1informers.EndpointsInformer, metricClient autoscaler.MetricClient) func(decider *autoscaler.Decider) (autoscaler.UniScaler, error) {
	return func(decider *autoscaler.Decider) (autoscaler.UniScaler, error) {
		if v, ok := decider.Labels[serving.ConfigurationLabelKey]; !ok || v == "" {
			return nil, fmt.Errorf("label %q not found or empty in Decider %s", serving.ConfigurationLabelKey, decider.Name)
		}
		if decider.Spec.ServiceName == "" {
			return nil, fmt.Errorf("%s decider has empty ServiceName", decider.Name)
		}

		serviceName := decider.Labels[serving.ServiceLabelKey] // This can be empty.
		configName := decider.Labels[serving.ConfigurationLabelKey]

		// Create a stats reporter which tags statistics by PA namespace, configuration name, and PA name.
		reporter, err := autoscaler.NewStatsReporter(decider.Namespace, serviceName, configName, decider.Name)
		if err != nil {
			return nil, err
		}

		podCounter := resources.NewScopedEndpointsCounter(endpointsInformer.Lister(), decider.Namespace, decider.Name)
		return autoscaler.New(decider.Namespace, decider.Name, metricClient, podCounter, decider.Spec, reporter)
	}
}

func statsScraperFactoryFunc(endpointsLister corev1listers.EndpointsLister) func(metric *autoscaler.Metric) (autoscaler.StatsScraper, error) {
	return func(metric *autoscaler.Metric) (autoscaler.StatsScraper, error) {
		podCounter := resources.NewScopedEndpointsCounter(endpointsLister, metric.Namespace, metric.Spec.ScrapeTarget)
		return autoscaler.NewServiceScraper(metric, podCounter)
	}
}

func flush(logger *zap.SugaredLogger) {
	logger.Sync()
	pkgmetrics.FlushExporter()
}
<|endoftext|>


<commit_before>package xxhash_test

import (
	"hash/adler32"
	"hash/crc32"
	"hash/fnv"
	"testing"

	C "github.com/OneOfOne/xxhash"
	N "github.com/OneOfOne/xxhash/native"
)

var (
	in = []byte(`Lorem ipsum dolor sit amet, consectetur adipisicing elit, sed do eiusmod tempor incididunt ut labore et dolore magna aliqua. Ut enim ad minim veniam, quis nostrud exercitation ullamco laboris nisi ut aliquip ex ea commodo consequat. Duis aute irure dolor in reprehenderit in voluptate velit esse cillum dolore eu fugiat nulla pariatur. Excepteur sint occaecat cupidatat non proident, sunt in culpa qui officia deserunt mollit anim id est laborum.
Lorem ipsum dolor sit amet, consectetur adipisicing elit, sed do eiusmod tempor incididunt ut labore et dolore magna aliqua. Ut enim ad minim veniam, quis nostrud exercitation ullamco laboris nisi ut aliquip ex ea commodo consequat. Duis aute irure dolor in reprehenderit in voluptate velit esse cillum dolore eu fugiat nulla pariatur. Excepteur sint occaecat cupidatat non proident, sunt in culpa qui officia deserunt mollit anim id est laborum.
Lorem ipsum dolor sit amet, consectetur adipisicing elit, sed do eiusmod tempor incididunt ut labore et dolore magna aliqua. Ut enim ad minim veniam, quis nostrud exercitation ullamco laboris nisi ut aliquip ex ea commodo consequat. Duis aute irure dolor in reprehenderit in voluptate velit esse cillum dolore eu fugiat nulla pariatur. Excepteur sint occaecat cupidatat non proident, sunt in culpa qui officia deserunt mollit anim id est laborum.
Lorem ipsum dolor sit amet, consectetur adipisicing elit, sed do eiusmod tempor incididunt ut labore et dolore magna aliqua. Ut enim ad minim veniam, quis nostrud exercitation ullamco laboris nisi ut aliquip ex ea commodo consequat. Duis aute irure dolor in reprehenderit in voluptate velit esse cillum dolore eu fugiat nulla pariatur. Excepteur sint occaecat cupidatat non proident, sunt in culpa qui officia deserunt mollit anim id est laborum.
Lorem ipsum dolor sit amet, consectetur adipisicing elit, sed do eiusmod tempor incididunt ut labore et dolore magna aliqua. Ut enim ad minim veniam, quis nostrud exercitation ullamco laboris nisi ut aliquip ex ea commodo consequat. Duis aute irure dolor in reprehenderit in voluptate velit esse cillum dolore eu fugiat nulla pariatur. Excepteur sint occaecat cupidatat non proident, sunt in culpa qui officia deserunt mollit anim id est laborum.
`)
)

const (
	expected32 uint32 = 0x6101218F
	expected64 uint64 = 0xFFAE31BEBFED7652
)

func Test(t *testing.T) {
	t.Logf("CGO version's backend: %s", C.Backend)
	t.Logf("Native version's backend: %s", N.Backend)
}

func TestHash32(t *testing.T) {
	h := N.New32()
	h.Write(in)
	r := h.Sum32()
	if r != expected32 {
		t.Errorf("expected 0x%x, got 0x%x.", expected32, r)
	}
}

func TestHash32Cgo(t *testing.T) {
	h := C.New32()
	h.Write(in)
	r := h.Sum32()
	if r != expected32 {
		t.Errorf("expected 0x%x, got 0x%x.", expected32, r)
	}
}

func TestHash32Short(t *testing.T) {
	r := N.Checksum32(in)
	if r != expected32 {
		t.Errorf("expected 0x%x, got 0x%x.", expected32, r)
	}
}

func TestHash32CgoShort(t *testing.T) {
	r := C.Checksum32(in)
	if r != expected32 {
		t.Errorf("expected 0x%x, got 0x%x.", expected32, r)
	}
}

func TestHash64(t *testing.T) {
	h := N.New64()
	h.Write(in)
	r := h.Sum64()
	if r != expected64 {
		t.Errorf("expected 0x%x, got 0x%x.", expected64, r)
	}
}

func TestHash64Cgo(t *testing.T) {
	h := C.New64()
	h.Write(in)
	r := h.Sum64()
	if r != expected64 {
		t.Errorf("expected 0x%x, got 0x%x.", expected64, r)
	}
}

func TestHash64Short(t *testing.T) {
	r := N.Checksum64(in)
	if r != expected64 {
		t.Errorf("expected 0x%x, got 0x%x.", expected64, r)
	}
}

func TestHash64CgoShort(t *testing.T) {
	r := C.Checksum64(in)
	if r != expected64 {
		t.Errorf("expected 0x%x, got 0x%x.", expected64, r)
	}
}

func BenchmarkXxhash32(b *testing.B) {
	for i := 0; i < b.N; i++ {
		N.Checksum32(in)
	}
}

func BenchmarkXxhash32Cgo(b *testing.B) {
	for i := 0; i < b.N; i++ {
		C.Checksum32(in)
	}
}

func BenchmarkXxhash64(b *testing.B) {
	for i := 0; i < b.N; i++ {
		N.Checksum64(in)
	}
}

func BenchmarkXxhash64Cgo(b *testing.B) {
	for i := 0; i < b.N; i++ {
		C.Checksum64(in)
	}
}

func BenchmarkFnv32(b *testing.B) {
	h := fnv.New32()
	for i := 0; i < b.N; i++ {
		h.Sum(in)
	}
}

func BenchmarkFnv64(b *testing.B) {
	h := fnv.New64()
	for i := 0; i < b.N; i++ {
		h.Sum(in)
	}
}

func BenchmarkAdler32(b *testing.B) {
	for i := 0; i < b.N; i++ {
		adler32.Checksum(in)
	}
}

func BenchmarkCRC32IEEE(b *testing.B) {
	for i := 0; i < b.N; i++ {
		crc32.ChecksumIEEE(in)
	}
}

func BenchmarkXxhash64VeryShort(b *testing.B) {
	k := []byte("Test-key-1")
	for i := 0; i < b.N; i++ {
		N.Checksum64(k)
	}
}

func BenchmarkXxhash64CgoVeryShort(b *testing.B) {
	k := []byte("Test-key-100")
	for i := 0; i < b.N; i++ {
		C.Checksum64(k)
	}
}

func BenchmarkXxhash64MultiWrites(b *testing.B) {
	h := N.New64()
	for i := 0; i < b.N; i++ {
		h.Write(in)
	}
	_ = h.Sum64()
}

func BenchmarkXxhash64CgoMultiWrites(b *testing.B) {
	h := C.New64()
	for i := 0; i < b.N; i++ {
		h.Write(in)
	}
	_ = h.Sum64()
}
<commit_msg>Fix a typo in BenchmarkXxhash64CgoVeryShort<commit_after>package xxhash_test

import (
	"hash/adler32"
	"hash/crc32"
	"hash/fnv"
	"testing"

	C "github.com/OneOfOne/xxhash"
	N "github.com/OneOfOne/xxhash/native"
)

var (
	in = []byte(`Lorem ipsum dolor sit amet, consectetur adipisicing elit, sed do eiusmod tempor incididunt ut labore et dolore magna aliqua. Ut enim ad minim veniam, quis nostrud exercitation ullamco laboris nisi ut aliquip ex ea commodo consequat. Duis aute irure dolor in reprehenderit in voluptate velit esse cillum dolore eu fugiat nulla pariatur. Excepteur sint occaecat cupidatat non proident, sunt in culpa qui officia deserunt mollit anim id est laborum.
Lorem ipsum dolor sit amet, consectetur adipisicing elit, sed do eiusmod tempor incididunt ut labore et dolore magna aliqua. Ut enim ad minim veniam, quis nostrud exercitation ullamco laboris nisi ut aliquip ex ea commodo consequat. Duis aute irure dolor in reprehenderit in voluptate velit esse cillum dolore eu fugiat nulla pariatur. Excepteur sint occaecat cupidatat non proident, sunt in culpa qui officia deserunt mollit anim id est laborum.
Lorem ipsum dolor sit amet, consectetur adipisicing elit, sed do eiusmod tempor incididunt ut labore et dolore magna aliqua. Ut enim ad minim veniam, quis nostrud exercitation ullamco laboris nisi ut aliquip ex ea commodo consequat. Duis aute irure dolor in reprehenderit in voluptate velit esse cillum dolore eu fugiat nulla pariatur. Excepteur sint occaecat cupidatat non proident, sunt in culpa qui officia deserunt mollit anim id est laborum.
Lorem ipsum dolor sit amet, consectetur adipisicing elit, sed do eiusmod tempor incididunt ut labore et dolore magna aliqua. Ut enim ad minim veniam, quis nostrud exercitation ullamco laboris nisi ut aliquip ex ea commodo consequat. Duis aute irure dolor in reprehenderit in voluptate velit esse cillum dolore eu fugiat nulla pariatur. Excepteur sint occaecat cupidatat non proident, sunt in culpa qui officia deserunt mollit anim id est laborum.
Lorem ipsum dolor sit amet, consectetur adipisicing elit, sed do eiusmod tempor incididunt ut labore et dolore magna aliqua. Ut enim ad minim veniam, quis nostrud exercitation ullamco laboris nisi ut aliquip ex ea commodo consequat. Duis aute irure dolor in reprehenderit in voluptate velit esse cillum dolore eu fugiat nulla pariatur. Excepteur sint occaecat cupidatat non proident, sunt in culpa qui officia deserunt mollit anim id est laborum.
`)
)

const (
	expected32 uint32 = 0x6101218F
	expected64 uint64 = 0xFFAE31BEBFED7652
)

func Test(t *testing.T) {
	t.Logf("CGO version's backend: %s", C.Backend)
	t.Logf("Native version's backend: %s", N.Backend)
}

func TestHash32(t *testing.T) {
	h := N.New32()
	h.Write(in)
	r := h.Sum32()
	if r != expected32 {
		t.Errorf("expected 0x%x, got 0x%x.", expected32, r)
	}
}

func TestHash32Cgo(t *testing.T) {
	h := C.New32()
	h.Write(in)
	r := h.Sum32()
	if r != expected32 {
		t.Errorf("expected 0x%x, got 0x%x.", expected32, r)
	}
}

func TestHash32Short(t *testing.T) {
	r := N.Checksum32(in)
	if r != expected32 {
		t.Errorf("expected 0x%x, got 0x%x.", expected32, r)
	}
}

func TestHash32CgoShort(t *testing.T) {
	r := C.Checksum32(in)
	if r != expected32 {
		t.Errorf("expected 0x%x, got 0x%x.", expected32, r)
	}
}

func TestHash64(t *testing.T) {
	h := N.New64()
	h.Write(in)
	r := h.Sum64()
	if r != expected64 {
		t.Errorf("expected 0x%x, got 0x%x.", expected64, r)
	}
}

func TestHash64Cgo(t *testing.T) {
	h := C.New64()
	h.Write(in)
	r := h.Sum64()
	if r != expected64 {
		t.Errorf("expected 0x%x, got 0x%x.", expected64, r)
	}
}

func TestHash64Short(t *testing.T) {
	r := N.Checksum64(in)
	if r != expected64 {
		t.Errorf("expected 0x%x, got 0x%x.", expected64, r)
	}
}

func TestHash64CgoShort(t *testing.T) {
	r := C.Checksum64(in)
	if r != expected64 {
		t.Errorf("expected 0x%x, got 0x%x.", expected64, r)
	}
}

func BenchmarkXxhash32(b *testing.B) {
	for i := 0; i < b.N; i++ {
		N.Checksum32(in)
	}
}

func BenchmarkXxhash32Cgo(b *testing.B) {
	for i := 0; i < b.N; i++ {
		C.Checksum32(in)
	}
}

func BenchmarkXxhash64(b *testing.B) {
	for i := 0; i < b.N; i++ {
		N.Checksum64(in)
	}
}

func BenchmarkXxhash64Cgo(b *testing.B) {
	for i := 0; i < b.N; i++ {
		C.Checksum64(in)
	}
}

func BenchmarkFnv32(b *testing.B) {
	h := fnv.New32()
	for i := 0; i < b.N; i++ {
		h.Sum(in)
	}
}

func BenchmarkFnv64(b *testing.B) {
	h := fnv.New64()
	for i := 0; i < b.N; i++ {
		h.Sum(in)
	}
}

func BenchmarkAdler32(b *testing.B) {
	for i := 0; i < b.N; i++ {
		adler32.Checksum(in)
	}
}

func BenchmarkCRC32IEEE(b *testing.B) {
	for i := 0; i < b.N; i++ {
		crc32.ChecksumIEEE(in)
	}
}

func BenchmarkXxhash64VeryShort(b *testing.B) {
	k := []byte("Test-key-100")
	for i := 0; i < b.N; i++ {
		N.Checksum64(k)
	}
}

func BenchmarkXxhash64CgoVeryShort(b *testing.B) {
	k := []byte("Test-key-100")
	for i := 0; i < b.N; i++ {
		C.Checksum64(k)
	}
}

func BenchmarkXxhash64MultiWrites(b *testing.B) {
	h := N.New64()
	for i := 0; i < b.N; i++ {
		h.Write(in)
	}
	_ = h.Sum64()
}

func BenchmarkXxhash64CgoMultiWrites(b *testing.B) {
	h := C.New64()
	for i := 0; i < b.N; i++ {
		h.Write(in)
	}
	_ = h.Sum64()
}
<|endoftext|>


<commit_before>// +build windows

/*
Copyright 2017 The Kubernetes Authors.

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
*/

package mount

import (
	"fmt"
	"os"
	"os/exec"
	"path/filepath"
	"strings"

	"k8s.io/klog/v2"
	"k8s.io/utils/keymutex"
)

const (
	accessDenied string = "access is denied"
)

// Mounter provides the default implementation of mount.Interface
// for the windows platform.  This implementation assumes that the
// kubelet is running in the host's root mount namespace.
type Mounter struct {
	mounterPath string
}

// New returns a mount.Interface for the current system.
// It provides options to override the default mounter behavior.
// mounterPath allows using an alternative to `/bin/mount` for mounting.
func New(mounterPath string) Interface {
	return &Mounter{
		mounterPath: mounterPath,
	}
}

// acquire lock for smb mount
var getSMBMountMutex = keymutex.NewHashed(0)

// Mount : mounts source to target with given options.
// currently only supports cifs(smb), bind mount(for disk)
func (mounter *Mounter) Mount(source string, target string, fstype string, options []string) error {
	return mounter.MountSensitive(source, target, fstype, options, nil /* sensitiveOptions */)
}

// MountSensitiveWithoutSystemd is the same as MountSensitive() but disable using ssytemd mount.
// Windows not supported systemd mount, this function degrades to MountSensitive().
func (mounter *Mounter) MountSensitiveWithoutSystemd(source string, target string, fstype string, options []string, sensitiveOptions []string) error {
	return mounter.MountSensitive(source, target, fstype, options, sensitiveOptions /* sensitiveOptions */)
}

// MountSensitive is the same as Mount() but this method allows
// sensitiveOptions to be passed in a separate parameter from the normal
// mount options and ensures the sensitiveOptions are never logged. This
// method should be used by callers that pass sensitive material (like
// passwords) as mount options.
func (mounter *Mounter) MountSensitive(source string, target string, fstype string, options []string, sensitiveOptions []string) error {
	target = NormalizeWindowsPath(target)
	sanitizedOptionsForLogging := sanitizedOptionsForLogging(options, sensitiveOptions)

	if source == "tmpfs" {
		klog.V(3).Infof("mounting source (%q), target (%q), with options (%q)", source, target, sanitizedOptionsForLogging)
		return os.MkdirAll(target, 0755)
	}

	parentDir := filepath.Dir(target)
	if err := os.MkdirAll(parentDir, 0755); err != nil {
		return err
	}

	klog.V(4).Infof("mount options(%q) source:%q, target:%q, fstype:%q, begin to mount",
		sanitizedOptionsForLogging, source, target, fstype)
	bindSource := source

	if bind, _, _, _ := MakeBindOptsSensitive(options, sensitiveOptions); bind {
		bindSource = NormalizeWindowsPath(source)
	} else {
		allOptions := []string{}
		allOptions = append(allOptions, options...)
		allOptions = append(allOptions, sensitiveOptions...)
		if len(allOptions) < 2 {
			return fmt.Errorf("mount options(%q) should have at least 2 options, current number:%d, source:%q, target:%q",
				sanitizedOptionsForLogging, len(allOptions), source, target)
		}

		// currently only cifs mount is supported
		if strings.ToLower(fstype) != "cifs" {
			return fmt.Errorf("only cifs mount is supported now, fstype: %q, mounting source (%q), target (%q), with options (%q)", fstype, source, target, sanitizedOptionsForLogging)
		}

		// lock smb mount for the same source
		getSMBMountMutex.LockKey(source)
		defer getSMBMountMutex.UnlockKey(source)

		username := allOptions[0]
		password := allOptions[1]
		if output, err := newSMBMapping(username, password, source); err != nil {
			klog.Warningf("SMB Mapping(%s) returned with error(%v), output(%s)", source, err, string(output))
			if isSMBMappingExist(source) {
				valid, err := isValidPath(source)
				if !valid {
					if err == nil || isAccessDeniedError(err) {
						klog.V(2).Infof("SMB Mapping(%s) already exists while it's not valid, return error: %v, now begin to remove and remount", source, err)
						if output, err = removeSMBMapping(source); err != nil {
							return fmt.Errorf("Remove-SmbGlobalMapping failed: %v, output: %q", err, output)
						}
						if output, err := newSMBMapping(username, password, source); err != nil {
							return fmt.Errorf("New-SmbGlobalMapping(%s) failed: %v, output: %q", source, err, output)
						}
					}
				} else {
					klog.V(2).Infof("SMB Mapping(%s) already exists and is still valid, skip error(%v)", source, err)
				}
			} else {
				return fmt.Errorf("New-SmbGlobalMapping(%s) failed: %v, output: %q", source, err, output)
			}
		}
	}

	// There is an issue in golang where EvalSymlinks fails on Windows when passed a
	// UNC share root path without a trailing backslash.
	// Ex: \\SERVER\share will fail to resolve but \\SERVER\share\ will resolve
	// containerD on Windows calls EvalSymlinks so we'll add the backslash when making the symlink if it is missing.
	// https://github.com/golang/go/pull/42096 fixes this issue in golang but a fix will not be available until
	// golang v1.16
	mklinkSource := bindSource
	if !strings.HasSuffix(mklinkSource, "\\") {
		mklinkSource = mklinkSource + "\\"
	}

	output, err := exec.Command("cmd", "/c", "mklink", "/D", target, mklinkSource).CombinedOutput()
	if err != nil {
		klog.Errorf("mklink failed: %v, source(%q) target(%q) output: %q", err, mklinkSource, target, string(output))
		return err
	}
	klog.V(2).Infof("mklink source(%q) on target(%q) successfully, output: %q", mklinkSource, target, string(output))

	return nil
}

// do the SMB mount with username, password, remotepath
// return (output, error)
func newSMBMapping(username, password, remotepath string) (string, error) {
	if username == "" || password == "" || remotepath == "" {
		return "", fmt.Errorf("invalid parameter(username: %s, password: %s, remoteapth: %s)", username, sensitiveOptionsRemoved, remotepath)
	}

	// use PowerShell Environment Variables to store user input string to prevent command line injection
	// https://docs.microsoft.com/en-us/powershell/module/microsoft.powershell.core/about/about_environment_variables?view=powershell-5.1
	cmdLine := `$PWord = ConvertTo-SecureString -String $Env:smbpassword -AsPlainText -Force` +
		`;$Credential = New-Object -TypeName System.Management.Automation.PSCredential -ArgumentList $Env:smbuser, $PWord` +
		`;New-SmbGlobalMapping -RemotePath $Env:smbremotepath -Credential $Credential`
	cmd := exec.Command("powershell", "/c", cmdLine)
	cmd.Env = append(os.Environ(),
		fmt.Sprintf("smbuser=%s", username),
		fmt.Sprintf("smbpassword=%s", password),
		fmt.Sprintf("smbremotepath=%s", remotepath))

	output, err := cmd.CombinedOutput()
	return string(output), err
}

// check whether remotepath is already mounted
func isSMBMappingExist(remotepath string) bool {
	cmd := exec.Command("powershell", "/c", `Get-SmbGlobalMapping -RemotePath $Env:smbremotepath`)
	cmd.Env = append(os.Environ(), fmt.Sprintf("smbremotepath=%s", remotepath))
	_, err := cmd.CombinedOutput()
	return err == nil
}

// check whether remotepath is valid
// return (true, nil) if remotepath is valid
func isValidPath(remotepath string) (bool, error) {
	cmd := exec.Command("powershell", "/c", `Test-Path $Env:remoteapth`)
	cmd.Env = append(os.Environ(), fmt.Sprintf("remoteapth=%s", remotepath))
	output, err := cmd.CombinedOutput()
	if err != nil {
		return false, fmt.Errorf("returned output: %s, error: %v", string(output), err)
	}

	return strings.HasPrefix(strings.ToLower(string(output)), "true"), nil
}

func isAccessDeniedError(err error) bool {
	return err != nil && strings.Contains(strings.ToLower(err.Error()), accessDenied)
}

// remove SMB mapping
func removeSMBMapping(remotepath string) (string, error) {
	cmd := exec.Command("powershell", "/c", `Remove-SmbGlobalMapping -RemotePath $Env:smbremotepath -Force`)
	cmd.Env = append(os.Environ(), fmt.Sprintf("smbremotepath=%s", remotepath))
	output, err := cmd.CombinedOutput()
	return string(output), err
}

// Unmount unmounts the target.
func (mounter *Mounter) Unmount(target string) error {
	klog.V(4).Infof("azureMount: Unmount target (%q)", target)
	target = NormalizeWindowsPath(target)
	if output, err := exec.Command("cmd", "/c", "rmdir", target).CombinedOutput(); err != nil {
		klog.Errorf("rmdir failed: %v, output: %q", err, string(output))
		return err
	}
	return nil
}

// List returns a list of all mounted filesystems. todo
func (mounter *Mounter) List() ([]MountPoint, error) {
	return []MountPoint{}, nil
}

// IsLikelyNotMountPoint determines if a directory is not a mountpoint.
func (mounter *Mounter) IsLikelyNotMountPoint(file string) (bool, error) {
	stat, err := os.Lstat(file)
	if err != nil {
		return true, err
	}

	if stat.Mode()&os.ModeSymlink != 0 {
		return false, err
	}
	return true, nil
}

// GetMountRefs : empty implementation here since there is no place to query all mount points on Windows
func (mounter *Mounter) GetMountRefs(pathname string) ([]string, error) {
	windowsPath := NormalizeWindowsPath(pathname)
	pathExists, pathErr := PathExists(windowsPath)
	if !pathExists {
		return []string{}, nil
	} else if IsCorruptedMnt(pathErr) {
		klog.Warningf("GetMountRefs found corrupted mount at %s, treating as unmounted path", windowsPath)
		return []string{}, nil
	} else if pathErr != nil {
		return nil, fmt.Errorf("error checking path %s: %v", windowsPath, pathErr)
	}
	return []string{pathname}, nil
}

func (mounter *SafeFormatAndMount) formatAndMountSensitive(source string, target string, fstype string, options []string, sensitiveOptions []string) error {
	// Try to mount the disk
	klog.V(4).Infof("Attempting to formatAndMount disk: %s %s %s", fstype, source, target)

	if err := ValidateDiskNumber(source); err != nil {
		klog.Errorf("diskMount: formatAndMount failed, err: %v", err)
		return err
	}

	if len(fstype) == 0 {
		// Use 'NTFS' as the default
		fstype = "NTFS"
	}

	// format disk if it is unformatted(raw)
	cmd := fmt.Sprintf("Get-Disk -Number %s | Where partitionstyle -eq 'raw' | Initialize-Disk -PartitionStyle MBR -PassThru"+
		" | New-Partition -UseMaximumSize | Format-Volume -FileSystem %s -Confirm:$false", source, fstype)
	if output, err := mounter.Exec.Command("powershell", "/c", cmd).CombinedOutput(); err != nil {
		return fmt.Errorf("diskMount: format disk failed, error: %v, output: %q", err, string(output))
	}
	klog.V(4).Infof("diskMount: Disk successfully formatted, disk: %q, fstype: %q", source, fstype)

	volumeIds, err := listVolumesOnDisk(source)
	if err != nil {
		return err
	}
	driverPath := volumeIds[0]
	target = NormalizeWindowsPath(target)
	output, err := mounter.Exec.Command("cmd", "/c", "mklink", "/D", target, driverPath).CombinedOutput()
	if err != nil {
		klog.Errorf("mklink(%s, %s) failed: %v, output: %q", target, driverPath, err, string(output))
		return err
	}
	klog.V(2).Infof("formatAndMount disk(%s) fstype(%s) on(%s) with output(%s) successfully", driverPath, fstype, target, string(output))
	return nil
}

// ListVolumesOnDisk - returns back list of volumes(volumeIDs) in the disk (requested in diskID).
func listVolumesOnDisk(diskID string) (volumeIDs []string, err error) {
	cmd := fmt.Sprintf("(Get-Disk -DeviceId %s | Get-Partition | Get-Volume).UniqueId", diskID)
	output, err := exec.Command("powershell", "/c", cmd).CombinedOutput()
	klog.V(4).Infof("listVolumesOnDisk id from %s: %s", diskID, string(output))
	if err != nil {
		return []string{}, fmt.Errorf("error list volumes on disk. cmd: %s, output: %s, error: %v", cmd, string(output), err)
	}

	volumeIds := strings.Split(strings.TrimSpace(string(output)), "\r\n")
	return volumeIds, nil
}

// getAllParentLinks walks all symbolic links and return all the parent targets recursively
func getAllParentLinks(path string) ([]string, error) {
	const maxIter = 255
	links := []string{}
	for {
		links = append(links, path)
		if len(links) > maxIter {
			return links, fmt.Errorf("unexpected length of parent links: %v", links)
		}

		fi, err := os.Lstat(path)
		if err != nil {
			return links, fmt.Errorf("Lstat: %v", err)
		}
		if fi.Mode()&os.ModeSymlink == 0 {
			break
		}

		path, err = os.Readlink(path)
		if err != nil {
			return links, fmt.Errorf("Readlink error: %v", err)
		}
	}

	return links, nil
}
<commit_msg>fix smb mount issue on Windows<commit_after>// +build windows

/*
Copyright 2017 The Kubernetes Authors.

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
*/

package mount

import (
	"fmt"
	"os"
	"os/exec"
	"path/filepath"
	"strings"

	"k8s.io/klog/v2"
	"k8s.io/utils/keymutex"
)

const (
	accessDenied string = "access is denied"
)

// Mounter provides the default implementation of mount.Interface
// for the windows platform.  This implementation assumes that the
// kubelet is running in the host's root mount namespace.
type Mounter struct {
	mounterPath string
}

// New returns a mount.Interface for the current system.
// It provides options to override the default mounter behavior.
// mounterPath allows using an alternative to `/bin/mount` for mounting.
func New(mounterPath string) Interface {
	return &Mounter{
		mounterPath: mounterPath,
	}
}

// acquire lock for smb mount
var getSMBMountMutex = keymutex.NewHashed(0)

// Mount : mounts source to target with given options.
// currently only supports cifs(smb), bind mount(for disk)
func (mounter *Mounter) Mount(source string, target string, fstype string, options []string) error {
	return mounter.MountSensitive(source, target, fstype, options, nil /* sensitiveOptions */)
}

// MountSensitiveWithoutSystemd is the same as MountSensitive() but disable using ssytemd mount.
// Windows not supported systemd mount, this function degrades to MountSensitive().
func (mounter *Mounter) MountSensitiveWithoutSystemd(source string, target string, fstype string, options []string, sensitiveOptions []string) error {
	return mounter.MountSensitive(source, target, fstype, options, sensitiveOptions /* sensitiveOptions */)
}

// MountSensitive is the same as Mount() but this method allows
// sensitiveOptions to be passed in a separate parameter from the normal
// mount options and ensures the sensitiveOptions are never logged. This
// method should be used by callers that pass sensitive material (like
// passwords) as mount options.
func (mounter *Mounter) MountSensitive(source string, target string, fstype string, options []string, sensitiveOptions []string) error {
	target = NormalizeWindowsPath(target)
	sanitizedOptionsForLogging := sanitizedOptionsForLogging(options, sensitiveOptions)

	if source == "tmpfs" {
		klog.V(3).Infof("mounting source (%q), target (%q), with options (%q)", source, target, sanitizedOptionsForLogging)
		return os.MkdirAll(target, 0755)
	}

	parentDir := filepath.Dir(target)
	if err := os.MkdirAll(parentDir, 0755); err != nil {
		return err
	}

	klog.V(4).Infof("mount options(%q) source:%q, target:%q, fstype:%q, begin to mount",
		sanitizedOptionsForLogging, source, target, fstype)
	bindSource := source

	if bind, _, _, _ := MakeBindOptsSensitive(options, sensitiveOptions); bind {
		bindSource = NormalizeWindowsPath(source)
	} else {
		allOptions := []string{}
		allOptions = append(allOptions, options...)
		allOptions = append(allOptions, sensitiveOptions...)
		if len(allOptions) < 2 {
			return fmt.Errorf("mount options(%q) should have at least 2 options, current number:%d, source:%q, target:%q",
				sanitizedOptionsForLogging, len(allOptions), source, target)
		}

		// currently only cifs mount is supported
		if strings.ToLower(fstype) != "cifs" {
			return fmt.Errorf("only cifs mount is supported now, fstype: %q, mounting source (%q), target (%q), with options (%q)", fstype, source, target, sanitizedOptionsForLogging)
		}

		// lock smb mount for the same source
		getSMBMountMutex.LockKey(source)
		defer getSMBMountMutex.UnlockKey(source)

		username := allOptions[0]
		password := allOptions[1]
		if output, err := newSMBMapping(username, password, source); err != nil {
			klog.Warningf("SMB Mapping(%s) returned with error(%v), output(%s)", source, err, string(output))
			if isSMBMappingExist(source) {
				valid, err := isValidPath(source)
				if !valid {
					if err == nil || isAccessDeniedError(err) {
						klog.V(2).Infof("SMB Mapping(%s) already exists while it's not valid, return error: %v, now begin to remove and remount", source, err)
						if output, err = removeSMBMapping(source); err != nil {
							return fmt.Errorf("Remove-SmbGlobalMapping failed: %v, output: %q", err, output)
						}
						if output, err := newSMBMapping(username, password, source); err != nil {
							return fmt.Errorf("New-SmbGlobalMapping(%s) failed: %v, output: %q", source, err, output)
						}
					}
				} else {
					klog.V(2).Infof("SMB Mapping(%s) already exists and is still valid, skip error(%v)", source, err)
				}
			} else {
				return fmt.Errorf("New-SmbGlobalMapping(%s) failed: %v, output: %q", source, err, output)
			}
		}
	}

	// There is an issue in golang where EvalSymlinks fails on Windows when passed a
	// UNC share root path without a trailing backslash.
	// Ex: \\SERVER\share will fail to resolve but \\SERVER\share\ will resolve
	// containerD on Windows calls EvalSymlinks so we'll add the backslash when making the symlink if it is missing.
	// https://github.com/golang/go/pull/42096 fixes this issue in golang but a fix will not be available until
	// golang v1.16
	mklinkSource := bindSource
	if !strings.HasSuffix(mklinkSource, "\\") {
		mklinkSource = mklinkSource + "\\"
	}

	output, err := exec.Command("cmd", "/c", "mklink", "/D", target, mklinkSource).CombinedOutput()
	if err != nil {
		klog.Errorf("mklink failed: %v, source(%q) target(%q) output: %q", err, mklinkSource, target, string(output))
		return err
	}
	klog.V(2).Infof("mklink source(%q) on target(%q) successfully, output: %q", mklinkSource, target, string(output))

	return nil
}

// do the SMB mount with username, password, remotepath
// return (output, error)
func newSMBMapping(username, password, remotepath string) (string, error) {
	if username == "" || password == "" || remotepath == "" {
		return "", fmt.Errorf("invalid parameter(username: %s, password: %s, remoteapth: %s)", username, sensitiveOptionsRemoved, remotepath)
	}

	// use PowerShell Environment Variables to store user input string to prevent command line injection
	// https://docs.microsoft.com/en-us/powershell/module/microsoft.powershell.core/about/about_environment_variables?view=powershell-5.1
	cmdLine := `$PWord = ConvertTo-SecureString -String $Env:smbpassword -AsPlainText -Force` +
		`;$Credential = New-Object -TypeName System.Management.Automation.PSCredential -ArgumentList $Env:smbuser, $PWord` +
		`;New-SmbGlobalMapping -RemotePath $Env:smbremotepath -Credential $Credential -RequirePrivacy $true`
	cmd := exec.Command("powershell", "/c", cmdLine)
	cmd.Env = append(os.Environ(),
		fmt.Sprintf("smbuser=%s", username),
		fmt.Sprintf("smbpassword=%s", password),
		fmt.Sprintf("smbremotepath=%s", remotepath))

	output, err := cmd.CombinedOutput()
	return string(output), err
}

// check whether remotepath is already mounted
func isSMBMappingExist(remotepath string) bool {
	cmd := exec.Command("powershell", "/c", `Get-SmbGlobalMapping -RemotePath $Env:smbremotepath`)
	cmd.Env = append(os.Environ(), fmt.Sprintf("smbremotepath=%s", remotepath))
	_, err := cmd.CombinedOutput()
	return err == nil
}

// check whether remotepath is valid
// return (true, nil) if remotepath is valid
func isValidPath(remotepath string) (bool, error) {
	cmd := exec.Command("powershell", "/c", `Test-Path $Env:remoteapth`)
	cmd.Env = append(os.Environ(), fmt.Sprintf("remoteapth=%s", remotepath))
	output, err := cmd.CombinedOutput()
	if err != nil {
		return false, fmt.Errorf("returned output: %s, error: %v", string(output), err)
	}

	return strings.HasPrefix(strings.ToLower(string(output)), "true"), nil
}

func isAccessDeniedError(err error) bool {
	return err != nil && strings.Contains(strings.ToLower(err.Error()), accessDenied)
}

// remove SMB mapping
func removeSMBMapping(remotepath string) (string, error) {
	cmd := exec.Command("powershell", "/c", `Remove-SmbGlobalMapping -RemotePath $Env:smbremotepath -Force`)
	cmd.Env = append(os.Environ(), fmt.Sprintf("smbremotepath=%s", remotepath))
	output, err := cmd.CombinedOutput()
	return string(output), err
}

// Unmount unmounts the target.
func (mounter *Mounter) Unmount(target string) error {
	klog.V(4).Infof("azureMount: Unmount target (%q)", target)
	target = NormalizeWindowsPath(target)
	if output, err := exec.Command("cmd", "/c", "rmdir", target).CombinedOutput(); err != nil {
		klog.Errorf("rmdir failed: %v, output: %q", err, string(output))
		return err
	}
	return nil
}

// List returns a list of all mounted filesystems. todo
func (mounter *Mounter) List() ([]MountPoint, error) {
	return []MountPoint{}, nil
}

// IsLikelyNotMountPoint determines if a directory is not a mountpoint.
func (mounter *Mounter) IsLikelyNotMountPoint(file string) (bool, error) {
	stat, err := os.Lstat(file)
	if err != nil {
		return true, err
	}

	if stat.Mode()&os.ModeSymlink != 0 {
		return false, err
	}
	return true, nil
}

// GetMountRefs : empty implementation here since there is no place to query all mount points on Windows
func (mounter *Mounter) GetMountRefs(pathname string) ([]string, error) {
	windowsPath := NormalizeWindowsPath(pathname)
	pathExists, pathErr := PathExists(windowsPath)
	if !pathExists {
		return []string{}, nil
	} else if IsCorruptedMnt(pathErr) {
		klog.Warningf("GetMountRefs found corrupted mount at %s, treating as unmounted path", windowsPath)
		return []string{}, nil
	} else if pathErr != nil {
		return nil, fmt.Errorf("error checking path %s: %v", windowsPath, pathErr)
	}
	return []string{pathname}, nil
}

func (mounter *SafeFormatAndMount) formatAndMountSensitive(source string, target string, fstype string, options []string, sensitiveOptions []string) error {
	// Try to mount the disk
	klog.V(4).Infof("Attempting to formatAndMount disk: %s %s %s", fstype, source, target)

	if err := ValidateDiskNumber(source); err != nil {
		klog.Errorf("diskMount: formatAndMount failed, err: %v", err)
		return err
	}

	if len(fstype) == 0 {
		// Use 'NTFS' as the default
		fstype = "NTFS"
	}

	// format disk if it is unformatted(raw)
	cmd := fmt.Sprintf("Get-Disk -Number %s | Where partitionstyle -eq 'raw' | Initialize-Disk -PartitionStyle MBR -PassThru"+
		" | New-Partition -UseMaximumSize | Format-Volume -FileSystem %s -Confirm:$false", source, fstype)
	if output, err := mounter.Exec.Command("powershell", "/c", cmd).CombinedOutput(); err != nil {
		return fmt.Errorf("diskMount: format disk failed, error: %v, output: %q", err, string(output))
	}
	klog.V(4).Infof("diskMount: Disk successfully formatted, disk: %q, fstype: %q", source, fstype)

	volumeIds, err := listVolumesOnDisk(source)
	if err != nil {
		return err
	}
	driverPath := volumeIds[0]
	target = NormalizeWindowsPath(target)
	output, err := mounter.Exec.Command("cmd", "/c", "mklink", "/D", target, driverPath).CombinedOutput()
	if err != nil {
		klog.Errorf("mklink(%s, %s) failed: %v, output: %q", target, driverPath, err, string(output))
		return err
	}
	klog.V(2).Infof("formatAndMount disk(%s) fstype(%s) on(%s) with output(%s) successfully", driverPath, fstype, target, string(output))
	return nil
}

// ListVolumesOnDisk - returns back list of volumes(volumeIDs) in the disk (requested in diskID).
func listVolumesOnDisk(diskID string) (volumeIDs []string, err error) {
	cmd := fmt.Sprintf("(Get-Disk -DeviceId %s | Get-Partition | Get-Volume).UniqueId", diskID)
	output, err := exec.Command("powershell", "/c", cmd).CombinedOutput()
	klog.V(4).Infof("listVolumesOnDisk id from %s: %s", diskID, string(output))
	if err != nil {
		return []string{}, fmt.Errorf("error list volumes on disk. cmd: %s, output: %s, error: %v", cmd, string(output), err)
	}

	volumeIds := strings.Split(strings.TrimSpace(string(output)), "\r\n")
	return volumeIds, nil
}

// getAllParentLinks walks all symbolic links and return all the parent targets recursively
func getAllParentLinks(path string) ([]string, error) {
	const maxIter = 255
	links := []string{}
	for {
		links = append(links, path)
		if len(links) > maxIter {
			return links, fmt.Errorf("unexpected length of parent links: %v", links)
		}

		fi, err := os.Lstat(path)
		if err != nil {
			return links, fmt.Errorf("Lstat: %v", err)
		}
		if fi.Mode()&os.ModeSymlink == 0 {
			break
		}

		path, err = os.Readlink(path)
		if err != nil {
			return links, fmt.Errorf("Readlink error: %v", err)
		}
	}

	return links, nil
}
<|endoftext|>


<commit_before>// Copyright 2020 Google LLC
//
// 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.

// This server is a simple webserver that triggers the e2e-test binary.
package main

import (
	"context"
	"crypto/rand"
	"fmt"
	"math/big"
	"net/http"
	"os"
	"strconv"
	"time"

	"github.com/google/exposure-notifications-server/pkg/logging"
	"github.com/google/exposure-notifications-server/pkg/observability"
	"github.com/google/exposure-notifications-server/pkg/server"

	"github.com/google/exposure-notifications-verification-server/pkg/buildinfo"
	"github.com/google/exposure-notifications-verification-server/pkg/clients"
	"github.com/google/exposure-notifications-verification-server/pkg/config"
	"github.com/google/exposure-notifications-verification-server/pkg/database"

	"github.com/gorilla/handlers"
	"github.com/gorilla/mux"
	"github.com/sethvargo/go-signalcontext"
)

const (
	realmName       = "e2e-test-realm"
	realmRegionCode = "e2e-test"
	adminKeyName    = "e2e-admin-key."
	deviceKeyName   = "e2e-device-key."
)

func main() {
	ctx, done := signalcontext.OnInterrupt()

	debug, _ := strconv.ParseBool(os.Getenv("LOG_DEBUG"))
	logger := logging.NewLogger(debug)
	logger = logger.With("build_id", buildinfo.BuildID)
	logger = logger.With("build_tag", buildinfo.BuildTag)

	ctx = logging.WithLogger(ctx, logger)

	err := realMain(ctx)
	done()

	if err != nil {
		logger.Fatal(err)
	}
	logger.Info("successful shutdown")
}

func randomString() (string, error) {
	n, err := rand.Int(rand.Reader, big.NewInt(10000))
	if err != nil {
		return "", err
	}
	return fmt.Sprintf("%04x", n), nil
}

func realMain(ctx context.Context) error {
	logger := logging.FromContext(ctx)

	// load configs
	e2eConfig, err := config.NewE2ERunnerConfig(ctx)
	if err != nil {
		return fmt.Errorf("failed to process e2e-runner config: %w", err)
	}

	// Setup monitoring
	logger.Info("configuring observability exporter")
	oe, err := observability.NewFromEnv(ctx, e2eConfig.Observability)
	if err != nil {
		return fmt.Errorf("unable to create ObservabilityExporter provider: %w", err)
	}
	if err := oe.StartExporter(); err != nil {
		return fmt.Errorf("error initializing observability exporter: %w", err)
	}
	defer oe.Close()
	logger.Infow("observability exporter", "config", e2eConfig.Observability)

	db, err := e2eConfig.Database.Load(ctx)
	if err != nil {
		return fmt.Errorf("failed to load database config: %w", err)
	}
	if err := db.Open(ctx); err != nil {
		return fmt.Errorf("failed to connect to database: %w", err)
	}
	defer db.Close()

	// Create or reuse the existing realm
	realm, err := db.FindRealmByName(realmName)
	if err != nil {
		if !database.IsNotFound(err) {
			return fmt.Errorf("error when finding the realm %q: %w", realmName, err)
		}
		realm = database.NewRealmWithDefaults(realmName)
		realm.RegionCode = realmRegionCode
		if err := db.SaveRealm(realm, database.System); err != nil {
			return fmt.Errorf("failed to create realm %+v: %w: %v", realm, err, realm.ErrorMessages())
		}
	}

	// Create new API keys
	suffix, err := randomString()
	if err != nil {
		return fmt.Errorf("failed to create suffix string for API keys: %w", err)
	}

	adminKey, err := realm.CreateAuthorizedApp(db, &database.AuthorizedApp{
		Name:       adminKeyName + suffix,
		APIKeyType: database.APIKeyTypeAdmin,
	}, database.System)
	if err != nil {
		return fmt.Errorf("error trying to create a new Admin API Key: %w", err)
	}

	defer func() {
		app, err := db.FindAuthorizedAppByAPIKey(adminKey)
		if err != nil {
			logger.Errorf("admin API key cleanup failed: %w", err)
		}
		now := time.Now().UTC()
		app.DeletedAt = &now
		if err := db.SaveAuthorizedApp(app, database.System); err != nil {
			logger.Errorf("admin API key disable failed: %w", err)
		}
		logger.Info("successfully cleaned up e2e test admin key")
	}()

	deviceKey, err := realm.CreateAuthorizedApp(db, &database.AuthorizedApp{
		Name:       deviceKeyName + suffix,
		APIKeyType: database.APIKeyTypeDevice,
	}, database.System)
	if err != nil {
		return fmt.Errorf("error trying to create a new Device API Key: %w", err)
	}

	defer func() {
		app, err := db.FindAuthorizedAppByAPIKey(deviceKey)
		if err != nil {
			logger.Errorf("device API key cleanup failed: %w", err)
		}
		now := time.Now().UTC()
		app.DeletedAt = &now
		if err := db.SaveAuthorizedApp(app, database.System); err != nil {
			logger.Errorf("device API key disable failed: %w", err)
		}
		logger.Info("successfully cleaned up e2e test device key")
	}()

	e2eConfig.TestConfig.VerificationAdminAPIKey = adminKey
	e2eConfig.TestConfig.VerificationAPIServerKey = deviceKey

	// Create the router
	r := mux.NewRouter()
	r.HandleFunc("/default", defaultHandler(ctx, &e2eConfig.TestConfig))
	r.HandleFunc("/revise", reviseHandler(ctx, &e2eConfig.TestConfig))

	srv, err := server.New(e2eConfig.Port)
	if err != nil {
		return fmt.Errorf("failed to create server: %w", err)
	}
	logger.Infow("server listening", "port", e2eConfig.Port)
	return srv.ServeHTTPHandler(ctx, handlers.CombinedLoggingHandler(os.Stdout, r))
}

func defaultHandler(ctx context.Context, c *config.E2ETestConfig) func(http.ResponseWriter, *http.Request) {
	return func(w http.ResponseWriter, r *http.Request) {
		c.DoRevise = false
		if err := clients.RunEndToEnd(ctx, c); err != nil {
			http.Error(w, "failed (check server logs for more details): "+err.Error(), http.StatusInternalServerError)
		} else {
			fmt.Fprint(w, "ok")
		}
	}
}

func reviseHandler(ctx context.Context, c *config.E2ETestConfig) func(http.ResponseWriter, *http.Request) {
	return func(w http.ResponseWriter, r *http.Request) {
		c.DoRevise = true
		if err := clients.RunEndToEnd(ctx, c); err != nil {
			http.Error(w, "failed (check server logs for more details): "+err.Error(), http.StatusInternalServerError)
		} else {
			fmt.Fprint(w, "ok")
		}
	}
}
<commit_msg>Add log errors for e2e fail (#786)<commit_after>// Copyright 2020 Google LLC
//
// 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.

// This server is a simple webserver that triggers the e2e-test binary.
package main

import (
	"context"
	"crypto/rand"
	"fmt"
	"math/big"
	"net/http"
	"os"
	"strconv"
	"time"

	"github.com/google/exposure-notifications-server/pkg/logging"
	"github.com/google/exposure-notifications-server/pkg/observability"
	"github.com/google/exposure-notifications-server/pkg/server"

	"github.com/google/exposure-notifications-verification-server/pkg/buildinfo"
	"github.com/google/exposure-notifications-verification-server/pkg/clients"
	"github.com/google/exposure-notifications-verification-server/pkg/config"
	"github.com/google/exposure-notifications-verification-server/pkg/database"

	"github.com/gorilla/handlers"
	"github.com/gorilla/mux"
	"github.com/sethvargo/go-signalcontext"
)

const (
	realmName       = "e2e-test-realm"
	realmRegionCode = "e2e-test"
	adminKeyName    = "e2e-admin-key."
	deviceKeyName   = "e2e-device-key."
)

func main() {
	ctx, done := signalcontext.OnInterrupt()

	debug, _ := strconv.ParseBool(os.Getenv("LOG_DEBUG"))
	logger := logging.NewLogger(debug)
	logger = logger.With("build_id", buildinfo.BuildID)
	logger = logger.With("build_tag", buildinfo.BuildTag)

	ctx = logging.WithLogger(ctx, logger)

	err := realMain(ctx)
	done()

	if err != nil {
		logger.Fatal(err)
	}
	logger.Info("successful shutdown")
}

func randomString() (string, error) {
	n, err := rand.Int(rand.Reader, big.NewInt(10000))
	if err != nil {
		return "", err
	}
	return fmt.Sprintf("%04x", n), nil
}

func realMain(ctx context.Context) error {
	logger := logging.FromContext(ctx)

	// load configs
	e2eConfig, err := config.NewE2ERunnerConfig(ctx)
	if err != nil {
		return fmt.Errorf("failed to process e2e-runner config: %w", err)
	}

	// Setup monitoring
	logger.Info("configuring observability exporter")
	oe, err := observability.NewFromEnv(ctx, e2eConfig.Observability)
	if err != nil {
		return fmt.Errorf("unable to create ObservabilityExporter provider: %w", err)
	}
	if err := oe.StartExporter(); err != nil {
		return fmt.Errorf("error initializing observability exporter: %w", err)
	}
	defer oe.Close()
	logger.Infow("observability exporter", "config", e2eConfig.Observability)

	db, err := e2eConfig.Database.Load(ctx)
	if err != nil {
		return fmt.Errorf("failed to load database config: %w", err)
	}
	if err := db.Open(ctx); err != nil {
		return fmt.Errorf("failed to connect to database: %w", err)
	}
	defer db.Close()

	// Create or reuse the existing realm
	realm, err := db.FindRealmByName(realmName)
	if err != nil {
		if !database.IsNotFound(err) {
			return fmt.Errorf("error when finding the realm %q: %w", realmName, err)
		}
		realm = database.NewRealmWithDefaults(realmName)
		realm.RegionCode = realmRegionCode
		if err := db.SaveRealm(realm, database.System); err != nil {
			return fmt.Errorf("failed to create realm %+v: %w: %v", realm, err, realm.ErrorMessages())
		}
	}

	// Create new API keys
	suffix, err := randomString()
	if err != nil {
		return fmt.Errorf("failed to create suffix string for API keys: %w", err)
	}

	adminKey, err := realm.CreateAuthorizedApp(db, &database.AuthorizedApp{
		Name:       adminKeyName + suffix,
		APIKeyType: database.APIKeyTypeAdmin,
	}, database.System)
	if err != nil {
		return fmt.Errorf("error trying to create a new Admin API Key: %w", err)
	}

	defer func() {
		app, err := db.FindAuthorizedAppByAPIKey(adminKey)
		if err != nil {
			logger.Errorf("admin API key cleanup failed: %w", err)
		}
		now := time.Now().UTC()
		app.DeletedAt = &now
		if err := db.SaveAuthorizedApp(app, database.System); err != nil {
			logger.Errorf("admin API key disable failed: %w", err)
		}
		logger.Info("successfully cleaned up e2e test admin key")
	}()

	deviceKey, err := realm.CreateAuthorizedApp(db, &database.AuthorizedApp{
		Name:       deviceKeyName + suffix,
		APIKeyType: database.APIKeyTypeDevice,
	}, database.System)
	if err != nil {
		return fmt.Errorf("error trying to create a new Device API Key: %w", err)
	}

	defer func() {
		app, err := db.FindAuthorizedAppByAPIKey(deviceKey)
		if err != nil {
			logger.Errorf("device API key cleanup failed: %w", err)
		}
		now := time.Now().UTC()
		app.DeletedAt = &now
		if err := db.SaveAuthorizedApp(app, database.System); err != nil {
			logger.Errorf("device API key disable failed: %w", err)
		}
		logger.Info("successfully cleaned up e2e test device key")
	}()

	e2eConfig.TestConfig.VerificationAdminAPIKey = adminKey
	e2eConfig.TestConfig.VerificationAPIServerKey = deviceKey

	// Create the router
	r := mux.NewRouter()
	r.HandleFunc("/default", defaultHandler(ctx, &e2eConfig.TestConfig))
	r.HandleFunc("/revise", reviseHandler(ctx, &e2eConfig.TestConfig))

	srv, err := server.New(e2eConfig.Port)
	if err != nil {
		return fmt.Errorf("failed to create server: %w", err)
	}
	logger.Infow("server listening", "port", e2eConfig.Port)
	return srv.ServeHTTPHandler(ctx, handlers.CombinedLoggingHandler(os.Stdout, r))
}

func defaultHandler(ctx context.Context, c *config.E2ETestConfig) func(http.ResponseWriter, *http.Request) {
	logger := logging.FromContext(ctx)
	return func(w http.ResponseWriter, r *http.Request) {
		c.DoRevise = false
		if err := clients.RunEndToEnd(ctx, c); err != nil {
			logger.Errorw("could not run default end to end", "error", err)
			http.Error(w, "failed (check server logs for more details): "+err.Error(), http.StatusInternalServerError)
			return
		}

		fmt.Fprint(w, "ok")
	}
}

func reviseHandler(ctx context.Context, c *config.E2ETestConfig) func(http.ResponseWriter, *http.Request) {
	logger := logging.FromContext(ctx)
	return func(w http.ResponseWriter, r *http.Request) {
		c.DoRevise = true
		if err := clients.RunEndToEnd(ctx, c); err != nil {
			logger.Errorw("could not run revise end to end", "error", err)
			http.Error(w, "failed (check server logs for more details): "+err.Error(), http.StatusInternalServerError)
			return
		}

		fmt.Fprint(w, "ok")
	}
}
<|endoftext|>


<commit_before>/*
 * Copyright 2013 Murali Suriar
 *
 * 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 in the LICENSE file, or at:
 *
 *     https://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

package main

import (
	"flag"
	"fmt"
	"github.com/msuriar/garbler"
	"os"
	"time"
)

var (
	prefix,cmd string
	advInt, cmd_timeout time.Duration
	successes, failures int
)

func init() {
	flag.StringVar(&prefix, "prefix", "", "Prefix to advertise.")
	flag.StringVar(&cmd, "cmd", "", "Healthcheck command to run.")
	flag.DurationVar(&advInt, "advInt", 30 * time.Second, `Interval between
	updates.`)
	flag.IntVar(&successes, "successes", 3, `Number of consecutive
	successes before garbler considers the service healthy.`)
	flag.IntVar(&failures, "failures", 3, `Number of consecutive
	failures before garbler considers the service unhealthy.`)
	flag.DurationVar(&cmd_timeout, "timeout", 1 * time.Second,
	`Timeout for healthcheck command.`)
}

func main() {
	flag.Parse()

	err := false

	if prefix == "" {
		fmt.Println("prefix is a required flag.")
		err = true
	}

	if cmd == "" {
		fmt.Println("cmd is a required flag.")
		err = true
	}

	if advInt <= cmd_timeout {
		fmt.Println("advInt is less than timeout. This is a bad idea.")
		err = true
	}

	if err { os.Exit(1) }

	garbler.StartService(cmd, prefix, advInt, cmd_timeout)
}
<commit_msg>Flag cleanup.<commit_after>/*
 * Copyright 2013 Murali Suriar
 *
 * 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 in the LICENSE file, or at:
 *
 *     https://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

package main

import (
	"flag"
	"fmt"
	"github.com/msuriar/garbler"
	"os"
	"time"
)

var (
	prefix,cmd string
	advInt, cmd_timeout time.Duration
	successes, failures int
)

func init() {
	flag.StringVar(&prefix, "prefix", "", "Prefix to advertise.")
	flag.StringVar(&cmd, "cmd", "", "Healthcheck command to run.")
	flag.DurationVar(&advInt, "advInt", 30 * time.Second,
	"Interval between updates.")
	flag.IntVar(&successes, "successes", 3,
	"Number of consecutive successes before service healthy.")
	flag.IntVar(&failures, "failures", 3,
	"Number of consecutive failures before service unhealthy.")
	flag.DurationVar(&cmd_timeout, "timeout", 1 * time.Second,
	"Timeout for healthcheck command.")
}

func main() {
	flag.Parse()

	err := false

	if prefix == "" {
		fmt.Println("prefix is a required flag.")
		err = true
	}

	if cmd == "" {
		fmt.Println("cmd is a required flag.")
		err = true
	}

	if advInt <= cmd_timeout {
		fmt.Println("advInt is less than timeout. This is a bad idea.")
		err = true
	}

	if err { os.Exit(1) }

	garbler.StartService(cmd, prefix, advInt, cmd_timeout)
}
<|endoftext|>


<commit_before>package main

import (
	"flag"
	"fmt"
	"log"

	"github.com/akhenakh/gozim"
	"github.com/blevesearch/bleve"
	"github.com/blevesearch/bleve/registry"
)

type ArticleIndex struct {
	Title string
}

var (
	path       = flag.String("path", "", "path for the zim file")
	indexPath  = flag.String("index", "", "path for the index file")
	cpuprofile = flag.String("cpuprofile", "", "write cpu profile to file")
	z          *zim.ZimReader
	lang       = flag.String("lang", "en", "language for indexation")
)

func inList(s []string, value string) bool {
	for _, v := range s {
		if v == value {
			return true
		}
	}
	return false
}

// Type return the Article type (used for bleve indexer)
func (a *ArticleIndex) Type() string {
	return "Article"
}

func main() {

	flag.Parse()

	if *path == "" {
		panic("provide a zim file path")
	}

	z, err := zim.NewReader(*path, false)
	if err != nil {
		panic(err)
	}

	if *indexPath == "" {
		panic("Please provide a path for the index")
	}

	switch *lang {
	case "en":
		//TODO: create a simple language support for stop word
	default:
		panic("unsupported language")
	}

	mapping := bleve.NewIndexMapping()
	mapping.DefaultType = "Article"

	articleMapping := bleve.NewDocumentMapping()
	mapping.AddDocumentMapping("Article", articleMapping)

	titleMapping := bleve.NewTextFieldMapping()
	titleMapping.Store = false
	titleMapping.Index = true
	titleMapping.Analyzer = "standard"
	articleMapping.AddFieldMappingsAt("Title", titleMapping)

	fmt.Println(registry.AnalyzerTypesAndInstances())

	index, err := bleve.New(*indexPath, mapping)
	if err != nil {
		panic(err)
	}

	i := 0

	batch := bleve.NewBatch()
	batchCount := 0
	idoc := ArticleIndex{}

	z.ListTitlesPtrIterator(func(idx uint32) {

		if i == 10000 {
			fmt.Print("*")
			i = 0
		}

		offset := z.OffsetAtURLIdx(idx)
		a := z.ArticleAt(offset)
		if a.EntryType == zim.DeletedEntry {
			i++
			return
		}

		if a.Namespace == 'A' {
			idoc.Title = a.Title
			// index the idoc with the idx as key
			batch.Index(fmt.Sprint(idx), idoc)
			batchCount++
		}

		i++

		// send a batch to bleve
		if batchCount >= 10000 {
			err = index.Batch(batch)
			if err != nil {
				log.Fatal(err.Error())
			}
			batch = bleve.NewBatch()
			batchCount = 0
		}
	})

	// batch the rest
	if batchCount > 0 {
		err = index.Batch(batch)
		if err != nil {
			log.Fatal(err.Error())
		}
	}

	index.Close()

}
<commit_msg>Use log.Fatal instead of panic like gozimhttpd<commit_after>package main

import (
	"errors"
	"flag"
	"fmt"
	"log"

	"github.com/akhenakh/gozim"
	"github.com/blevesearch/bleve"
	"github.com/blevesearch/bleve/registry"
)

type ArticleIndex struct {
	Title string
}

var (
	path       = flag.String("path", "", "path for the zim file")
	indexPath  = flag.String("index", "", "path for the index file")
	cpuprofile = flag.String("cpuprofile", "", "write cpu profile to file")
	z          *zim.ZimReader
	lang       = flag.String("lang", "en", "language for indexation")
)

func inList(s []string, value string) bool {
	for _, v := range s {
		if v == value {
			return true
		}
	}
	return false
}

// Type return the Article type (used for bleve indexer)
func (a *ArticleIndex) Type() string {
	return "Article"
}

func main() {

	flag.Parse()

	if *path == "" {
		log.Fatal(errors.New("provide a zim file path"))
	}

	z, err := zim.NewReader(*path, false)
	if err != nil {
		panic(err)
	}

	if *indexPath == "" {
		panic("Please provide a path for the index")
	}

	switch *lang {
	case "en":
		//TODO: create a simple language support for stop word
	default:
		panic("unsupported language")
	}

	mapping := bleve.NewIndexMapping()
	mapping.DefaultType = "Article"

	articleMapping := bleve.NewDocumentMapping()
	mapping.AddDocumentMapping("Article", articleMapping)

	titleMapping := bleve.NewTextFieldMapping()
	titleMapping.Store = false
	titleMapping.Index = true
	titleMapping.Analyzer = "standard"
	articleMapping.AddFieldMappingsAt("Title", titleMapping)

	fmt.Println(registry.AnalyzerTypesAndInstances())

	index, err := bleve.New(*indexPath, mapping)
	if err != nil {
		panic(err)
	}

	i := 0

	batch := bleve.NewBatch()
	batchCount := 0
	idoc := ArticleIndex{}

	z.ListTitlesPtrIterator(func(idx uint32) {

		if i == 10000 {
			fmt.Print("*")
			i = 0
		}

		offset := z.OffsetAtURLIdx(idx)
		a := z.ArticleAt(offset)
		if a.EntryType == zim.DeletedEntry {
			i++
			return
		}

		if a.Namespace == 'A' {
			idoc.Title = a.Title
			// index the idoc with the idx as key
			batch.Index(fmt.Sprint(idx), idoc)
			batchCount++
		}

		i++

		// send a batch to bleve
		if batchCount >= 10000 {
			err = index.Batch(batch)
			if err != nil {
				log.Fatal(err.Error())
			}
			batch = bleve.NewBatch()
			batchCount = 0
		}
	})

	// batch the rest
	if batchCount > 0 {
		err = index.Batch(batch)
		if err != nil {
			log.Fatal(err.Error())
		}
	}

	index.Close()

}
<|endoftext|>