<commit_before>package main

import (
	"errors"
	"fmt"
	"github.com/fzzy/radix/redis"
	"time"
)

const (
	REDIS_KEY     = "hchecker"
	REDIS_ADDRESS = "localhost:6379"
)

var (
	redisAddress string
)

type Cache struct {
	redisConn *redis.Client
	redisSub  *redis.Subscription
	// Maintain a mapping between a backends and several frontend
	// -> map[BACKEND_URL][FRONTEND_NAME][BACKEND_ID]
	backendsMapping map[string]map[string][]int
}

func NewCache() (*Cache, error) {
	conf := redis.DefaultConfig()
	conf.Address = redisAddress
	redisConn := redis.NewClient(conf)
	cache := &Cache{
		redisConn:       redisConn,
		backendsMapping: make(map[string]map[string][]int)}
	// We're starting, let's clear any previous meta-data
	// WARNING: This can be a problem if there are several processes sharing
	// the same redis on the same machine. If one of them is restarted, it'll
	// clear the meta-data of everyone...
	redisConn.Del(REDIS_KEY)
	return cache, nil
}

/*
 * Maintain a mapping between Frontends and Backends ID
 */
func (c *Cache) updateFrontendMapping(check *Check) {
	m, exists := c.backendsMapping[check.BackendUrl]
	if !exists {
		m = make(map[string][]int)
	}
	if val, ok := m[check.FrontendKey]; ok {
		for _, item := range val {
			if item == check.BackendId {
				// The backendId has been already stored, ignoring
				return
			}
		}
		// The frontend already exists, adding the new backendId
		m[check.FrontendKey] = append(val, check.BackendId)
	} else {
		// The frontend is new
		m[check.FrontendKey] = []int{check.BackendId}
	}
	c.backendsMapping[check.BackendUrl] = m
}

/*
 * Lock a backend in Redis by its URL
 */
func (c *Cache) LockBackend(check *Check) bool {
	// The syncKey makes sure an entire backend mapping is keep in the same
	// process (we never update a backend mapping from 2 different processes)
	syncKey := check.BackendUrl + ";" + myId
	// Lock the backend with a temporary value, we'll update this with the
	// goroutine signature later
	resp := c.redisConn.Transaction(func(mc *redis.MultiCall) {
		mc.Hsetnx(REDIS_KEY, check.BackendUrl, 1)
		mc.Hexists(REDIS_KEY, syncKey)
	})
	locked, _ := resp.Elems[0].Bool()
	isMine, _ := resp.Elems[1].Bool()
	if locked == false && isMine == false {
		// The backend is being monitored by someone else
		return false
	}
	if locked == true {
		// we got the lock, let's create a unique sig for the goroutine
		t := time.Now()
		// This one is done in the lock, this will garanty that no routine
		// will get the same sig
		sig := fmt.Sprintf("%s;%d.%d", myId, t.Unix(), t.Nanosecond())
		c.redisConn.Hset(REDIS_KEY, check.BackendUrl, sig)
		c.redisConn.Hset(REDIS_KEY, syncKey, 1)
		check.routineSig = sig
	}
	// Let's update the mapping in case this is a new frontend
	c.updateFrontendMapping(check)
	return locked
}

func (c *Cache) IsUnlockedBackend(check *Check) bool {
	// On top of checking the lock, we compare the lock content to make sure
	// we still own the lock
	resp, _ := c.redisConn.Hget(REDIS_KEY, check.BackendUrl).Str()
	return (resp != check.routineSig)
}

func (c *Cache) UnlockBackend(check *Check) {
	c.redisConn.Hdel(REDIS_KEY, check.BackendUrl,
		check.BackendUrl+":"+myId)
	delete(c.backendsMapping, check.BackendUrl)
}

func (c *Cache) MarkBackendDead(check *Check) {
	m, exists := c.backendsMapping[check.BackendUrl]
	if !exists {
		return
	}
	c.redisConn.Transaction(func(mc *redis.MultiCall) {
		for frontendKey, ids := range m {
			deadKey := "dead:" + frontendKey
			for _, id := range ids {
				mc.Sadd(deadKey, id)
			}
			// Better way would be to set the same TTL than Hipache. Not critical
			// since we'll clean the backend list
			mc.Expire(deadKey, 60)
		}
	})
}

func (c *Cache) MarkBackendAlive(check *Check) {
	m, exists := c.backendsMapping[check.BackendUrl]
	if !exists {
		return
	}
	c.redisConn.Transaction(func(mc *redis.MultiCall) {
		for frontendKey, ids := range m {
			for _, id := range ids {
				mc.Srem("dead:"+frontendKey, id)
			}
		}
	})
}

func (c *Cache) ListenToChannel(channel string, callback func(line string)) error {
	// Listening on the "dead" channel to get dead notifications by Hipache
	// Format received on the channel is:
	// -> frontend_key;backend_url;backend_id;number_of_backends
	// Example: "localhost;http://localhost:4242;0;1"
	msgHandler := func(msg *redis.Message) {
		switch msg.Type {
		case redis.MessageMessage:
			callback(msg.Payload)
		}
	}
	sub, err := c.redisConn.Subscription(msgHandler)
	if err != nil {
		return errors.New(fmt.Sprintf("Error: cannot subscribe to "+
			"the \"dead\" channel: %#v", err))
	}
	sub.Subscribe(channel)
	c.redisSub = sub
	return nil
}

func (c *Cache) PingAlive() {
	c.redisConn.Set("hchecker_ping", time.Now().Unix())
}
<commit_msg>small code cleaning<commit_after>package main

import (
	"errors"
	"fmt"
	"github.com/fzzy/radix/redis"
	"time"
)

const (
	REDIS_KEY     = "hchecker"
	REDIS_ADDRESS = "localhost:6379"
)

var (
	redisAddress string
)

type Cache struct {
	redisConn *redis.Client
	redisSub  *redis.Subscription
	// Maintain a mapping between a backends and several frontend
	// -> map[BACKEND_URL][FRONTEND_NAME][BACKEND_ID]
	backendsMapping map[string]map[string][]int
}

func NewCache() (*Cache, error) {
	conf := redis.DefaultConfig()
	conf.Address = redisAddress
	redisConn := redis.NewClient(conf)
	cache := &Cache{
		redisConn:       redisConn,
		backendsMapping: make(map[string]map[string][]int)}
	// We're starting, let's clear any previous meta-data
	// WARNING: This can be a problem if there are several processes sharing
	// the same redis on the same machine. If one of them is restarted, it'll
	// clear the meta-data of everyone...
	redisConn.Del(REDIS_KEY)
	return cache, nil
}

/*
 * Maintain a mapping between Frontends and Backends ID
 */
func (c *Cache) updateFrontendMapping(check *Check) {
	m, exists := c.backendsMapping[check.BackendUrl]
	if !exists {
		m = make(map[string][]int)
	}
	if val, ok := m[check.FrontendKey]; ok {
		for _, item := range val {
			if item == check.BackendId {
				// The backendId has been already stored, ignoring
				return
			}
		}
		// The frontend already exists, adding the new backendId
		m[check.FrontendKey] = append(val, check.BackendId)
	} else {
		// The frontend is new
		m[check.FrontendKey] = []int{check.BackendId}
	}
	c.backendsMapping[check.BackendUrl] = m
}

/*
 * Lock a backend in Redis by its URL
 */
func (c *Cache) LockBackend(check *Check) bool {
	// The syncKey makes sure an entire backend mapping is keep in the same
	// process (we never update a backend mapping from 2 different processes)
	syncKey := check.BackendUrl + ";" + myId
	// Lock the backend with a temporary value, we'll update this with the
	// goroutine signature later
	resp := c.redisConn.Transaction(func(mc *redis.MultiCall) {
		mc.Hsetnx(REDIS_KEY, check.BackendUrl, 1)
		mc.Hexists(REDIS_KEY, syncKey)
	})
	locked, _ := resp.Elems[0].Bool()
	isMine, _ := resp.Elems[1].Bool()
	if locked == false && isMine == false {
		// The backend is being monitored by someone else
		return false
	}
	if locked == true {
		// we got the lock, let's create a unique sig for the goroutine
		t := time.Now()
		// This one is done in the lock, this will garanty that no routine
		// will get the same sig
		sig := fmt.Sprintf("%s;%d.%d", myId, t.Unix(), t.Nanosecond())
		c.redisConn.Hset(REDIS_KEY, check.BackendUrl, sig)
		c.redisConn.Hset(REDIS_KEY, syncKey, 1)
		check.routineSig = sig
	}
	// Let's update the mapping in case this is a new frontend
	c.updateFrontendMapping(check)
	return locked
}

func (c *Cache) IsUnlockedBackend(check *Check) bool {
	// On top of checking the lock, we compare the lock content to make sure
	// we still own the lock
	resp, _ := c.redisConn.Hget(REDIS_KEY, check.BackendUrl).Str()
	return (resp != check.routineSig)
}

func (c *Cache) UnlockBackend(check *Check) {
	c.redisConn.Hdel(REDIS_KEY, check.BackendUrl,
		check.BackendUrl + ":" + myId)
	delete(c.backendsMapping, check.BackendUrl)
}

func (c *Cache) MarkBackendDead(check *Check) {
	m, exists := c.backendsMapping[check.BackendUrl]
	if !exists {
		return
	}
	c.redisConn.Transaction(func(mc *redis.MultiCall) {
		for frontendKey, ids := range m {
			deadKey := "dead:" + frontendKey
			for _, id := range ids {
				mc.Sadd(deadKey, id)
			}
			// Better way would be to set the same TTL than Hipache. Not critical
			// since we'll clean the backend list
			mc.Expire(deadKey, 60)
		}
	})
}

func (c *Cache) MarkBackendAlive(check *Check) {
	m, exists := c.backendsMapping[check.BackendUrl]
	if !exists {
		return
	}
	c.redisConn.Transaction(func(mc *redis.MultiCall) {
		for frontendKey, ids := range m {
			for _, id := range ids {
				mc.Srem("dead:"+frontendKey, id)
			}
		}
	})
}

func (c *Cache) ListenToChannel(channel string, callback func(line string)) error {
	// Listening on the "dead" channel to get dead notifications by Hipache
	// Format received on the channel is:
	// -> frontend_key;backend_url;backend_id;number_of_backends
	// Example: "localhost;http://localhost:4242;0;1"
	msgHandler := func(msg *redis.Message) {
		switch msg.Type {
		case redis.MessageMessage:
			callback(msg.Payload)
		}
	}
	sub, err := c.redisConn.Subscription(msgHandler)
	if err != nil {
		return errors.New(fmt.Sprintf("Error: cannot subscribe to "+
			"the \"dead\" channel: %#v", err))
	}
	sub.Subscribe(channel)
	c.redisSub = sub
	return nil
}

func (c *Cache) PingAlive() {
	c.redisConn.Set("hchecker_ping", time.Now().Unix())
}
<|endoftext|>


<commit_before>package fronted

import (
	"encoding/json"
	"io/ioutil"
	"os"
	"time"
)

var (
	// Nil value indicates end of cache filling
	fillSentinel masquerade
)

func (d *direct) initCaching(cacheFile string) int {
	cache := d.prepopulateMasquerades(cacheFile)
	prevetted := len(cache)
	go d.fillCache(cache, cacheFile)
	return prevetted
}

func (d *direct) prepopulateMasquerades(cacheFile string) []masquerade {
	var cache []masquerade
	file, err := os.Open(cacheFile)
	if err == nil {
		log.Debugf("Attempting to prepopulate masquerades from cache")
		defer file.Close()
		var masquerades []masquerade
		err := json.NewDecoder(file).Decode(&masquerades)
		if err != nil {
			log.Errorf("Error prepopulating cached masquerades: %v", err)
			return cache
		}

		log.Debugf("Cache contained %d masquerades", len(masquerades))
		now := time.Now()
		for _, m := range masquerades {
			if now.Sub(m.LastVetted) < d.maxAllowedCachedAge {
				select {
				case d.masquerades <- m:
					// submitted
					cache = append(cache, m)
				default:
					// channel full, that's okay
				}
			}
		}
	}

	return cache
}

func (d *direct) fillCache(cache []masquerade, cacheFile string) {
	saveTimer := time.NewTimer(d.cacheSaveInterval)
	cacheChanged := false
	for {
		select {
		case m := <-d.toCache:
			if m == fillSentinel {
				log.Debug("Cache closed, stop filling")
				return
			}
			log.Debugf("Caching vetted masquerade for %v (%v)", m.Domain, m.IpAddress)
			cache = append(cache, m)
			cacheChanged = true
		case <-saveTimer.C:
			if !cacheChanged {
				continue
			}
			log.Debug("Saving updated masquerade cache")
			// Truncate cache to max length if necessary
			if len(cache) > d.maxCacheSize {
				truncated := make([]masquerade, d.maxCacheSize)
				copy(truncated, cache[len(cache)-d.maxCacheSize:])
				cache = truncated
			}
			b, err := json.Marshal(cache)
			if err != nil {
				log.Errorf("Unable to marshal cache to JSON: %v", err)
				break
			}
			err = ioutil.WriteFile(cacheFile, b, 0644)
			if err != nil {
				log.Errorf("Unable to save cache to disk: %v", err)
			}
			cacheChanged = false
			saveTimer.Reset(d.cacheSaveInterval)
		}
	}
}

// CloseCache closes any existing file cache.
func CloseCache() {
	_existing, ok := _instance.Get(0)
	if ok && _existing != nil {
		existing := _existing.(*direct)
		log.Debug("Closing cache from existing instance")
		existing.closeCache()
	}
}

func (d *direct) closeCache() {
	d.toCache <- fillSentinel
}
<commit_msg>comment<commit_after>package fronted

import (
	"encoding/json"
	"io/ioutil"
	"os"
	"time"
)

var (
	// zero value indicates end of cache filling
	fillSentinel masquerade
)

func (d *direct) initCaching(cacheFile string) int {
	cache := d.prepopulateMasquerades(cacheFile)
	prevetted := len(cache)
	go d.fillCache(cache, cacheFile)
	return prevetted
}

func (d *direct) prepopulateMasquerades(cacheFile string) []masquerade {
	var cache []masquerade
	file, err := os.Open(cacheFile)
	if err == nil {
		log.Debugf("Attempting to prepopulate masquerades from cache")
		defer file.Close()
		var masquerades []masquerade
		err := json.NewDecoder(file).Decode(&masquerades)
		if err != nil {
			log.Errorf("Error prepopulating cached masquerades: %v", err)
			return cache
		}

		log.Debugf("Cache contained %d masquerades", len(masquerades))
		now := time.Now()
		for _, m := range masquerades {
			if now.Sub(m.LastVetted) < d.maxAllowedCachedAge {
				select {
				case d.masquerades <- m:
					// submitted
					cache = append(cache, m)
				default:
					// channel full, that's okay
				}
			}
		}
	}

	return cache
}

func (d *direct) fillCache(cache []masquerade, cacheFile string) {
	saveTimer := time.NewTimer(d.cacheSaveInterval)
	cacheChanged := false
	for {
		select {
		case m := <-d.toCache:
			if m == fillSentinel {
				log.Debug("Cache closed, stop filling")
				return
			}
			log.Debugf("Caching vetted masquerade for %v (%v)", m.Domain, m.IpAddress)
			cache = append(cache, m)
			cacheChanged = true
		case <-saveTimer.C:
			if !cacheChanged {
				continue
			}
			log.Debug("Saving updated masquerade cache")
			// Truncate cache to max length if necessary
			if len(cache) > d.maxCacheSize {
				truncated := make([]masquerade, d.maxCacheSize)
				copy(truncated, cache[len(cache)-d.maxCacheSize:])
				cache = truncated
			}
			b, err := json.Marshal(cache)
			if err != nil {
				log.Errorf("Unable to marshal cache to JSON: %v", err)
				break
			}
			err = ioutil.WriteFile(cacheFile, b, 0644)
			if err != nil {
				log.Errorf("Unable to save cache to disk: %v", err)
			}
			cacheChanged = false
			saveTimer.Reset(d.cacheSaveInterval)
		}
	}
}

// CloseCache closes any existing file cache.
func CloseCache() {
	_existing, ok := _instance.Get(0)
	if ok && _existing != nil {
		existing := _existing.(*direct)
		log.Debug("Closing cache from existing instance")
		existing.closeCache()
	}
}

func (d *direct) closeCache() {
	d.toCache <- fillSentinel
}
<|endoftext|>


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

package godspeed

import (
	"io"
	"log"
	"net/http"
	"os"
	"strings"

	"github.com/hraban/lrucache"
	"github.com/hraban/mkdtemp"
)

type cacheWrapper struct {
	Basedir string
	idx     *lrucache.Cache
	wrapped http.Handler
}

type diskEntry struct {
	path string
	size int64
}

func (e diskEntry) Size() int64 {
	return e.size
}

func (e diskEntry) OnPurge(why lrucache.PurgeReason) {
	err := os.Remove(e.path)
	if err != nil {
		// Unexpected and seemingly harmless so I don't really care
		log.Print("Failed to remove", e.path, "from godspeed cache:", err)
	}
}

func cachepath(c *cacheWrapper, r *http.Request) string {
	if r.URL.RawQuery != "" || strings.HasSuffix(r.URL.Path, "/") {
		return ""
	}
	host := r.Host
	if host == "" {
		host = "localhost"
	}
	return c.Basedir + host + "/" + r.URL.Path
}

// Uses X-Cache header to determine cacheability (POC)
// Obvious TODO: real HTTP/1.1 conforming caching
func (c *cacheWrapper) ServeHTTP(w http.ResponseWriter, r *http.Request) {
	path := cachepath(c, r)
	if path == "" {
		c.wrapped.ServeHTTP(w, r)
		return
	}
	_, err := c.idx.Get(path)
	if err == nil {
		// Element is cached
		// TODO: Race condition (could be deleted by now)
		// TODO: Breaks (semi transparently) for paths ending with index.html
		w.Header().Set("X-Cache", "Hit")
		http.ServeFile(w, r, path)
		return
	}
	var cachef *os.File
	f := func(w http.ResponseWriter) io.Writer {
		head := w.Header()
		cached := "0"
		defer func() {
			head.Set("X-Cached", cached)
		}()
		if head.Get("X-Cache") == "" {
			return w
		}
		head.Set("X-Cache", "Miss")
		cachedir := dirname(path)
		err := os.MkdirAll(cachedir, 0700)
		if err != nil {
			// Probrem? Just continue as if nothing happened
			log.Printf("Could not create cache file dir %q: %s",
				cachedir, err.Error())
			return w
		}
		cachef, err = os.OpenFile(path, os.O_WRONLY|os.O_CREATE, 0600)
		if err != nil {
			log.Printf("Could not open cache file %q for writing: %s",
				path, err.Error())
			return w
		}
		cached = "1"
		return io.MultiWriter(cachef, w)
	}
	c.wrapped.ServeHTTP(wrapBody(w, f), r)
	if cachef == nil {
		return
	}
	stat, err := cachef.Stat()
	if err != nil {
		log.Printf("Failed to obtain stat info for %q: %v", path, err)
		cachef.Close()
		os.Remove(path)
		return
	}
	err = cachef.Close()
	if err != nil {
		log.Printf("Could not save cache file %q: %v", path, err)
		// No problem
		return
	}
	c.idx.Set(path, diskEntry{path: path, size: stat.Size()})
	return
}

func mustmkdtemp(tmpl string) string {
	dname, err := mkdtemp.Mkdtemp(tmpl)
	if err != nil {
		panic("Failed to create temporary caching directory: " + err.Error())
	}
	return dname
}

// Store cacheable resources on disk. Naive (and non-conforming) implementation
// of a HTTP cache. Note that this is really not transparent caching:
//
// - Headers are not cached
// - As a consequence, encoded data (gzip etc) is served from cache as raw data
// - Upstream status code is completely ignored
// - Caching directives from client are ignored
// - HTTP caching directives from upstream are ignored
// - Non-standard upstream "X-Cache" header is used to determine cacheability
// - Cache hits are served as raw files, introducing some illegal headers
// - and probably more...
//
// All of that notwithstanding, this is a proof of concept worth exploring.
func Cache(h http.Handler) http.Handler {
	return &cacheWrapper{
		Basedir: mustmkdtemp("godspeed-cache-XXXXXXXX") + "/",
		// TODO: Manage maximum size
		idx:     lrucache.New(1 << 20),
		wrapped: h,
	}
}
<commit_msg>Fix formatting of one error logging statement<commit_after>// Copyright © 2013 Hraban Luyat <hraban@0brg.net>
//
// Permission is hereby granted, free of charge, to any person obtaining a copy
// of this software and associated documentation files (the "Software"), to
// deal in the Software without restriction, including without limitation the
// rights to use, copy, modify, merge, publish, distribute, sublicense, and/or
// sell copies of the Software, and to permit persons to whom the Software is
// furnished to do so, subject to the following conditions:
//
// The above copyright notice and this permission notice shall be included in
// all copies or substantial portions of the Software.
//
// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING
// FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS
// IN THE SOFTWARE.

package godspeed

import (
	"io"
	"log"
	"net/http"
	"os"
	"strings"

	"github.com/hraban/lrucache"
	"github.com/hraban/mkdtemp"
)

type cacheWrapper struct {
	Basedir string
	idx     *lrucache.Cache
	wrapped http.Handler
}

type diskEntry struct {
	path string
	size int64
}

func (e diskEntry) Size() int64 {
	return e.size
}

func (e diskEntry) OnPurge(why lrucache.PurgeReason) {
	err := os.Remove(e.path)
	if err != nil {
		// Unexpected and seemingly harmless so I don't really care
		log.Print("Failed to remove ", e.path, " from godspeed cache: ", err)
	}
}

func cachepath(c *cacheWrapper, r *http.Request) string {
	if r.URL.RawQuery != "" || strings.HasSuffix(r.URL.Path, "/") {
		return ""
	}
	host := r.Host
	if host == "" {
		host = "localhost"
	}
	return c.Basedir + host + "/" + r.URL.Path
}

// Uses X-Cache header to determine cacheability (POC)
// Obvious TODO: real HTTP/1.1 conforming caching
func (c *cacheWrapper) ServeHTTP(w http.ResponseWriter, r *http.Request) {
	path := cachepath(c, r)
	if path == "" {
		c.wrapped.ServeHTTP(w, r)
		return
	}
	_, err := c.idx.Get(path)
	if err == nil {
		// Element is cached
		// TODO: Race condition (could be deleted by now)
		// TODO: Breaks (semi transparently) for paths ending with index.html
		w.Header().Set("X-Cache", "Hit")
		http.ServeFile(w, r, path)
		return
	}
	var cachef *os.File
	f := func(w http.ResponseWriter) io.Writer {
		head := w.Header()
		cached := "0"
		defer func() {
			head.Set("X-Cached", cached)
		}()
		if head.Get("X-Cache") == "" {
			return w
		}
		head.Set("X-Cache", "Miss")
		cachedir := dirname(path)
		err := os.MkdirAll(cachedir, 0700)
		if err != nil {
			// Probrem? Just continue as if nothing happened
			log.Printf("Could not create cache file dir %q: %s",
				cachedir, err.Error())
			return w
		}
		cachef, err = os.OpenFile(path, os.O_WRONLY|os.O_CREATE, 0600)
		if err != nil {
			log.Printf("Could not open cache file %q for writing: %s",
				path, err.Error())
			return w
		}
		cached = "1"
		return io.MultiWriter(cachef, w)
	}
	c.wrapped.ServeHTTP(wrapBody(w, f), r)
	if cachef == nil {
		return
	}
	stat, err := cachef.Stat()
	if err != nil {
		log.Printf("Failed to obtain stat info for %q: %v", path, err)
		cachef.Close()
		os.Remove(path)
		return
	}
	err = cachef.Close()
	if err != nil {
		log.Printf("Could not save cache file %q: %v", path, err)
		// No problem
		return
	}
	c.idx.Set(path, diskEntry{path: path, size: stat.Size()})
	return
}

func mustmkdtemp(tmpl string) string {
	dname, err := mkdtemp.Mkdtemp(tmpl)
	if err != nil {
		panic("Failed to create temporary caching directory: " + err.Error())
	}
	return dname
}

// Store cacheable resources on disk. Naive (and non-conforming) implementation
// of a HTTP cache. Note that this is really not transparent caching:
//
// - Headers are not cached
// - As a consequence, encoded data (gzip etc) is served from cache as raw data
// - Upstream status code is completely ignored
// - Caching directives from client are ignored
// - HTTP caching directives from upstream are ignored
// - Non-standard upstream "X-Cache" header is used to determine cacheability
// - Cache hits are served as raw files, introducing some illegal headers
// - and probably more...
//
// All of that notwithstanding, this is a proof of concept worth exploring.
func Cache(h http.Handler) http.Handler {
	return &cacheWrapper{
		Basedir: mustmkdtemp("godspeed-cache-XXXXXXXX") + "/",
		// TODO: Manage maximum size
		idx:     lrucache.New(1 << 20),
		wrapped: h,
	}
}
<|endoftext|>


<commit_before>package main

import (
	"bufio"
	"fmt"
	"io"
	"log"
	"os"
	"strings"
	"syscall"
	"time"
)

func main() {
	parentpid, err := os.FindProcess(os.Getppid())
	if err != nil {
		log.Fatal(err)
	}

	// check last ping and reconnect if needed
	ping := time.Now()
	go func() {
		for {
			if time.Since(ping) > 255*time.Second {
				parentpid.Signal(syscall.SIGUSR1)
			}
			time.Sleep(100 * time.Second)
		}
	}()

	bio := bufio.NewReader(os.Stdin)

	for {
		line, err := bio.ReadString('\n')
		if err != nil {
			if err == io.EOF {
				parentpid.Signal(syscall.SIGUSR1)
			} else {
				fmt.Fprintln(os.Stderr, err)
			}
			time.Sleep(2 * time.Second)
		} else {
			if len(line) > 4 && line[:4] == "PING" {
				fmt.Printf(strings.Replace(line, "PING", "PONG", 1))
				ping = time.Now()
			}
		}
	}
}
<commit_msg>ping plugin: reduce import list<commit_after>package main

import (
	"bufio"
	"fmt"
	"io"
	"os"
	"syscall"
	"time"
)

func main() {
	parentpid, err := os.FindProcess(os.Getppid())
	if err != nil {
		fmt.Fprintln(os.Stderr, err)
		os.Exit(2)
	}

	// check last ping and reconnect if needed
	ping := time.Now()
	go func() {
		for {
			if time.Since(ping) > 255*time.Second {
				parentpid.Signal(syscall.SIGUSR1)
			}
			time.Sleep(100 * time.Second)
		}
	}()

	bio := bufio.NewReader(os.Stdin)

	for {
		line, err := bio.ReadString('\n')
		if err != nil {
			if err == io.EOF {
				parentpid.Signal(syscall.SIGUSR1)
			} else {
				fmt.Fprintln(os.Stderr, err)
			}
			time.Sleep(2 * time.Second)
		} else {
			if len(line) > 4 && line[:4] == "PING" {
				fmt.Printf(fmt.Sprintf("PONG%s", line[4:]))
				ping = time.Now()
			}
		}
	}
}
<|endoftext|>


<commit_before>// Copyright 2015 CoreOS, Inc.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package ignition

import (
	"fmt"
	"strings"

	"github.com/coreos/mantle/kola/register"
	"github.com/coreos/mantle/platform"
)

func init() {
	register.Register(&register.Test{
		Name:        "coreos.metadata.aws",
		Run:         verifyAWS,
		ClusterSize: 1,
		Platforms:   []string{"aws"},
		UserData: `{
		               "ignitionVersion": 1,
		               "systemd": {
		                   "units": [
		                       {
		                           "name": "coreos-metadata.service",
		                           "enable": true
		                       },
		                       {
		                           "name": "metadata.target",
		                           "enable": true,
		                           "contents": "[Install]\nWantedBy=multi-user.target"
		                       }
		                   ]
		               }
		           }`,
	})

	register.Register(&register.Test{
		Name:        "coreos.metadata.azure",
		Run:         verifyAzure,
		ClusterSize: 1,
		Platforms:   []string{"azure"},
		UserData: `{
		               "ignitionVersion": 1,
		               "systemd": {
		                   "units": [
		                       {
		                           "name": "coreos-metadata.service",
		                           "enable": true
		                       },
		                       {
		                           "name": "metadata.target",
		                           "enable": true,
		                           "contents": "[Install]\nWantedBy=multi-user.target"
		                       }
		                   ]
		               }
		           }`,
	})
}

func verifyAWS(c platform.TestCluster) error {
	return verify(c, "COREOS_IPV4_LOCAL", "COREOS_IPV4_PUBLIC", "COREOS_HOSTNAME")
}

func verifyAzure(c platform.TestCluster) error {
	return verify(c, "COREOS_IPV4_LOCAL", "COREOS_IPV4_PUBLIC")
}

func verify(c platform.TestCluster, keys ...string) error {
	m := c.Machines()[0]

	out, err := m.SSH("cat /run/metadata/coreos")
	if err != nil {
		return fmt.Errorf("failed to cat /run/metadata/coreos: %s: %v", out, err)
	}

	for _, key := range keys {
		if !strings.Contains(string(out), key) {
			return fmt.Errorf("%q wasn't found in %q", key, string(out))
		}
	}

	return nil
}
<commit_msg>kola/test/metadata: add support for new names<commit_after>// Copyright 2015 CoreOS, Inc.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package ignition

import (
	"fmt"
	"strings"

	"github.com/coreos/mantle/kola/register"
	"github.com/coreos/mantle/platform"

	"github.com/coreos/mantle/Godeps/_workspace/src/github.com/coreos/go-semver/semver"
)

func init() {
	register.Register(&register.Test{
		Name:        "coreos.metadata.aws",
		Run:         verifyAWS,
		ClusterSize: 1,
		Platforms:   []string{"aws"},
		UserData: `{
		               "ignitionVersion": 1,
		               "systemd": {
		                   "units": [
		                       {
		                           "name": "coreos-metadata.service",
		                           "enable": true
		                       },
		                       {
		                           "name": "metadata.target",
		                           "enable": true,
		                           "contents": "[Install]\nWantedBy=multi-user.target"
		                       }
		                   ]
		               }
		           }`,
	})

	register.Register(&register.Test{
		Name:        "coreos.metadata.azure",
		Run:         verifyAzure,
		ClusterSize: 1,
		Platforms:   []string{"azure"},
		UserData: `{
		               "ignitionVersion": 1,
		               "systemd": {
		                   "units": [
		                       {
		                           "name": "coreos-metadata.service",
		                           "enable": true
		                       },
		                       {
		                           "name": "metadata.target",
		                           "enable": true,
		                           "contents": "[Install]\nWantedBy=multi-user.target"
		                       }
		                   ]
		               }
		           }`,
	})
}

func verifyAWS(c platform.TestCluster) error {
	m := c.Machines()[0]

	out, err := m.SSH("coreos-metadata --version")
	if err != nil {
		return fmt.Errorf("failed to cat /run/metadata/coreos: %s: %v", out, err)
	}

	versionStr := strings.TrimPrefix(string(out), "coreos-metadata v")
	version, err := semver.NewVersion(versionStr)
	if err != nil {
		return fmt.Errorf("failed to parse coreos-metadata version: %v", err)
	}

	if version.LessThan(semver.Version{Minor: 3}) {
		return verify(c, "COREOS_IPV4_LOCAL", "COREOS_IPV4_PUBLIC", "COREOS_HOSTNAME")
	} else {
		return verify(c, "COREOS_EC2_IPV4_LOCAL", "COREOS_EC2_IPV4_PUBLIC", "COREOS_EC2_HOSTNAME")
	}
}

func verifyAzure(c platform.TestCluster) error {
	return verify(c, "COREOS_AZURE_IPV4_DYNAMIC", "COREOS_AZURE_IPV4_VIRTUAL")
}

func verify(c platform.TestCluster, keys ...string) error {
	m := c.Machines()[0]

	out, err := m.SSH("cat /run/metadata/coreos")
	if err != nil {
		return fmt.Errorf("failed to cat /run/metadata/coreos: %s: %v", out, err)
	}

	for _, key := range keys {
		if !strings.Contains(string(out), key) {
			return fmt.Errorf("%q wasn't found in %q", key, string(out))
		}
	}

	return nil
}
<|endoftext|>


<commit_before>package baudio

import (
	"bytes"
	"encoding/binary"
	"fmt"
	"io"
	"math"
	//"os"
	"os/exec"
	"runtime"
	"strconv"
	//"time"
)

const (
	FuncValueTypeFloat    = 0
	FuncValueTypeNotFloat = 1
)

type BChannel struct {
	funcValueType int
	funcs         []func(float64, int) float64
}

func newBChannel(fvt int) *BChannel {
	bc := &BChannel{
		funcValueType: fvt,
		funcs:         make([]func(float64, int) float64, 0),
	}
	return bc
}

func (bc *BChannel) push(fn func(float64, int) float64) {
	bc.funcs = append(bc.funcs, fn)
}

type bOptions struct {
	size int
	rate int
}

func NewBOptions() *bOptions {
	return &bOptions{
		size: 2048,
		rate: 44000,
	}
}

type B struct {
	readable   bool
	size       int
	rate       int
	t          float64
	i          int
	paused     bool
	ended      bool
	destroyed  bool
	channels   []*BChannel
	chEnd      chan bool
	chEndSox   chan bool
	chResume   chan func()
	chNextTick chan bool
	pipeReader *io.PipeReader
	pipeWriter *io.PipeWriter
}

func New(opts *bOptions, fn func(float64, int) float64) *B {
	b := &B{
		readable:   true,
		size:       2048,
		rate:       44000,
		t:          0,
		i:          0,
		paused:     false,
		ended:      false,
		destroyed:  false,
		chEnd:      make(chan bool),
		chEndSox:   make(chan bool),
		chResume:   make(chan func()),
		chNextTick: make(chan bool),
	}
	b.pipeReader, b.pipeWriter = io.Pipe()
	if opts != nil {
		b.size = opts.size
		b.rate = opts.rate
	}
	if fn != nil {
		b.Push(fn)
	}
	go func() {
		if b.paused {
			b.chResume <- func() {
				go b.loop()
				b.main()
			}
		} else {
			go b.loop()
			b.main()
		}
	}()
	//go b.loop()
	return b
}

func (b *B) main() {
	for {
		// 2013-02-28 koyachi ここで何かしないとループまわらないのなぜ
		// => fmt.PrinfすることでnodeのnextTick的なものがつまれててそのうちPlay()のread待ちまで進めるのでは。
		//L1:
		//fmt.Println("main loop header")
		//fmt.Printf(".")
		//time.Sleep(1 * time.Millisecond)
		runtime.Gosched()
		select {
		case <-b.chEnd:
			fmt.Println("main chEnd")
			b.terminateMain()
			break
		case fn := <-b.chResume:
			//fmt.Println("main chResume")
			fn()
		case <-b.chNextTick:
			//fmt.Println("main chNextTick")
			go b.loop()
			//b.loop()
		default:
			//fmt.Println("main default")
			//go b.loop()
			//goto L1
		}
	}
}

func (b *B) terminateMain() {
	b.pipeWriter.Close()
	b.ended = true
	b.chEndSox <- true
}

func (b *B) End() {
	b.ended = true
}

func (b *B) Destroy() {
	b.destroyed = true
	b.chEnd <- true
}

func (b *B) Pause() {
	b.paused = true
}

func (b *B) Resume() {
	if !b.paused {
		return
	}
	b.paused = false
	b.chResume <- func() {}
}

func (b *B) AddChannel(funcValueType int, fn func(float64, int) float64) {
	bc := newBChannel(funcValueType)
	bc.push(fn)
	b.channels = append(b.channels, bc)
}

func (b *B) PushTo(index int, fn func(float64, int) float64) {
	if len(b.channels) <= index {
		bc := newBChannel(FuncValueTypeFloat)
		b.channels = append(b.channels, bc)
	}
	b.channels[index].funcs = append(b.channels[index].funcs, fn)
}

func (b *B) Push(fn func(float64, int) float64) {
	b.PushTo(len(b.channels), fn)
}

func (b *B) loop() {
	buf := b.tick()
	if b.destroyed {
		// no more events
		//fmt.Println("loop destroyed")
	} else if b.paused {
		//fmt.Println("loop paused")
		b.chResume <- func() {
			b.pipeWriter.Write(buf.Bytes())
			b.chNextTick <- true
		}
	} else {
		//fmt.Println("loop !(destroyed || paused)")
		b.pipeWriter.Write(buf.Bytes())
		if b.ended {
			//fmt.Println("loop ended")
			b.chEnd <- true
		} else {
			//fmt.Println("loop !ended")
			b.chNextTick <- true
		}
	}
}

func (b *B) tick() *bytes.Buffer {
	bufSize := b.size * len(b.channels)
	byteBuffer := make([]byte, 0)
	buf := bytes.NewBuffer(byteBuffer)
	for i := 0; i < bufSize; i += 2 {
		lrIndex := int(i / 2)
		lenCh := len(b.channels)
		ch := b.channels[lrIndex%lenCh]
		t := float64(b.t) + math.Floor(float64(lrIndex))/float64(b.rate)/float64(lenCh)
		counter := b.i + int(math.Floor(float64(lrIndex)/float64(lenCh)))

		value := float64(0)
		n := float64(0)
		for j := 0; j < len(ch.funcs); j++ {
			x := ch.funcs[j](float64(t), counter)
			n += x
		}
		n /= float64(len(ch.funcs))

		if ch.funcValueType == FuncValueTypeFloat {
			value = signed(n)
		} else {
			b_ := math.Pow(2, float64(ch.funcValueType))
			x := math.Mod(math.Floor(n), b_) / b_ * math.Pow(2, 15)
			value = x
		}
		if err := binary.Write(buf, binary.LittleEndian, int16(clamp(value))); err != nil {
			panic(err)
		}
	}
	b.i += b.size / 2
	b.t += float64(b.size) / float64(2) / float64(b.rate)
	return buf
}

func clamp(x float64) float64 {
	return math.Max(math.Min(x, math.Pow(2, 15)-1), -math.Pow(2, 15))
}

func signed(n float64) float64 {
	b := math.Pow(2, 15)
	if n > 0 {
		return math.Min(b-1, math.Floor(b*n-1))
	}
	return math.Max(-b, math.Ceil(b*n-1))
}

func mergeArgs(opts, args map[string]string) []string {
	for k, _ := range opts {
		args[k] = opts[k]
	}
	var resultsLast []string
	var results []string
	for k, _ := range args {
		switch k {
		case "-":
			resultsLast = append(resultsLast, k)
		case "-o":
			resultsLast = append(resultsLast, k, args[k])
		default:
			var dash string
			if len(k) == 1 {
				dash = "-"
			} else {
				dash = "--"
			}
			results = append(results, dash+k, args[k])
		}
	}
	results = append(results, resultsLast...)
	fmt.Printf("results = %v\n", results)
	return results
}

func (b *B) runCommand(command string, mergedArgs []string) {
	cmd := exec.Command(command, mergedArgs...)
	stdin, err := cmd.StdinPipe()
	if err != nil {
		panic(err)
	}
	defer func() {
		fmt.Println("runCommand: before stdin.Close()")
		stdin.Close()
	}()
	var out bytes.Buffer
	cmd.Stdout = &out
	// TODO: option
	//cmd.Stdout = os.Stdout
	//cmd.Stderr = os.Stderr
	if err := cmd.Start(); err != nil {
		panic(err)
	}
	defer func() {
		if p := cmd.Process; p != nil {
			fmt.Println("runCommand: before p.Kill()")
			p.Kill()
		}
	}()

	readBuf := make([]byte, b.size*len(b.channels))
	for {
		//fmt.Println("play loop header")
		if _, err := b.pipeReader.Read(readBuf); err != nil {
			panic(err)
		}
		if _, err = stdin.Write(readBuf); err != nil {
			// TODO: more better error handling
			if err.Error() == "write |1: broken pipe" {
				fmt.Printf("ERR: stdin.Write(readBuf): err = %v\n", err)
				runtime.Gosched()
				break
			}
			panic(err)
		}
	}
}

func (b *B) Play(opts map[string]string) {
	go b.runCommand("play", mergeArgs(opts, map[string]string{
		"c": strconv.Itoa(len(b.channels)),
		"r": strconv.Itoa(b.rate),
		"t": "s16",
		"-": "DUMMY",
	}))
	<-b.chEndSox
	b.pipeReader.Close()
}

func (b *B) Record(file string, opts map[string]string) {
	go b.runCommand("sox", mergeArgs(opts, map[string]string{
		"c":  strconv.Itoa(len(b.channels)),
		"r":  strconv.Itoa(b.rate),
		"t":  "s16",
		"-":  "DUMMY",
		"-o": file,
	}))
	<-b.chEndSox
	b.pipeReader.Close()
}
<commit_msg>fix BOptions.<commit_after>package baudio

import (
	"bytes"
	"encoding/binary"
	"fmt"
	"io"
	"math"
	//"os"
	"os/exec"
	"runtime"
	"strconv"
	//"time"
)

const (
	FuncValueTypeFloat    = 0
	FuncValueTypeNotFloat = 1
)

type BChannel struct {
	funcValueType int
	funcs         []func(float64, int) float64
}

func newBChannel(fvt int) *BChannel {
	bc := &BChannel{
		funcValueType: fvt,
		funcs:         make([]func(float64, int) float64, 0),
	}
	return bc
}

func (bc *BChannel) push(fn func(float64, int) float64) {
	bc.funcs = append(bc.funcs, fn)
}

type BOptions struct {
	Size int
	Rate int
}

func NewBOptions() *BOptions {
	return &BOptions{
		Size: 2048,
		Rate: 44000,
	}
}

type B struct {
	readable   bool
	size       int
	rate       int
	t          float64
	i          int
	paused     bool
	ended      bool
	destroyed  bool
	channels   []*BChannel
	chEnd      chan bool
	chEndSox   chan bool
	chResume   chan func()
	chNextTick chan bool
	pipeReader *io.PipeReader
	pipeWriter *io.PipeWriter
}

func New(opts *BOptions, fn func(float64, int) float64) *B {
	b := &B{
		readable:   true,
		size:       2048,
		rate:       44000,
		t:          0,
		i:          0,
		paused:     false,
		ended:      false,
		destroyed:  false,
		chEnd:      make(chan bool),
		chEndSox:   make(chan bool),
		chResume:   make(chan func()),
		chNextTick: make(chan bool),
	}
	b.pipeReader, b.pipeWriter = io.Pipe()
	if opts != nil {
		b.size = opts.Size
		b.rate = opts.Rate
	}
	if fn != nil {
		b.Push(fn)
	}
	go func() {
		if b.paused {
			b.chResume <- func() {
				go b.loop()
				b.main()
			}
		} else {
			go b.loop()
			b.main()
		}
	}()
	//go b.loop()
	return b
}

func (b *B) main() {
	for {
		// 2013-02-28 koyachi ここで何かしないとループまわらないのなぜ
		// => fmt.PrinfすることでnodeのnextTick的なものがつまれててそのうちPlay()のread待ちまで進めるのでは。
		//L1:
		//fmt.Println("main loop header")
		//fmt.Printf(".")
		//time.Sleep(1 * time.Millisecond)
		runtime.Gosched()
		select {
		case <-b.chEnd:
			fmt.Println("main chEnd")
			b.terminateMain()
			break
		case fn := <-b.chResume:
			//fmt.Println("main chResume")
			fn()
		case <-b.chNextTick:
			//fmt.Println("main chNextTick")
			go b.loop()
			//b.loop()
		default:
			//fmt.Println("main default")
			//go b.loop()
			//goto L1
		}
	}
}

func (b *B) terminateMain() {
	b.pipeWriter.Close()
	b.ended = true
	b.chEndSox <- true
}

func (b *B) End() {
	b.ended = true
}

func (b *B) Destroy() {
	b.destroyed = true
	b.chEnd <- true
}

func (b *B) Pause() {
	b.paused = true
}

func (b *B) Resume() {
	if !b.paused {
		return
	}
	b.paused = false
	b.chResume <- func() {}
}

func (b *B) AddChannel(funcValueType int, fn func(float64, int) float64) {
	bc := newBChannel(funcValueType)
	bc.push(fn)
	b.channels = append(b.channels, bc)
}

func (b *B) PushTo(index int, fn func(float64, int) float64) {
	if len(b.channels) <= index {
		bc := newBChannel(FuncValueTypeFloat)
		b.channels = append(b.channels, bc)
	}
	b.channels[index].funcs = append(b.channels[index].funcs, fn)
}

func (b *B) Push(fn func(float64, int) float64) {
	b.PushTo(len(b.channels), fn)
}

func (b *B) loop() {
	buf := b.tick()
	if b.destroyed {
		// no more events
		//fmt.Println("loop destroyed")
	} else if b.paused {
		//fmt.Println("loop paused")
		b.chResume <- func() {
			b.pipeWriter.Write(buf.Bytes())
			b.chNextTick <- true
		}
	} else {
		//fmt.Println("loop !(destroyed || paused)")
		b.pipeWriter.Write(buf.Bytes())
		if b.ended {
			//fmt.Println("loop ended")
			b.chEnd <- true
		} else {
			//fmt.Println("loop !ended")
			b.chNextTick <- true
		}
	}
}

func (b *B) tick() *bytes.Buffer {
	bufSize := b.size * len(b.channels)
	byteBuffer := make([]byte, 0)
	buf := bytes.NewBuffer(byteBuffer)
	for i := 0; i < bufSize; i += 2 {
		lrIndex := int(i / 2)
		lenCh := len(b.channels)
		ch := b.channels[lrIndex%lenCh]
		t := float64(b.t) + math.Floor(float64(lrIndex))/float64(b.rate)/float64(lenCh)
		counter := b.i + int(math.Floor(float64(lrIndex)/float64(lenCh)))

		value := float64(0)
		n := float64(0)
		for j := 0; j < len(ch.funcs); j++ {
			x := ch.funcs[j](float64(t), counter)
			n += x
		}
		n /= float64(len(ch.funcs))

		if ch.funcValueType == FuncValueTypeFloat {
			value = signed(n)
		} else {
			b_ := math.Pow(2, float64(ch.funcValueType))
			x := math.Mod(math.Floor(n), b_) / b_ * math.Pow(2, 15)
			value = x
		}
		if err := binary.Write(buf, binary.LittleEndian, int16(clamp(value))); err != nil {
			panic(err)
		}
	}
	b.i += b.size / 2
	b.t += float64(b.size) / float64(2) / float64(b.rate)
	return buf
}

func clamp(x float64) float64 {
	return math.Max(math.Min(x, math.Pow(2, 15)-1), -math.Pow(2, 15))
}

func signed(n float64) float64 {
	b := math.Pow(2, 15)
	if n > 0 {
		return math.Min(b-1, math.Floor(b*n-1))
	}
	return math.Max(-b, math.Ceil(b*n-1))
}

func mergeArgs(opts, args map[string]string) []string {
	for k, _ := range opts {
		args[k] = opts[k]
	}
	var resultsLast []string
	var results []string
	for k, _ := range args {
		switch k {
		case "-":
			resultsLast = append(resultsLast, k)
		case "-o":
			resultsLast = append(resultsLast, k, args[k])
		default:
			var dash string
			if len(k) == 1 {
				dash = "-"
			} else {
				dash = "--"
			}
			results = append(results, dash+k, args[k])
		}
	}
	results = append(results, resultsLast...)
	fmt.Printf("results = %v\n", results)
	return results
}

func (b *B) runCommand(command string, mergedArgs []string) {
	cmd := exec.Command(command, mergedArgs...)
	stdin, err := cmd.StdinPipe()
	if err != nil {
		panic(err)
	}
	defer func() {
		fmt.Println("runCommand: before stdin.Close()")
		stdin.Close()
	}()
	var out bytes.Buffer
	cmd.Stdout = &out
	// TODO: option
	//cmd.Stdout = os.Stdout
	//cmd.Stderr = os.Stderr
	if err := cmd.Start(); err != nil {
		panic(err)
	}
	defer func() {
		if p := cmd.Process; p != nil {
			fmt.Println("runCommand: before p.Kill()")
			p.Kill()
		}
	}()

	readBuf := make([]byte, b.size*len(b.channels))
	for {
		//fmt.Println("play loop header")
		if _, err := b.pipeReader.Read(readBuf); err != nil {
			panic(err)
		}
		if _, err = stdin.Write(readBuf); err != nil {
			// TODO: more better error handling
			if err.Error() == "write |1: broken pipe" {
				fmt.Printf("ERR: stdin.Write(readBuf): err = %v\n", err)
				runtime.Gosched()
				break
			}
			panic(err)
		}
	}
}

func (b *B) Play(opts map[string]string) {
	go b.runCommand("play", mergeArgs(opts, map[string]string{
		"c": strconv.Itoa(len(b.channels)),
		"r": strconv.Itoa(b.rate),
		"t": "s16",
		"-": "DUMMY",
	}))
	<-b.chEndSox
	b.pipeReader.Close()
}

func (b *B) Record(file string, opts map[string]string) {
	go b.runCommand("sox", mergeArgs(opts, map[string]string{
		"c":  strconv.Itoa(len(b.channels)),
		"r":  strconv.Itoa(b.rate),
		"t":  "s16",
		"-":  "DUMMY",
		"-o": file,
	}))
	<-b.chEndSox
	b.pipeReader.Close()
}
<|endoftext|>


<commit_before>package baudio

import (
	"bytes"
	"encoding/binary"
	"fmt"
	"io"
	"math"
	"runtime"
	"strconv"
	//"time"
)

const (
	FuncValueTypeFloat    = 0
	FuncValueTypeNotFloat = 1
)

type BChannel struct {
	funcValueType int
	funcs         []func(float64, int) float64
}

func newBChannel(fvt int) *BChannel {
	bc := &BChannel{
		funcValueType: fvt,
		funcs:         make([]func(float64, int) float64, 0),
	}
	return bc
}

func (bc *BChannel) push(fn func(float64, int) float64) {
	bc.funcs = append(bc.funcs, fn)
}

type BOptions struct {
	Size int
	Rate int
}

func NewBOptions() *BOptions {
	return &BOptions{
		Size: 2048,
		Rate: 44000,
	}
}

type B struct {
	readable   bool
	size       int
	rate       int
	t          float64
	i          int
	paused     bool
	ended      bool
	destroyed  bool
	channels   []*BChannel
	chEnd      chan bool
	chEndSox   chan bool
	chResume   chan func()
	chNextTick chan bool
	pipeReader *io.PipeReader
	pipeWriter *io.PipeWriter
	sox        *Sox
}

func New(opts *BOptions, fn func(float64, int) float64) *B {
	b := &B{
		readable:   true,
		size:       2048,
		rate:       44000,
		t:          0,
		i:          0,
		paused:     false,
		ended:      false,
		destroyed:  false,
		chEnd:      make(chan bool),
		chEndSox:   make(chan bool),
		chResume:   make(chan func()),
		chNextTick: make(chan bool),
	}
	b.pipeReader, b.pipeWriter = io.Pipe()
	if opts != nil {
		b.size = opts.Size
		b.rate = opts.Rate
	}
	if fn != nil {
		b.Push(fn)
	}
	go func() {
		if b.paused {
			b.chResume <- func() {
				go b.loop()
				b.main()
			}
		} else {
			go b.loop()
			b.main()
		}
	}()
	//go b.loop()
	return b
}

func (b *B) main() {
	for {
		// 2013-02-28 koyachi ここで何かしないとループまわらないのなぜ
		// => fmt.PrinfすることでnodeのnextTick的なものがつまれててそのうちPlay()のread待ちまで進めるのでは。
		//L1:
		//fmt.Println("main loop header")
		//fmt.Printf(".")
		//time.Sleep(1 * time.Millisecond)
		runtime.Gosched()
		select {
		case <-b.chEnd:
			fmt.Println("main chEnd")
			b.terminateMain()
			break
		case fn := <-b.chResume:
			//fmt.Println("main chResume")
			fn()
		case <-b.chNextTick:
			//fmt.Println("main chNextTick")
			go b.loop()
			//b.loop()
		default:
			//fmt.Println("main default")
			//go b.loop()
			//goto L1
		}
	}
}

func (b *B) terminateMain() {
	b.pipeWriter.Close()
	b.ended = true
	b.chEndSox <- true
}

func (b *B) End() {
	b.ended = true
}

func (b *B) Destroy() {
	b.destroyed = true
	b.chEnd <- true
}

func (b *B) Pause() {
	b.paused = true
}

func (b *B) Resume() {
	if !b.paused {
		return
	}
	b.paused = false
	b.chResume <- func() {}
}

func (b *B) AddChannel(funcValueType int, fn func(float64, int) float64) {
	bc := newBChannel(funcValueType)
	bc.push(fn)
	b.channels = append(b.channels, bc)
}

func (b *B) Push(fn func(float64, int) float64) {
	index := len(b.channels)
	if len(b.channels) <= index {
		bc := newBChannel(FuncValueTypeFloat)
		b.channels = append(b.channels, bc)
	}
	b.channels[index].funcs = append(b.channels[index].funcs, fn)
}

func (b *B) loop() {
	buf := b.tick()
	if b.destroyed {
		// no more events
		//fmt.Println("loop destroyed")
	} else if b.paused {
		//fmt.Println("loop paused")
		b.chResume <- func() {
			b.pipeWriter.Write(buf.Bytes())
			b.chNextTick <- true
		}
	} else {
		//fmt.Println("loop !(destroyed || paused)")
		b.pipeWriter.Write(buf.Bytes())
		if b.ended {
			//fmt.Println("loop ended")
			b.chEnd <- true
		} else {
			//fmt.Println("loop !ended")
			b.chNextTick <- true
		}
	}
}

func (b *B) tick() *bytes.Buffer {
	bufSize := b.size * len(b.channels)
	byteBuffer := make([]byte, 0)
	buf := bytes.NewBuffer(byteBuffer)
	for i := 0; i < bufSize; i += 2 {
		lrIndex := int(i / 2)
		lenCh := len(b.channels)
		ch := b.channels[lrIndex%lenCh]
		t := float64(b.t) + math.Floor(float64(lrIndex))/float64(b.rate)/float64(lenCh)
		counter := b.i + int(math.Floor(float64(lrIndex)/float64(lenCh)))

		value := float64(0)
		n := float64(0)
		for j := 0; j < len(ch.funcs); j++ {
			x := ch.funcs[j](float64(t), counter)
			n += x
		}
		n /= float64(len(ch.funcs))

		if ch.funcValueType == FuncValueTypeFloat {
			value = signed(n)
		} else {
			b_ := math.Pow(2, float64(ch.funcValueType))
			x := math.Mod(math.Floor(n), b_) / b_ * math.Pow(2, 15)
			value = x
		}
		if err := binary.Write(buf, binary.LittleEndian, int16(clamp(value))); err != nil {
			panic(err)
		}
	}
	b.i += b.size / 2
	b.t += float64(b.size) / float64(2) / float64(b.rate)
	return buf
}

func clamp(x float64) float64 {
	return math.Max(math.Min(x, math.Pow(2, 15)-1), -math.Pow(2, 15))
}

func signed(n float64) float64 {
	b := math.Pow(2, 15)
	if n > 0 {
		return math.Min(b-1, math.Floor(b*n-1))
	}
	return math.Max(-b, math.Ceil(b*n-1))
}

func (b *B) Play(opts map[string]string) {
	go SoxPlay(mergeArgs(opts, map[string]string{
		"c": strconv.Itoa(len(b.channels)),
		"r": strconv.Itoa(b.rate),
		"t": "s16",
		"-": "DUMMY",
	}), b.waveReceiver())
	<-b.chEndSox
	b.pipeReader.Close()
}

func (b *B) Record(file string, opts map[string]string) {
	go SoxRecord(file, mergeArgs(opts, map[string]string{
		"c": strconv.Itoa(len(b.channels)),
		"r": strconv.Itoa(b.rate),
		"t": "s16",
		"-": "DUMMY",
	}), b.waveReceiver())
	<-b.chEndSox
	b.pipeReader.Close()
}

func (b *B) waveReceiver() func() []byte {
	readBuf := make([]byte, b.size*len(b.channels))
	return func() []byte {
		if _, err := b.pipeReader.Read(readBuf); err != nil {
			panic(err)
		}
		return readBuf
	}
}
<commit_msg>add comment.<commit_after>package baudio

import (
	"bytes"
	"encoding/binary"
	"fmt"
	"io"
	"math"
	"runtime"
	"strconv"
	//"time"
)

const (
	FuncValueTypeFloat    = 0
	FuncValueTypeNotFloat = 1
)

type BChannel struct {
	funcValueType int
	funcs         []func(float64, int) float64
}

func newBChannel(fvt int) *BChannel {
	bc := &BChannel{
		funcValueType: fvt,
		funcs:         make([]func(float64, int) float64, 0),
	}
	return bc
}

func (bc *BChannel) push(fn func(float64, int) float64) {
	bc.funcs = append(bc.funcs, fn)
}

type BOptions struct {
	Size int
	Rate int
}

func NewBOptions() *BOptions {
	return &BOptions{
		Size: 2048,
		Rate: 44000,
	}
}

type B struct {
	readable   bool
	size       int
	rate       int
	t          float64
	i          int
	paused     bool
	ended      bool
	destroyed  bool
	channels   []*BChannel
	chEnd      chan bool
	chEndSox   chan bool
	chResume   chan func()
	chNextTick chan bool
	pipeReader *io.PipeReader
	pipeWriter *io.PipeWriter
	sox        *Sox
}

func New(opts *BOptions, fn func(float64, int) float64) *B {
	b := &B{
		readable:   true,
		size:       2048,
		rate:       44000,
		t:          0,
		i:          0,
		paused:     false,
		ended:      false,
		destroyed:  false,
		chEnd:      make(chan bool),
		chEndSox:   make(chan bool),
		chResume:   make(chan func()),
		chNextTick: make(chan bool),
	}
	b.pipeReader, b.pipeWriter = io.Pipe()
	if opts != nil {
		b.size = opts.Size
		b.rate = opts.Rate
	}
	if fn != nil {
		b.Push(fn)
	}
	go func() {
		if b.paused {
			b.chResume <- func() {
				go b.loop()
				b.main()
			}
		} else {
			go b.loop()
			b.main()
		}
	}()
	//go b.loop()
	return b
}

func (b *B) main() {
	for {
		// 2013-02-28 koyachi ここで何かしないとループまわらないのなぜ
		// => fmt.PrinfすることでnodeのnextTick的なものがつまれててそのうちPlay()のread待ちまで進めるのでは。
		//L1:
		//fmt.Println("main loop header")
		//fmt.Printf(".")
		//time.Sleep(1 * time.Millisecond)
		runtime.Gosched()
		select {
		case <-b.chEnd:
			fmt.Println("main chEnd")
			b.terminateMain()
			break
		case fn := <-b.chResume:
			//fmt.Println("main chResume")
			fn()
		case <-b.chNextTick:
			//fmt.Println("main chNextTick")
			go b.loop()
			//b.loop()
		default:
			//fmt.Println("main default")
			//go b.loop()
			//goto L1
		}
	}
}

func (b *B) terminateMain() {
	b.pipeWriter.Close()
	b.ended = true
	b.chEndSox <- true
}

// TODO: To Go Style (end,destroy,pause,resume are node.js's Stream interface.)
func (b *B) End() {
	b.ended = true
}

func (b *B) Destroy() {
	b.destroyed = true
	b.chEnd <- true
}

func (b *B) Pause() {
	b.paused = true
}

func (b *B) Resume() {
	if !b.paused {
		return
	}
	b.paused = false
	b.chResume <- func() {}
}

func (b *B) AddChannel(funcValueType int, fn func(float64, int) float64) {
	bc := newBChannel(funcValueType)
	bc.push(fn)
	b.channels = append(b.channels, bc)
}

func (b *B) Push(fn func(float64, int) float64) {
	index := len(b.channels)
	if len(b.channels) <= index {
		bc := newBChannel(FuncValueTypeFloat)
		b.channels = append(b.channels, bc)
	}
	b.channels[index].funcs = append(b.channels[index].funcs, fn)
}

func (b *B) loop() {
	buf := b.tick()
	if b.destroyed {
		// no more events
		//fmt.Println("loop destroyed")
	} else if b.paused {
		//fmt.Println("loop paused")
		b.chResume <- func() {
			b.pipeWriter.Write(buf.Bytes())
			b.chNextTick <- true
		}
	} else {
		//fmt.Println("loop !(destroyed || paused)")
		b.pipeWriter.Write(buf.Bytes())
		if b.ended {
			//fmt.Println("loop ended")
			b.chEnd <- true
		} else {
			//fmt.Println("loop !ended")
			b.chNextTick <- true
		}
	}
}

func (b *B) tick() *bytes.Buffer {
	bufSize := b.size * len(b.channels)
	byteBuffer := make([]byte, 0)
	buf := bytes.NewBuffer(byteBuffer)
	for i := 0; i < bufSize; i += 2 {
		lrIndex := int(i / 2)
		lenCh := len(b.channels)
		ch := b.channels[lrIndex%lenCh]
		t := float64(b.t) + math.Floor(float64(lrIndex))/float64(b.rate)/float64(lenCh)
		counter := b.i + int(math.Floor(float64(lrIndex)/float64(lenCh)))

		value := float64(0)
		n := float64(0)
		for j := 0; j < len(ch.funcs); j++ {
			x := ch.funcs[j](float64(t), counter)
			n += x
		}
		n /= float64(len(ch.funcs))

		if ch.funcValueType == FuncValueTypeFloat {
			value = signed(n)
		} else {
			b_ := math.Pow(2, float64(ch.funcValueType))
			x := math.Mod(math.Floor(n), b_) / b_ * math.Pow(2, 15)
			value = x
		}
		if err := binary.Write(buf, binary.LittleEndian, int16(clamp(value))); err != nil {
			panic(err)
		}
	}
	b.i += b.size / 2
	b.t += float64(b.size) / float64(2) / float64(b.rate)
	return buf
}

func clamp(x float64) float64 {
	return math.Max(math.Min(x, math.Pow(2, 15)-1), -math.Pow(2, 15))
}

func signed(n float64) float64 {
	b := math.Pow(2, 15)
	if n > 0 {
		return math.Min(b-1, math.Floor(b*n-1))
	}
	return math.Max(-b, math.Ceil(b*n-1))
}

func (b *B) Play(opts map[string]string) {
	go SoxPlay(mergeArgs(opts, map[string]string{
		"c": strconv.Itoa(len(b.channels)),
		"r": strconv.Itoa(b.rate),
		"t": "s16",
		"-": "DUMMY",
	}), b.waveReceiver())
	<-b.chEndSox
	b.pipeReader.Close()
}

func (b *B) Record(file string, opts map[string]string) {
	go SoxRecord(file, mergeArgs(opts, map[string]string{
		"c": strconv.Itoa(len(b.channels)),
		"r": strconv.Itoa(b.rate),
		"t": "s16",
		"-": "DUMMY",
	}), b.waveReceiver())
	<-b.chEndSox
	b.pipeReader.Close()
}

func (b *B) waveReceiver() func() []byte {
	readBuf := make([]byte, b.size*len(b.channels))
	return func() []byte {
		if _, err := b.pipeReader.Read(readBuf); err != nil {
			panic(err)
		}
		return readBuf
	}
}
<|endoftext|>


<commit_before>package manifest

import (
	"encoding/json"
	"fmt"

	compressiontypes "github.com/containers/image/v5/pkg/compression/types"
	"github.com/containers/image/v5/types"
	"github.com/sirupsen/logrus"
)

// dupStringSlice returns a deep copy of a slice of strings, or nil if the
// source slice is empty.
func dupStringSlice(list []string) []string {
	if len(list) == 0 {
		return nil
	}
	dup := make([]string, len(list))
	copy(dup, list)
	return dup
}

// dupStringStringMap returns a deep copy of a map[string]string, or nil if the
// passed-in map is nil or has no keys.
func dupStringStringMap(m map[string]string) map[string]string {
	if len(m) == 0 {
		return nil
	}
	result := make(map[string]string)
	for k, v := range m {
		result[k] = v
	}
	return result
}

// allowedManifestFields is a bit mask of “essential” manifest fields that validateUnambiguousManifestFormat
// can expect to be present.
type allowedManifestFields int

const (
	allowedFieldConfig allowedManifestFields = 1 << iota
	allowedFieldFSLayers
	allowedFieldHistory
	allowedFieldLayers
	allowedFieldManifests
	allowedFieldFirstUnusedBit // Keep this at the end!
)

// validateUnambiguousManifestFormat rejects manifests (incl. multi-arch) that look like more than
// one kind we currently recognize, i.e. if they contain any of the known “essential” format fields
// other than the ones the caller specifically allows.
// expectedMIMEType is used only for diagnostics.
// NOTE: The caller should do the non-heuristic validations (e.g. check for any specified format
// identification/version, or other “magic numbers”) before calling this, to cleanly reject unambiguous
// data that just isn’t what was expected, as opposed to actually ambiguous data.
func validateUnambiguousManifestFormat(manifest []byte, expectedMIMEType string,
	allowed allowedManifestFields) error {
	if allowed >= allowedFieldFirstUnusedBit {
		return fmt.Errorf("internal error: invalid allowedManifestFields value %#v", allowed)
	}
	// Use a private type to decode, not just a map[string]interface{}, because we want
	// to also reject case-insensitive matches (which would be used by Go when really decoding
	// the manifest).
	// (It is expected that as manifest formats are added or extended over time, more fields will be added
	// here.)
	detectedFields := struct {
		Config    interface{} `json:"config"`
		FSLayers  interface{} `json:"fsLayers"`
		History   interface{} `json:"history"`
		Layers    interface{} `json:"layers"`
		Manifests interface{} `json:"manifests"`
	}{}
	if err := json.Unmarshal(manifest, &detectedFields); err != nil {
		// The caller was supposed to already validate version numbers, so this should not happen;
		// let’s not bother with making this error “nice”.
		return err
	}
	unexpected := []string{}
	// Sadly this isn’t easy to automate in Go, without reflection. So, copy&paste.
	if detectedFields.Config != nil && (allowed&allowedFieldConfig) == 0 {
		unexpected = append(unexpected, "config")
	}
	if detectedFields.FSLayers != nil && (allowed&allowedFieldFSLayers) == 0 {
		unexpected = append(unexpected, "fsLayers")
	}
	if detectedFields.History != nil && (allowed&allowedFieldHistory) == 0 {
		unexpected = append(unexpected, "history")
	}
	if detectedFields.Layers != nil && (allowed&allowedFieldLayers) == 0 {
		unexpected = append(unexpected, "layers")
	}
	if detectedFields.Manifests != nil && (allowed&allowedFieldManifests) == 0 {
		unexpected = append(unexpected, "manifests")
	}
	if len(unexpected) != 0 {
		return fmt.Errorf(`rejecting ambiguous manifest, unexpected fields %#v in supposedly %s`,
			unexpected, expectedMIMEType)
	}
	return nil
}

// layerInfosToStrings converts a list of layer infos, presumably obtained from a Manifest.LayerInfos()
// method call, into a format suitable for inclusion in a types.ImageInspectInfo structure.
func layerInfosToStrings(infos []LayerInfo) []string {
	layers := make([]string, len(infos))
	for i, info := range infos {
		layers[i] = info.Digest.String()
	}
	return layers
}

// compressionMIMETypeSet describes a set of MIME type “variants” that represent differently-compressed
// versions of “the same kind of content”.
// The map key is the return value of compressiontypes.Algorithm.Name(), or mtsUncompressed;
// the map value is a MIME type, or mtsUnsupportedMIMEType to mean "recognized but unsupported".
type compressionMIMETypeSet map[string]string

const mtsUncompressed = ""        // A key in compressionMIMETypeSet for the uncompressed variant
const mtsUnsupportedMIMEType = "" // A value in compressionMIMETypeSet that means “recognized but unsupported”

// compressionVariantMIMEType returns a variant of mimeType for the specified algorithm (which may be nil
// to mean "no compression"), based on variantTable.
// The returned error will be a ManifestLayerCompressionIncompatibilityError if mimeType has variants
// that differ only in what type of compression is applied, but it can't be combined with this
// algorithm to produce an updated MIME type that complies with the standard that defines mimeType.
// If the compression algorithm is unrecognized, or mimeType is not known to have variants that
// differ from it only in what type of compression has been applied, the returned error will not be
// a ManifestLayerCompressionIncompatibilityError.
func compressionVariantMIMEType(variantTable []compressionMIMETypeSet, mimeType string, algorithm *compressiontypes.Algorithm) (string, error) {
	if mimeType == mtsUnsupportedMIMEType { // Prevent matching against the {algo:mtsUnsupportedMIMEType} entries
		return "", fmt.Errorf("cannot update unknown MIME type")
	}
	for _, variants := range variantTable {
		for _, mt := range variants {
			if mt == mimeType { // Found the variant
				name := mtsUncompressed
				if algorithm != nil {
					name = algorithm.InternalUnstableUndocumentedMIMEQuestionMark()
				}
				if res, ok := variants[name]; ok {
					if res != mtsUnsupportedMIMEType {
						return res, nil
					}
					if name != mtsUncompressed {
						return "", ManifestLayerCompressionIncompatibilityError{fmt.Sprintf("%s compression is not supported for type %q", name, mt)}
					}
					return "", ManifestLayerCompressionIncompatibilityError{fmt.Sprintf("uncompressed variant is not supported for type %q", mt)}
				}
				if name != mtsUncompressed {
					return "", ManifestLayerCompressionIncompatibilityError{fmt.Sprintf("unknown compressed with algorithm %s variant for type %s", name, mt)}
				}
				// We can't very well say “the idea of no compression is unknown”
				return "", ManifestLayerCompressionIncompatibilityError{fmt.Sprintf("uncompressed variant is not supported for type %q", mt)}
			}
		}
	}
	if algorithm != nil {
		return "", fmt.Errorf("unsupported MIME type for compression: %s", mimeType)
	}
	return "", fmt.Errorf("unsupported MIME type for decompression: %s", mimeType)
}

// updatedMIMEType returns the result of applying edits in updated (MediaType, CompressionOperation) to
// mimeType, based on variantTable.  It may use updated.Digest for error messages.
// The returned error will be a ManifestLayerCompressionIncompatibilityError if mimeType has variants
// that differ only in what type of compression is applied, but applying updated.CompressionOperation
// and updated.CompressionAlgorithm to it won't produce an updated MIME type that complies with the
// standard that defines mimeType.
func updatedMIMEType(variantTable []compressionMIMETypeSet, mimeType string, updated types.BlobInfo) (string, error) {
	// Note that manifests in containers-storage might be reporting the
	// wrong media type since the original manifests are stored while layers
	// are decompressed in storage.  Hence, we need to consider the case
	// that an already {de}compressed layer should be {de}compressed;
	// compressionVariantMIMEType does that by not caring whether the original is
	// {de}compressed.
	switch updated.CompressionOperation {
	case types.PreserveOriginal:
		// Force a change to the media type if we're being told to use a particular compressor,
		// since it might be different from the one associated with the media type.  Otherwise,
		// try to keep the original media type.
		if updated.CompressionAlgorithm != nil {
			return compressionVariantMIMEType(variantTable, mimeType, updated.CompressionAlgorithm)
		}
		// Keep the original media type.
		return mimeType, nil

	case types.Decompress:
		return compressionVariantMIMEType(variantTable, mimeType, nil)

	case types.Compress:
		if updated.CompressionAlgorithm == nil {
			logrus.Debugf("Error preparing updated manifest: blob %q was compressed but does not specify by which algorithm: falling back to use the original blob", updated.Digest)
			return mimeType, nil
		}
		return compressionVariantMIMEType(variantTable, mimeType, updated.CompressionAlgorithm)

	default:
		return "", fmt.Errorf("unknown compression operation (%d)", updated.CompressionOperation)
	}
}

// ManifestLayerCompressionIncompatibilityError indicates that a specified compression algorithm
// could not be applied to a layer MIME type.  A caller that receives this should either retry
// the call with a different compression algorithm, or attempt to use a different manifest type.
type ManifestLayerCompressionIncompatibilityError struct {
	text string
}

func (m ManifestLayerCompressionIncompatibilityError) Error() string {
	return m.text
}
<commit_msg>Split findCompressionMIMETypeSet from compressionVariantMIMEType<commit_after>package manifest

import (
	"encoding/json"
	"fmt"

	compressiontypes "github.com/containers/image/v5/pkg/compression/types"
	"github.com/containers/image/v5/types"
	"github.com/sirupsen/logrus"
)

// dupStringSlice returns a deep copy of a slice of strings, or nil if the
// source slice is empty.
func dupStringSlice(list []string) []string {
	if len(list) == 0 {
		return nil
	}
	dup := make([]string, len(list))
	copy(dup, list)
	return dup
}

// dupStringStringMap returns a deep copy of a map[string]string, or nil if the
// passed-in map is nil or has no keys.
func dupStringStringMap(m map[string]string) map[string]string {
	if len(m) == 0 {
		return nil
	}
	result := make(map[string]string)
	for k, v := range m {
		result[k] = v
	}
	return result
}

// allowedManifestFields is a bit mask of “essential” manifest fields that validateUnambiguousManifestFormat
// can expect to be present.
type allowedManifestFields int

const (
	allowedFieldConfig allowedManifestFields = 1 << iota
	allowedFieldFSLayers
	allowedFieldHistory
	allowedFieldLayers
	allowedFieldManifests
	allowedFieldFirstUnusedBit // Keep this at the end!
)

// validateUnambiguousManifestFormat rejects manifests (incl. multi-arch) that look like more than
// one kind we currently recognize, i.e. if they contain any of the known “essential” format fields
// other than the ones the caller specifically allows.
// expectedMIMEType is used only for diagnostics.
// NOTE: The caller should do the non-heuristic validations (e.g. check for any specified format
// identification/version, or other “magic numbers”) before calling this, to cleanly reject unambiguous
// data that just isn’t what was expected, as opposed to actually ambiguous data.
func validateUnambiguousManifestFormat(manifest []byte, expectedMIMEType string,
	allowed allowedManifestFields) error {
	if allowed >= allowedFieldFirstUnusedBit {
		return fmt.Errorf("internal error: invalid allowedManifestFields value %#v", allowed)
	}
	// Use a private type to decode, not just a map[string]interface{}, because we want
	// to also reject case-insensitive matches (which would be used by Go when really decoding
	// the manifest).
	// (It is expected that as manifest formats are added or extended over time, more fields will be added
	// here.)
	detectedFields := struct {
		Config    interface{} `json:"config"`
		FSLayers  interface{} `json:"fsLayers"`
		History   interface{} `json:"history"`
		Layers    interface{} `json:"layers"`
		Manifests interface{} `json:"manifests"`
	}{}
	if err := json.Unmarshal(manifest, &detectedFields); err != nil {
		// The caller was supposed to already validate version numbers, so this should not happen;
		// let’s not bother with making this error “nice”.
		return err
	}
	unexpected := []string{}
	// Sadly this isn’t easy to automate in Go, without reflection. So, copy&paste.
	if detectedFields.Config != nil && (allowed&allowedFieldConfig) == 0 {
		unexpected = append(unexpected, "config")
	}
	if detectedFields.FSLayers != nil && (allowed&allowedFieldFSLayers) == 0 {
		unexpected = append(unexpected, "fsLayers")
	}
	if detectedFields.History != nil && (allowed&allowedFieldHistory) == 0 {
		unexpected = append(unexpected, "history")
	}
	if detectedFields.Layers != nil && (allowed&allowedFieldLayers) == 0 {
		unexpected = append(unexpected, "layers")
	}
	if detectedFields.Manifests != nil && (allowed&allowedFieldManifests) == 0 {
		unexpected = append(unexpected, "manifests")
	}
	if len(unexpected) != 0 {
		return fmt.Errorf(`rejecting ambiguous manifest, unexpected fields %#v in supposedly %s`,
			unexpected, expectedMIMEType)
	}
	return nil
}

// layerInfosToStrings converts a list of layer infos, presumably obtained from a Manifest.LayerInfos()
// method call, into a format suitable for inclusion in a types.ImageInspectInfo structure.
func layerInfosToStrings(infos []LayerInfo) []string {
	layers := make([]string, len(infos))
	for i, info := range infos {
		layers[i] = info.Digest.String()
	}
	return layers
}

// compressionMIMETypeSet describes a set of MIME type “variants” that represent differently-compressed
// versions of “the same kind of content”.
// The map key is the return value of compressiontypes.Algorithm.Name(), or mtsUncompressed;
// the map value is a MIME type, or mtsUnsupportedMIMEType to mean "recognized but unsupported".
type compressionMIMETypeSet map[string]string

const mtsUncompressed = ""        // A key in compressionMIMETypeSet for the uncompressed variant
const mtsUnsupportedMIMEType = "" // A value in compressionMIMETypeSet that means “recognized but unsupported”

// findCompressionMIMETypeSet returns a pointer to a compressionMIMETypeSet in variantTable that contains a value of mimeType, or nil if not found
func findCompressionMIMETypeSet(variantTable []compressionMIMETypeSet, mimeType string) compressionMIMETypeSet {
	for _, variants := range variantTable {
		for _, mt := range variants {
			if mt == mimeType {
				return variants
			}
		}
	}
	return nil
}

// compressionVariantMIMEType returns a variant of mimeType for the specified algorithm (which may be nil
// to mean "no compression"), based on variantTable.
// The returned error will be a ManifestLayerCompressionIncompatibilityError if mimeType has variants
// that differ only in what type of compression is applied, but it can't be combined with this
// algorithm to produce an updated MIME type that complies with the standard that defines mimeType.
// If the compression algorithm is unrecognized, or mimeType is not known to have variants that
// differ from it only in what type of compression has been applied, the returned error will not be
// a ManifestLayerCompressionIncompatibilityError.
func compressionVariantMIMEType(variantTable []compressionMIMETypeSet, mimeType string, algorithm *compressiontypes.Algorithm) (string, error) {
	if mimeType == mtsUnsupportedMIMEType { // Prevent matching against the {algo:mtsUnsupportedMIMEType} entries
		return "", fmt.Errorf("cannot update unknown MIME type")
	}
	variants := findCompressionMIMETypeSet(variantTable, mimeType)
	if variants != nil {
		name := mtsUncompressed
		if algorithm != nil {
			name = algorithm.InternalUnstableUndocumentedMIMEQuestionMark()
		}
		if res, ok := variants[name]; ok {
			if res != mtsUnsupportedMIMEType {
				return res, nil
			}
			if name != mtsUncompressed {
				return "", ManifestLayerCompressionIncompatibilityError{fmt.Sprintf("%s compression is not supported for type %q", name, mimeType)}
			}
			return "", ManifestLayerCompressionIncompatibilityError{fmt.Sprintf("uncompressed variant is not supported for type %q", mimeType)}
		}
		if name != mtsUncompressed {
			return "", ManifestLayerCompressionIncompatibilityError{fmt.Sprintf("unknown compressed with algorithm %s variant for type %s", name, mimeType)}
		}
		// We can't very well say “the idea of no compression is unknown”
		return "", ManifestLayerCompressionIncompatibilityError{fmt.Sprintf("uncompressed variant is not supported for type %q", mimeType)}
	}
	if algorithm != nil {
		return "", fmt.Errorf("unsupported MIME type for compression: %s", mimeType)
	}
	return "", fmt.Errorf("unsupported MIME type for decompression: %s", mimeType)
}

// updatedMIMEType returns the result of applying edits in updated (MediaType, CompressionOperation) to
// mimeType, based on variantTable.  It may use updated.Digest for error messages.
// The returned error will be a ManifestLayerCompressionIncompatibilityError if mimeType has variants
// that differ only in what type of compression is applied, but applying updated.CompressionOperation
// and updated.CompressionAlgorithm to it won't produce an updated MIME type that complies with the
// standard that defines mimeType.
func updatedMIMEType(variantTable []compressionMIMETypeSet, mimeType string, updated types.BlobInfo) (string, error) {
	// Note that manifests in containers-storage might be reporting the
	// wrong media type since the original manifests are stored while layers
	// are decompressed in storage.  Hence, we need to consider the case
	// that an already {de}compressed layer should be {de}compressed;
	// compressionVariantMIMEType does that by not caring whether the original is
	// {de}compressed.
	switch updated.CompressionOperation {
	case types.PreserveOriginal:
		// Force a change to the media type if we're being told to use a particular compressor,
		// since it might be different from the one associated with the media type.  Otherwise,
		// try to keep the original media type.
		if updated.CompressionAlgorithm != nil {
			return compressionVariantMIMEType(variantTable, mimeType, updated.CompressionAlgorithm)
		}
		// Keep the original media type.
		return mimeType, nil

	case types.Decompress:
		return compressionVariantMIMEType(variantTable, mimeType, nil)

	case types.Compress:
		if updated.CompressionAlgorithm == nil {
			logrus.Debugf("Error preparing updated manifest: blob %q was compressed but does not specify by which algorithm: falling back to use the original blob", updated.Digest)
			return mimeType, nil
		}
		return compressionVariantMIMEType(variantTable, mimeType, updated.CompressionAlgorithm)

	default:
		return "", fmt.Errorf("unknown compression operation (%d)", updated.CompressionOperation)
	}
}

// ManifestLayerCompressionIncompatibilityError indicates that a specified compression algorithm
// could not be applied to a layer MIME type.  A caller that receives this should either retry
// the call with a different compression algorithm, or attempt to use a different manifest type.
type ManifestLayerCompressionIncompatibilityError struct {
	text string
}

func (m ManifestLayerCompressionIncompatibilityError) Error() string {
	return m.text
}
<|endoftext|>


<commit_before>// Copyright 2017 Pilosa Corp.
//
// 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 pilosa

// #cgo  CFLAGS:-mpopcnt

import (
	"encoding/json"
	"sort"

	"github.com/pilosa/pilosa/internal"
	"github.com/pilosa/pilosa/roaring"
)

// Bitmap represents a set of bits.
type Bitmap struct {
	segments []BitmapSegment

	// Attributes associated with the bitmap.
	Attrs map[string]interface{}
}

// NewBitmap returns a new instance of Bitmap.
func NewBitmap(bits ...uint64) *Bitmap {
	bm := &Bitmap{}
	for _, i := range bits {
		bm.SetBit(i)
	}
	return bm
}

// Merge merges data from other into b.
func (b *Bitmap) Merge(other *Bitmap) {
	var segments []BitmapSegment

	itr := newMergeSegmentIterator(b.segments, other.segments)
	for s0, s1 := itr.next(); s0 != nil || s1 != nil; s0, s1 = itr.next() {
		// Use the other bitmap's data if segment is missing.
		if s0 == nil {
			segments = append(segments, *s1)
			continue
		} else if s1 == nil {
			segments = append(segments, *s0)
			continue
		}

		// Otherwise merge.
		s0.Merge(s1)
		segments = append(segments, *s0)
	}

	b.segments = segments
	b.InvalidateCount()
}

// IntersectionCount returns the number of intersections between b and other.
func (b *Bitmap) IntersectionCount(other *Bitmap) uint64 {
	var n uint64

	itr := newMergeSegmentIterator(b.segments, other.segments)
	for s0, s1 := itr.next(); s0 != nil || s1 != nil; s0, s1 = itr.next() {
		// Ignore non-overlapping segments.
		if s0 == nil || s1 == nil {
			continue
		}

		n += s0.IntersectionCount(s1)
	}
	return n
}

// Intersect returns the itersection of b and other.
func (b *Bitmap) Intersect(other *Bitmap) *Bitmap {
	var segments []BitmapSegment

	itr := newMergeSegmentIterator(b.segments, other.segments)
	for s0, s1 := itr.next(); s0 != nil || s1 != nil; s0, s1 = itr.next() {
		// Ignore non-overlapping segments.
		if s0 == nil || s1 == nil {
			continue
		}
		segments = append(segments, *s0.Intersect(s1))
	}

	return &Bitmap{segments: segments}
}

// Xor returns the xor of b and other.
func (b *Bitmap) Xor(other *Bitmap) *Bitmap {
	var segments []BitmapSegment

	itr := newMergeSegmentIterator(b.segments, other.segments)
	for s0, s1 := itr.next(); s0 != nil || s1 != nil; s0, s1 = itr.next() {
		if s1 == nil {
			segments = append(segments, *s0)
			continue
		} else if s0 == nil {
			segments = append(segments, *s1)
			continue
		}

		segments = append(segments, *s0.Xor(s1))
	}

	return &Bitmap{segments: segments}
}

// Union returns the bitwise union of b and other.
func (b *Bitmap) Union(other *Bitmap) *Bitmap {
	var segments []BitmapSegment
	itr := newMergeSegmentIterator(b.segments, other.segments)
	for s0, s1 := itr.next(); s0 != nil || s1 != nil; s0, s1 = itr.next() {
		if s1 == nil {
			segments = append(segments, *s0)
			continue
		} else if s0 == nil {
			segments = append(segments, *s1)
			continue
		}
		segments = append(segments, *s0.Union(s1))
	}

	return &Bitmap{segments: segments}
}

// Difference returns the diff of b and other.
func (b *Bitmap) Difference(other *Bitmap) *Bitmap {
	var segments []BitmapSegment

	itr := newMergeSegmentIterator(b.segments, other.segments)
	for s0, s1 := itr.next(); s0 != nil || s1 != nil; s0, s1 = itr.next() {
		if s0 == nil {
			continue
		} else if s1 == nil {
			segments = append(segments, *s0)
			continue
		}
		segments = append(segments, *s0.Difference(s1))
	}

	return &Bitmap{segments: segments}
}

// SetBit sets the i-th bit of the bitmap.
func (b *Bitmap) SetBit(i uint64) (changed bool) {
	return b.createSegmentIfNotExists(i / SliceWidth).SetBit(i)
}

// ClearBit clears the i-th bit of the bitmap.
func (b *Bitmap) ClearBit(i uint64) (changed bool) {
	s := b.segment(i / SliceWidth)
	if s == nil {
		return false
	}
	return s.ClearBit(i)
}

// segment returns a segment for a given slice.
// Returns nil if segment does not exist.
func (b *Bitmap) segment(slice uint64) *BitmapSegment {
	if i := sort.Search(len(b.segments), func(i int) bool {
		return b.segments[i].slice >= slice
	}); i < len(b.segments) && b.segments[i].slice == slice {
		return &b.segments[i]
	}
	return nil
}

func (b *Bitmap) createSegmentIfNotExists(slice uint64) *BitmapSegment {
	i := sort.Search(len(b.segments), func(i int) bool {
		return b.segments[i].slice >= slice
	})

	// Return exact match.
	if i < len(b.segments) && b.segments[i].slice == slice {
		return &b.segments[i]
	}

	// Insert new segment.
	b.segments = append(b.segments, BitmapSegment{})
	if i < len(b.segments) {
		copy(b.segments[i+1:], b.segments[i:])
	}
	b.segments[i] = BitmapSegment{
		slice:    slice,
		writable: true,
	}

	return &b.segments[i]
}

// InvalidateCount updates the cached count in the bitmap.
func (b *Bitmap) InvalidateCount() {
	for i := range b.segments {
		b.segments[i].InvalidateCount()
	}
}

// IncrementCount increments the bitmap cached counter, note this is an optimization that assumes that the caller is aware the size increased.
func (b *Bitmap) IncrementCount(i uint64) {
	seg := b.segment(i / SliceWidth)
	if seg != nil {
		seg.n++
	}

}

// DecrementCount decrements the bitmap cached counter.
func (b *Bitmap) DecrementCount(i uint64) {
	seg := b.segment(i / SliceWidth)
	if seg != nil {
		if seg.n > 0 {
			seg.n--
		}
	}
}

// Count returns the number of set bits in the bitmap.
func (b *Bitmap) Count() uint64 {
	var n uint64
	for i := range b.segments {
		n += b.segments[i].Count()
	}
	return n
}

// MarshalJSON returns a JSON-encoded byte slice of b.
func (b *Bitmap) MarshalJSON() ([]byte, error) {
	var o struct {
		Attrs map[string]interface{} `json:"attrs"`
		Bits  []uint64               `json:"bits"`
	}
	o.Bits = b.Bits()

	o.Attrs = b.Attrs
	if o.Attrs == nil {
		o.Attrs = make(map[string]interface{})
	}

	return json.Marshal(&o)
}

// Bits returns the bits in b as a slice of ints.
func (b *Bitmap) Bits() []uint64 {
	a := make([]uint64, 0, b.Count())
	for i := range b.segments {
		a = append(a, b.segments[i].Bits()...)
	}
	return a
}

// encodeBitmap converts b into its internal representation.
func encodeBitmap(b *Bitmap) *internal.Bitmap {
	if b == nil {
		return nil
	}

	return &internal.Bitmap{
		Bits:  b.Bits(),
		Attrs: encodeAttrs(b.Attrs),
	}
}

// decodeBitmap converts b from its internal representation.
func decodeBitmap(pb *internal.Bitmap) *Bitmap {
	if pb == nil {
		return nil
	}

	b := NewBitmap()
	b.Attrs = decodeAttrs(pb.Attrs)
	for _, v := range pb.Bits {
		b.SetBit(v)
	}
	return b
}

// Union performs a union on a slice of bitmaps.
func Union(bitmaps []*Bitmap) *Bitmap {
	other := bitmaps[0]
	for _, bm := range bitmaps[1:] {
		other = other.Union(bm)
	}
	return other
}

// BitmapSegment holds a subset of a bitmap.
// This could point to a mmapped roaring bitmap or an in-memory bitmap. The
// width of the segment will always match the slice width.
type BitmapSegment struct {
	// Slice this segment belongs to
	slice uint64

	// Underlying raw bitmap implementation.
	// This is an mmapped bitmap if writable is false. Otherwise
	// it is a heap allocated bitmap which can be manipulated.
	data     roaring.Bitmap
	writable bool

	// Bit count
	n uint64
}

// Merge adds chunks from other to s.
// Chunks in s are overwritten if they exist in other.
func (s *BitmapSegment) Merge(other *BitmapSegment) {
	s.ensureWritable()

	itr := other.data.Iterator()
	for v, eof := itr.Next(); !eof; v, eof = itr.Next() {
		s.SetBit(v)
	}
}

// IntersectionCount returns the number of intersections between s and other.
func (s *BitmapSegment) IntersectionCount(other *BitmapSegment) uint64 {
	return s.data.IntersectionCount(&other.data)
}

// Intersect returns the itersection of s and other.
func (s *BitmapSegment) Intersect(other *BitmapSegment) *BitmapSegment {
	data := s.data.Intersect(&other.data)

	return &BitmapSegment{
		data:  *data,
		slice: s.slice,
		n:     data.Count(),
	}
}

// Union returns the bitwise union of s and other.
func (s *BitmapSegment) Union(other *BitmapSegment) *BitmapSegment {
	data := s.data.Union(&other.data)

	return &BitmapSegment{
		data:  *data,
		slice: s.slice,
		n:     data.Count(),
	}
}

// Difference returns the diff of s and other.
func (s *BitmapSegment) Difference(other *BitmapSegment) *BitmapSegment {
	data := s.data.Difference(&other.data)

	return &BitmapSegment{
		data:  *data,
		slice: s.slice,
		n:     data.Count(),
	}
}

// Xor returns the xor of s and other.
func (s *BitmapSegment) Xor(other *BitmapSegment) *BitmapSegment {
	data := s.data.Xor(&other.data)

	return &BitmapSegment{
		data:  *data,
		slice: s.slice,
		n:     data.Count(),
	}
}

// SetBit sets the i-th bit of the bitmap.
func (s *BitmapSegment) SetBit(i uint64) (changed bool) {
	s.ensureWritable()
	changed, _ = s.data.Add(i)
	if changed {
		s.n++
	}
	return changed
}

// ClearBit clears the i-th bit of the bitmap.
func (s *BitmapSegment) ClearBit(i uint64) (changed bool) {
	s.ensureWritable()

	changed, _ = s.data.Remove(i)
	if changed {
		s.n--
	}
	return changed
}

// InvalidateCount updates the cached count in the bitmap.
func (s *BitmapSegment) InvalidateCount() {
	s.n = s.data.Count()
}

// Bits returns a list of all bits set in the segment.
func (s *BitmapSegment) Bits() []uint64 {
	a := make([]uint64, 0, s.Count())
	itr := s.data.Iterator()
	for v, eof := itr.Next(); !eof; v, eof = itr.Next() {
		a = append(a, v)
	}
	return a
}

// Count returns the number of set bits in the bitmap.
func (s *BitmapSegment) Count() uint64 { return s.n }

// ensureWritable clones the segment if it is pointing to non-writable data.
func (s *BitmapSegment) ensureWritable() {
	if s.writable {
		return
	}

	s.data = *s.data.Clone()
	s.writable = true
}

// mergeSegmentIterator produces an iterator that loops through two sets of segments.
type mergeSegmentIterator struct {
	a0, a1 []BitmapSegment
}

// newMergeSegmentIterator returns a new instance of mergeSegmentIterator.
func newMergeSegmentIterator(a0, a1 []BitmapSegment) mergeSegmentIterator {
	return mergeSegmentIterator{a0: a0, a1: a1}
}

// next returns the next set of segments.
func (itr *mergeSegmentIterator) next() (s0, s1 *BitmapSegment) {
	// Find current segments.
	if len(itr.a0) > 0 {
		s0 = &itr.a0[0]
	}
	if len(itr.a1) > 0 {
		s1 = &itr.a1[0]
	}

	// Return if either or both are nil.
	if s0 == nil && s1 == nil {
		return
	} else if s0 == nil {
		itr.a1 = itr.a1[1:]
		return
	} else if s1 == nil {
		itr.a0 = itr.a0[1:]
		return
	}

	// Otherwise determine which is first.
	if s0.slice < s1.slice {
		itr.a0 = itr.a0[1:]
		return s0, nil
	} else if s0.slice > s1.slice {
		itr.a1 = itr.a1[1:]
		return s1, nil
	}

	// Return both if slices are equal.
	itr.a0, itr.a1 = itr.a0[1:], itr.a1[1:]
	return s0, s1
}
<commit_msg>Remove unused cgo preamble<commit_after>// Copyright 2017 Pilosa Corp.
//
// 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 pilosa

import (
	"encoding/json"
	"sort"

	"github.com/pilosa/pilosa/internal"
	"github.com/pilosa/pilosa/roaring"
)

// Bitmap represents a set of bits.
type Bitmap struct {
	segments []BitmapSegment

	// Attributes associated with the bitmap.
	Attrs map[string]interface{}
}

// NewBitmap returns a new instance of Bitmap.
func NewBitmap(bits ...uint64) *Bitmap {
	bm := &Bitmap{}
	for _, i := range bits {
		bm.SetBit(i)
	}
	return bm
}

// Merge merges data from other into b.
func (b *Bitmap) Merge(other *Bitmap) {
	var segments []BitmapSegment

	itr := newMergeSegmentIterator(b.segments, other.segments)
	for s0, s1 := itr.next(); s0 != nil || s1 != nil; s0, s1 = itr.next() {
		// Use the other bitmap's data if segment is missing.
		if s0 == nil {
			segments = append(segments, *s1)
			continue
		} else if s1 == nil {
			segments = append(segments, *s0)
			continue
		}

		// Otherwise merge.
		s0.Merge(s1)
		segments = append(segments, *s0)
	}

	b.segments = segments
	b.InvalidateCount()
}

// IntersectionCount returns the number of intersections between b and other.
func (b *Bitmap) IntersectionCount(other *Bitmap) uint64 {
	var n uint64

	itr := newMergeSegmentIterator(b.segments, other.segments)
	for s0, s1 := itr.next(); s0 != nil || s1 != nil; s0, s1 = itr.next() {
		// Ignore non-overlapping segments.
		if s0 == nil || s1 == nil {
			continue
		}

		n += s0.IntersectionCount(s1)
	}
	return n
}

// Intersect returns the itersection of b and other.
func (b *Bitmap) Intersect(other *Bitmap) *Bitmap {
	var segments []BitmapSegment

	itr := newMergeSegmentIterator(b.segments, other.segments)
	for s0, s1 := itr.next(); s0 != nil || s1 != nil; s0, s1 = itr.next() {
		// Ignore non-overlapping segments.
		if s0 == nil || s1 == nil {
			continue
		}
		segments = append(segments, *s0.Intersect(s1))
	}

	return &Bitmap{segments: segments}
}

// Xor returns the xor of b and other.
func (b *Bitmap) Xor(other *Bitmap) *Bitmap {
	var segments []BitmapSegment

	itr := newMergeSegmentIterator(b.segments, other.segments)
	for s0, s1 := itr.next(); s0 != nil || s1 != nil; s0, s1 = itr.next() {
		if s1 == nil {
			segments = append(segments, *s0)
			continue
		} else if s0 == nil {
			segments = append(segments, *s1)
			continue
		}

		segments = append(segments, *s0.Xor(s1))
	}

	return &Bitmap{segments: segments}
}

// Union returns the bitwise union of b and other.
func (b *Bitmap) Union(other *Bitmap) *Bitmap {
	var segments []BitmapSegment
	itr := newMergeSegmentIterator(b.segments, other.segments)
	for s0, s1 := itr.next(); s0 != nil || s1 != nil; s0, s1 = itr.next() {
		if s1 == nil {
			segments = append(segments, *s0)
			continue
		} else if s0 == nil {
			segments = append(segments, *s1)
			continue
		}
		segments = append(segments, *s0.Union(s1))
	}

	return &Bitmap{segments: segments}
}

// Difference returns the diff of b and other.
func (b *Bitmap) Difference(other *Bitmap) *Bitmap {
	var segments []BitmapSegment

	itr := newMergeSegmentIterator(b.segments, other.segments)
	for s0, s1 := itr.next(); s0 != nil || s1 != nil; s0, s1 = itr.next() {
		if s0 == nil {
			continue
		} else if s1 == nil {
			segments = append(segments, *s0)
			continue
		}
		segments = append(segments, *s0.Difference(s1))
	}

	return &Bitmap{segments: segments}
}

// SetBit sets the i-th bit of the bitmap.
func (b *Bitmap) SetBit(i uint64) (changed bool) {
	return b.createSegmentIfNotExists(i / SliceWidth).SetBit(i)
}

// ClearBit clears the i-th bit of the bitmap.
func (b *Bitmap) ClearBit(i uint64) (changed bool) {
	s := b.segment(i / SliceWidth)
	if s == nil {
		return false
	}
	return s.ClearBit(i)
}

// segment returns a segment for a given slice.
// Returns nil if segment does not exist.
func (b *Bitmap) segment(slice uint64) *BitmapSegment {
	if i := sort.Search(len(b.segments), func(i int) bool {
		return b.segments[i].slice >= slice
	}); i < len(b.segments) && b.segments[i].slice == slice {
		return &b.segments[i]
	}
	return nil
}

func (b *Bitmap) createSegmentIfNotExists(slice uint64) *BitmapSegment {
	i := sort.Search(len(b.segments), func(i int) bool {
		return b.segments[i].slice >= slice
	})

	// Return exact match.
	if i < len(b.segments) && b.segments[i].slice == slice {
		return &b.segments[i]
	}

	// Insert new segment.
	b.segments = append(b.segments, BitmapSegment{})
	if i < len(b.segments) {
		copy(b.segments[i+1:], b.segments[i:])
	}
	b.segments[i] = BitmapSegment{
		slice:    slice,
		writable: true,
	}

	return &b.segments[i]
}

// InvalidateCount updates the cached count in the bitmap.
func (b *Bitmap) InvalidateCount() {
	for i := range b.segments {
		b.segments[i].InvalidateCount()
	}
}

// IncrementCount increments the bitmap cached counter, note this is an optimization that assumes that the caller is aware the size increased.
func (b *Bitmap) IncrementCount(i uint64) {
	seg := b.segment(i / SliceWidth)
	if seg != nil {
		seg.n++
	}

}

// DecrementCount decrements the bitmap cached counter.
func (b *Bitmap) DecrementCount(i uint64) {
	seg := b.segment(i / SliceWidth)
	if seg != nil {
		if seg.n > 0 {
			seg.n--
		}
	}
}

// Count returns the number of set bits in the bitmap.
func (b *Bitmap) Count() uint64 {
	var n uint64
	for i := range b.segments {
		n += b.segments[i].Count()
	}
	return n
}

// MarshalJSON returns a JSON-encoded byte slice of b.
func (b *Bitmap) MarshalJSON() ([]byte, error) {
	var o struct {
		Attrs map[string]interface{} `json:"attrs"`
		Bits  []uint64               `json:"bits"`
	}
	o.Bits = b.Bits()

	o.Attrs = b.Attrs
	if o.Attrs == nil {
		o.Attrs = make(map[string]interface{})
	}

	return json.Marshal(&o)
}

// Bits returns the bits in b as a slice of ints.
func (b *Bitmap) Bits() []uint64 {
	a := make([]uint64, 0, b.Count())
	for i := range b.segments {
		a = append(a, b.segments[i].Bits()...)
	}
	return a
}

// encodeBitmap converts b into its internal representation.
func encodeBitmap(b *Bitmap) *internal.Bitmap {
	if b == nil {
		return nil
	}

	return &internal.Bitmap{
		Bits:  b.Bits(),
		Attrs: encodeAttrs(b.Attrs),
	}
}

// decodeBitmap converts b from its internal representation.
func decodeBitmap(pb *internal.Bitmap) *Bitmap {
	if pb == nil {
		return nil
	}

	b := NewBitmap()
	b.Attrs = decodeAttrs(pb.Attrs)
	for _, v := range pb.Bits {
		b.SetBit(v)
	}
	return b
}

// Union performs a union on a slice of bitmaps.
func Union(bitmaps []*Bitmap) *Bitmap {
	other := bitmaps[0]
	for _, bm := range bitmaps[1:] {
		other = other.Union(bm)
	}
	return other
}

// BitmapSegment holds a subset of a bitmap.
// This could point to a mmapped roaring bitmap or an in-memory bitmap. The
// width of the segment will always match the slice width.
type BitmapSegment struct {
	// Slice this segment belongs to
	slice uint64

	// Underlying raw bitmap implementation.
	// This is an mmapped bitmap if writable is false. Otherwise
	// it is a heap allocated bitmap which can be manipulated.
	data     roaring.Bitmap
	writable bool

	// Bit count
	n uint64
}

// Merge adds chunks from other to s.
// Chunks in s are overwritten if they exist in other.
func (s *BitmapSegment) Merge(other *BitmapSegment) {
	s.ensureWritable()

	itr := other.data.Iterator()
	for v, eof := itr.Next(); !eof; v, eof = itr.Next() {
		s.SetBit(v)
	}
}

// IntersectionCount returns the number of intersections between s and other.
func (s *BitmapSegment) IntersectionCount(other *BitmapSegment) uint64 {
	return s.data.IntersectionCount(&other.data)
}

// Intersect returns the itersection of s and other.
func (s *BitmapSegment) Intersect(other *BitmapSegment) *BitmapSegment {
	data := s.data.Intersect(&other.data)

	return &BitmapSegment{
		data:  *data,
		slice: s.slice,
		n:     data.Count(),
	}
}

// Union returns the bitwise union of s and other.
func (s *BitmapSegment) Union(other *BitmapSegment) *BitmapSegment {
	data := s.data.Union(&other.data)

	return &BitmapSegment{
		data:  *data,
		slice: s.slice,
		n:     data.Count(),
	}
}

// Difference returns the diff of s and other.
func (s *BitmapSegment) Difference(other *BitmapSegment) *BitmapSegment {
	data := s.data.Difference(&other.data)

	return &BitmapSegment{
		data:  *data,
		slice: s.slice,
		n:     data.Count(),
	}
}

// Xor returns the xor of s and other.
func (s *BitmapSegment) Xor(other *BitmapSegment) *BitmapSegment {
	data := s.data.Xor(&other.data)

	return &BitmapSegment{
		data:  *data,
		slice: s.slice,
		n:     data.Count(),
	}
}

// SetBit sets the i-th bit of the bitmap.
func (s *BitmapSegment) SetBit(i uint64) (changed bool) {
	s.ensureWritable()
	changed, _ = s.data.Add(i)
	if changed {
		s.n++
	}
	return changed
}

// ClearBit clears the i-th bit of the bitmap.
func (s *BitmapSegment) ClearBit(i uint64) (changed bool) {
	s.ensureWritable()

	changed, _ = s.data.Remove(i)
	if changed {
		s.n--
	}
	return changed
}

// InvalidateCount updates the cached count in the bitmap.
func (s *BitmapSegment) InvalidateCount() {
	s.n = s.data.Count()
}

// Bits returns a list of all bits set in the segment.
func (s *BitmapSegment) Bits() []uint64 {
	a := make([]uint64, 0, s.Count())
	itr := s.data.Iterator()
	for v, eof := itr.Next(); !eof; v, eof = itr.Next() {
		a = append(a, v)
	}
	return a
}

// Count returns the number of set bits in the bitmap.
func (s *BitmapSegment) Count() uint64 { return s.n }

// ensureWritable clones the segment if it is pointing to non-writable data.
func (s *BitmapSegment) ensureWritable() {
	if s.writable {
		return
	}

	s.data = *s.data.Clone()
	s.writable = true
}

// mergeSegmentIterator produces an iterator that loops through two sets of segments.
type mergeSegmentIterator struct {
	a0, a1 []BitmapSegment
}

// newMergeSegmentIterator returns a new instance of mergeSegmentIterator.
func newMergeSegmentIterator(a0, a1 []BitmapSegment) mergeSegmentIterator {
	return mergeSegmentIterator{a0: a0, a1: a1}
}

// next returns the next set of segments.
func (itr *mergeSegmentIterator) next() (s0, s1 *BitmapSegment) {
	// Find current segments.
	if len(itr.a0) > 0 {
		s0 = &itr.a0[0]
	}
	if len(itr.a1) > 0 {
		s1 = &itr.a1[0]
	}

	// Return if either or both are nil.
	if s0 == nil && s1 == nil {
		return
	} else if s0 == nil {
		itr.a1 = itr.a1[1:]
		return
	} else if s1 == nil {
		itr.a0 = itr.a0[1:]
		return
	}

	// Otherwise determine which is first.
	if s0.slice < s1.slice {
		itr.a0 = itr.a0[1:]
		return s0, nil
	} else if s0.slice > s1.slice {
		itr.a1 = itr.a1[1:]
		return s1, nil
	}

	// Return both if slices are equal.
	itr.a0, itr.a1 = itr.a0[1:], itr.a1[1:]
	return s0, s1
}
<|endoftext|>