<commit_before>package request_strategy

import (
	"sort"

	"github.com/anacrolix/multiless"
	pp "github.com/anacrolix/torrent/peer_protocol"
	"github.com/anacrolix/torrent/types"
)

type (
	Request       = types.Request
	pieceIndex    = types.PieceIndex
	piecePriority = types.PiecePriority
	// This can be made into a type-param later, will be great for testing.
	ChunkSpec = types.ChunkSpec
)

type ClientPieceOrder struct {
	pieces []pieceRequestOrderPiece
}

type pieceRequestOrderPiece struct {
	t     *Torrent
	index pieceIndex
	Piece
}

func (me *ClientPieceOrder) Len() int {
	return len(me.pieces)
}

func (me ClientPieceOrder) sort() {
	sort.Slice(me.pieces, me.less)
}

func (me ClientPieceOrder) less(_i, _j int) bool {
	i := me.pieces[_i]
	j := me.pieces[_j]
	return multiless.New().Int(
		int(j.Priority), int(i.Priority),
	).Bool(
		j.Partial, i.Partial,
	).Int64(i.Availability, j.Availability).Int(i.index, j.index).Less()
}

type requestsPeer struct {
	Peer
	nextState                  PeerNextRequestState
	requestablePiecesRemaining int
}

func (rp *requestsPeer) canFitRequest() bool {
	return len(rp.nextState.Requests) < rp.MaxRequests
}

func (rp *requestsPeer) addNextRequest(r Request) {
	_, ok := rp.nextState.Requests[r]
	if ok {
		panic("should only add once")
	}
	rp.nextState.Requests[r] = struct{}{}
}

type peersForPieceRequests struct {
	requestsInPiece int
	*requestsPeer
}

func (me *peersForPieceRequests) addNextRequest(r Request) {
	me.requestsPeer.addNextRequest(r)
	me.requestsInPiece++
}

type Torrent struct {
	Pieces   []Piece
	Capacity *func() *int64
	Peers    []Peer // not closed.
}

func (requestOrder *ClientPieceOrder) DoRequests(torrents []*Torrent) map[PeerId]PeerNextRequestState {
	requestOrder.pieces = requestOrder.pieces[:0]
	allPeers := make(map[*Torrent][]*requestsPeer)
	// Storage capacity left for this run, keyed by the storage capacity pointer on the storage
	// TorrentImpl.
	storageLeft := make(map[*func() *int64]*int64)
	for _, t := range torrents {
		// TODO: We could do metainfo requests here.
		key := t.Capacity
		if key != nil {
			if _, ok := storageLeft[key]; !ok {
				storageLeft[key] = (*key)()
			}
		}
		var peers []*requestsPeer
		for _, p := range t.Peers {
			peers = append(peers, &requestsPeer{
				Peer: p,
				nextState: PeerNextRequestState{
					Requests: make(map[Request]struct{}),
				},
			})
		}
		for i, tp := range t.Pieces {
			requestOrder.pieces = append(requestOrder.pieces, pieceRequestOrderPiece{
				t:     t,
				index: i,
				Piece: tp,
			})
			if tp.Request {
				for _, p := range peers {
					if p.canRequestPiece(i) {
						p.requestablePiecesRemaining++
					}
				}
			}
		}
		allPeers[t] = peers
	}
	requestOrder.sort()
	for _, p := range requestOrder.pieces {
		torrentPiece := p
		if left := storageLeft[p.t.Capacity]; left != nil {
			if *left < int64(torrentPiece.Length) {
				continue
			}
			*left -= int64(torrentPiece.Length)
		}
		if !p.Request {
			continue
		}
		allocatePendingChunks(p, allPeers[p.t])
	}
	ret := make(map[PeerId]PeerNextRequestState)
	for _, peers := range allPeers {
		for _, rp := range peers {
			if rp.requestablePiecesRemaining != 0 {
				panic(rp.requestablePiecesRemaining)
			}
			ret[rp.Id] = rp.nextState
		}
	}
	return ret
}

func allocatePendingChunks(p pieceRequestOrderPiece, peers []*requestsPeer) {
	peersForPiece := make([]*peersForPieceRequests, 0, len(peers))
	for _, peer := range peers {
		peersForPiece = append(peersForPiece, &peersForPieceRequests{
			requestsInPiece: 0,
			requestsPeer:    peer,
		})
	}
	defer func() {
		for _, peer := range peersForPiece {
			if peer.canRequestPiece(p.index) {
				peer.requestablePiecesRemaining--
			}
		}
	}()
	sortPeersForPiece := func(byHasRequest *Request) {
		sort.Slice(peersForPiece, func(i, j int) bool {
			ml := multiless.New().Int(
				peersForPiece[i].requestsInPiece,
				peersForPiece[j].requestsInPiece,
			).Int(
				peersForPiece[i].requestablePiecesRemaining,
				peersForPiece[j].requestablePiecesRemaining,
			).Float64(
				peersForPiece[j].DownloadRate,
				peersForPiece[i].DownloadRate,
			)
			if byHasRequest != nil {
				_, iHas := peersForPiece[i].nextState.Requests[*byHasRequest]
				_, jHas := peersForPiece[j].nextState.Requests[*byHasRequest]
				ml = ml.Bool(jHas, iHas)
			}
			return ml.Int64(
				int64(peersForPiece[j].Age), int64(peersForPiece[i].Age),
				// TODO: Probably peer priority can come next
			).Uintptr(
				peersForPiece[i].Id.Uintptr(),
				peersForPiece[j].Id.Uintptr(),
			).MustLess()
		})
	}
	preallocated := make(map[ChunkSpec]*peersForPieceRequests, p.NumPendingChunks)
	p.iterPendingChunksWrapper(func(spec ChunkSpec) {
		req := Request{pp.Integer(p.index), spec}
		for _, p := range peersForPiece {
			if h := p.HasExistingRequest; h != nil && h(req) {
				preallocated[spec] = p
				p.addNextRequest(req)
			}
		}
	})
	pendingChunksRemaining := int(p.NumPendingChunks)
	p.iterPendingChunksWrapper(func(chunk types.ChunkSpec) {
		if _, ok := preallocated[chunk]; ok {
			return
		}
		req := Request{pp.Integer(p.index), chunk}
		defer func() { pendingChunksRemaining-- }()
		sortPeersForPiece(nil)
		for _, peer := range peersForPiece {
			if !peer.canFitRequest() {
				continue
			}
			if !peer.HasPiece(p.index) {
				continue
			}
			if !peer.pieceAllowedFastOrDefault(p.index) {
				// TODO: Verify that's okay to stay uninterested if we request allowed fast pieces.
				peer.nextState.Interested = true
				if peer.Choking {
					continue
				}
			}
			peer.addNextRequest(req)
			return
		}
	})
chunk:
	for chunk, prePeer := range preallocated {
		req := Request{pp.Integer(p.index), chunk}
		prePeer.requestsInPiece--
		sortPeersForPiece(&req)
		delete(prePeer.nextState.Requests, req)
		for _, peer := range peersForPiece {
			if !peer.canFitRequest() {
				continue
			}
			if !peer.HasPiece(p.index) {
				continue
			}
			if !peer.pieceAllowedFastOrDefault(p.index) {
				// TODO: Verify that's okay to stay uninterested if we request allowed fast pieces.
				peer.nextState.Interested = true
				if peer.Choking {
					continue
				}
			}
			pendingChunksRemaining--
			peer.addNextRequest(req)
			continue chunk
		}
	}
	if pendingChunksRemaining != 0 {
		panic(pendingChunksRemaining)
	}
}
<commit_msg>Do checks for preallocated requests too<commit_after>package request_strategy

import (
	"sort"

	"github.com/anacrolix/multiless"
	pp "github.com/anacrolix/torrent/peer_protocol"
	"github.com/anacrolix/torrent/types"
)

type (
	Request       = types.Request
	pieceIndex    = types.PieceIndex
	piecePriority = types.PiecePriority
	// This can be made into a type-param later, will be great for testing.
	ChunkSpec = types.ChunkSpec
)

type ClientPieceOrder struct {
	pieces []pieceRequestOrderPiece
}

type pieceRequestOrderPiece struct {
	t     *Torrent
	index pieceIndex
	Piece
}

func (me *ClientPieceOrder) Len() int {
	return len(me.pieces)
}

func (me ClientPieceOrder) sort() {
	sort.Slice(me.pieces, me.less)
}

func (me ClientPieceOrder) less(_i, _j int) bool {
	i := me.pieces[_i]
	j := me.pieces[_j]
	return multiless.New().Int(
		int(j.Priority), int(i.Priority),
	).Bool(
		j.Partial, i.Partial,
	).Int64(i.Availability, j.Availability).Int(i.index, j.index).Less()
}

type requestsPeer struct {
	Peer
	nextState                  PeerNextRequestState
	requestablePiecesRemaining int
}

func (rp *requestsPeer) canFitRequest() bool {
	return len(rp.nextState.Requests) < rp.MaxRequests
}

func (rp *requestsPeer) addNextRequest(r Request) {
	_, ok := rp.nextState.Requests[r]
	if ok {
		panic("should only add once")
	}
	rp.nextState.Requests[r] = struct{}{}
}

type peersForPieceRequests struct {
	requestsInPiece int
	*requestsPeer
}

func (me *peersForPieceRequests) addNextRequest(r Request) {
	me.requestsPeer.addNextRequest(r)
	me.requestsInPiece++
}

type Torrent struct {
	Pieces   []Piece
	Capacity *func() *int64
	Peers    []Peer // not closed.
}

func (requestOrder *ClientPieceOrder) DoRequests(torrents []*Torrent) map[PeerId]PeerNextRequestState {
	requestOrder.pieces = requestOrder.pieces[:0]
	allPeers := make(map[*Torrent][]*requestsPeer)
	// Storage capacity left for this run, keyed by the storage capacity pointer on the storage
	// TorrentImpl.
	storageLeft := make(map[*func() *int64]*int64)
	for _, t := range torrents {
		// TODO: We could do metainfo requests here.
		key := t.Capacity
		if key != nil {
			if _, ok := storageLeft[key]; !ok {
				storageLeft[key] = (*key)()
			}
		}
		var peers []*requestsPeer
		for _, p := range t.Peers {
			peers = append(peers, &requestsPeer{
				Peer: p,
				nextState: PeerNextRequestState{
					Requests: make(map[Request]struct{}),
				},
			})
		}
		for i, tp := range t.Pieces {
			requestOrder.pieces = append(requestOrder.pieces, pieceRequestOrderPiece{
				t:     t,
				index: i,
				Piece: tp,
			})
			if tp.Request {
				for _, p := range peers {
					if p.canRequestPiece(i) {
						p.requestablePiecesRemaining++
					}
				}
			}
		}
		allPeers[t] = peers
	}
	requestOrder.sort()
	for _, p := range requestOrder.pieces {
		torrentPiece := p
		if left := storageLeft[p.t.Capacity]; left != nil {
			if *left < int64(torrentPiece.Length) {
				continue
			}
			*left -= int64(torrentPiece.Length)
		}
		if !p.Request {
			continue
		}
		allocatePendingChunks(p, allPeers[p.t])
	}
	ret := make(map[PeerId]PeerNextRequestState)
	for _, peers := range allPeers {
		for _, rp := range peers {
			if rp.requestablePiecesRemaining != 0 {
				panic(rp.requestablePiecesRemaining)
			}
			ret[rp.Id] = rp.nextState
		}
	}
	return ret
}

func allocatePendingChunks(p pieceRequestOrderPiece, peers []*requestsPeer) {
	peersForPiece := make([]*peersForPieceRequests, 0, len(peers))
	for _, peer := range peers {
		peersForPiece = append(peersForPiece, &peersForPieceRequests{
			requestsInPiece: 0,
			requestsPeer:    peer,
		})
	}
	defer func() {
		for _, peer := range peersForPiece {
			if peer.canRequestPiece(p.index) {
				peer.requestablePiecesRemaining--
			}
		}
	}()
	sortPeersForPiece := func(byHasRequest *Request) {
		sort.Slice(peersForPiece, func(i, j int) bool {
			ml := multiless.New().Int(
				peersForPiece[i].requestsInPiece,
				peersForPiece[j].requestsInPiece,
			).Int(
				peersForPiece[i].requestablePiecesRemaining,
				peersForPiece[j].requestablePiecesRemaining,
			).Float64(
				peersForPiece[j].DownloadRate,
				peersForPiece[i].DownloadRate,
			)
			if byHasRequest != nil {
				_, iHas := peersForPiece[i].nextState.Requests[*byHasRequest]
				_, jHas := peersForPiece[j].nextState.Requests[*byHasRequest]
				ml = ml.Bool(jHas, iHas)
			}
			return ml.Int64(
				int64(peersForPiece[j].Age), int64(peersForPiece[i].Age),
				// TODO: Probably peer priority can come next
			).Uintptr(
				peersForPiece[i].Id.Uintptr(),
				peersForPiece[j].Id.Uintptr(),
			).MustLess()
		})
	}
	preallocated := make(map[ChunkSpec]*peersForPieceRequests, p.NumPendingChunks)
	p.iterPendingChunksWrapper(func(spec ChunkSpec) {
		req := Request{pp.Integer(p.index), spec}
		for _, peer := range peersForPiece {
			if h := peer.HasExistingRequest; h == nil || !h(req) {
				continue
			}
			if !peer.canFitRequest() {
				continue
			}
			if !peer.canRequestPiece(p.index) {
				continue
			}
			preallocated[spec] = peer
			peer.addNextRequest(req)
		}
	})
	pendingChunksRemaining := int(p.NumPendingChunks)
	p.iterPendingChunksWrapper(func(chunk types.ChunkSpec) {
		if _, ok := preallocated[chunk]; ok {
			return
		}
		req := Request{pp.Integer(p.index), chunk}
		defer func() { pendingChunksRemaining-- }()
		sortPeersForPiece(nil)
		for _, peer := range peersForPiece {
			if !peer.canFitRequest() {
				continue
			}
			if !peer.HasPiece(p.index) {
				continue
			}
			if !peer.pieceAllowedFastOrDefault(p.index) {
				// TODO: Verify that's okay to stay uninterested if we request allowed fast pieces.
				peer.nextState.Interested = true
				if peer.Choking {
					continue
				}
			}
			peer.addNextRequest(req)
			return
		}
	})
chunk:
	for chunk, prePeer := range preallocated {
		req := Request{pp.Integer(p.index), chunk}
		prePeer.requestsInPiece--
		sortPeersForPiece(&req)
		delete(prePeer.nextState.Requests, req)
		for _, peer := range peersForPiece {
			if !peer.canFitRequest() {
				continue
			}
			if !peer.HasPiece(p.index) {
				continue
			}
			if !peer.pieceAllowedFastOrDefault(p.index) {
				// TODO: Verify that's okay to stay uninterested if we request allowed fast pieces.
				peer.nextState.Interested = true
				if peer.Choking {
					continue
				}
			}
			pendingChunksRemaining--
			peer.addNextRequest(req)
			continue chunk
		}
	}
	if pendingChunksRemaining != 0 {
		panic(pendingChunksRemaining)
	}
}
<|endoftext|>


<commit_before>package main

import (
	"fmt"
	"github.com/ant0ine/go-json-rest/rest"
	"github.com/pmylund/go-cache"
	"github.com/stacktic/dropbox"
	"gopkg.in/yaml.v2"
	"io/ioutil"
	"log"
	"net/http"
	"path/filepath"
	"strings"
	"time"
)

type Config struct {
	Api_secret    string
	Api_key       string
	Client_token  string
	Context_root  string
	Server_listen string
	Server_port   string
}

type Archive struct {
	Software string
	Date     string
	Version  string
	Tag      string
}

type BasicResult struct {
	Tag  string
	Date string
}

func (a Archive) Init(path string) *Archive {
	_, filename := filepath.Split(path)
	parts := strings.Split(filename, "-")
	a.Software = parts[0]
	a.Version = strings.Join([]string{parts[1], parts[2]}, " ")
	a.Date = strings.Join([]string{parts[3], parts[4], parts[5]}, "-")
	a.Tag = parts[6]
	return &a
}

// The struct to store the configuration data
var config Config

// 12 hour caching that cleans up every 15 minutes
var cache_instance *cache.Cache

//Link to dropbox
var db *dropbox.Dropbox

var do_not_include []string

func main() {
	var err error

	config = Config{}
	cache_instance = cache.New(12*time.Hour, 15*time.Minute)
	db = dropbox.NewDropbox()

	do_not_include = []string{}
	do_not_include = append(do_not_include, ".txt")

	data, read_err := ioutil.ReadFile("config.yml")
	if read_err != nil {
		log.Fatal(read_err)
	}

	yaml_err := yaml.Unmarshal(data, &config)
	if yaml_err != nil {
		log.Fatal("error: %v", yaml_err)
	}
	fmt.Printf("--- config_file dump:\n%s\n\n", data)
	fmt.Printf("--- config dump:\n%s\n\n", config)

	db.SetAppInfo(config.Api_key, config.Api_secret)
	if len(config.Client_token) >= 1 {
		db.SetAccessToken(config.Client_token)
	} else {
		if err = db.Auth(); err != nil {
			fmt.Println(err)
			return
		}
		config.Client_token = db.AccessToken()
		db.SetAccessToken(config.Client_token)
		d, marshal_err := yaml.Marshal(&config)
		if marshal_err != nil {
			log.Fatal("error: %v", marshal_err)
		}
		ioutil.WriteFile("config.yml", []byte(d), 0644)
	}

	// root_paths := get_directories(cache, db, "")
	// fmt.Printf("--- paths dump:\n%s\n\n", root_paths)

	// nightly_files := get_files(cache, db, "ARMv7")
	// fmt.Printf("--- paths dump:\n%s\n\n", nightly_files)

	// setup server to link
	api := rest.NewApi()
	api.Use(rest.DefaultDevStack...)
	router, err := rest.MakeRouter(
		rest.Get(strings.Join([]string{config.Context_root, ""}, "/"), list_arch),
		rest.Get(strings.Join([]string{config.Context_root, "#arch"}, "/"), list_software),
		rest.Get(strings.Join([]string{config.Context_root, "#arch/#software"}, "/"), list_versions),
		rest.Get(strings.Join([]string{config.Context_root, "#arch/#software/#version"}, "/"), list_targets),
		rest.Get(strings.Join([]string{config.Context_root, "#arch/#software/#version/#target"}, "/"), link_target),
	)
	if err != nil {
		log.Fatal(err)
	}
	api.SetApp(router)
	s := []string{}
	s = append(s, config.Server_listen)
	s = append(s, config.Server_port)
	server_listen := strings.Join(s, ":")
	log.Fatal(http.ListenAndServe(server_listen, api.MakeHandler()))
}

func list_arch(w rest.ResponseWriter, r *rest.Request) {
	// Use caching to reduce calls to the Dropbox API
	cache_path := "arches"
	data, found := cache_instance.Get(cache_path)
	if found {
		if cached, ok := data.([]string); ok {
			w.WriteJson(cached)
			return
		} else {
			log.Println("Error: Unable to retrieve from cache")
		}
	}

	arches := []string{}
	directories := get_directories(cache_instance, db, "/")
	for _, arch := range directories {
		arches = append(arches, strings.Replace(arch.Path, "/", "", -1))
	}
	cache_instance.Set(cache_path, arches, 0)
	w.WriteJson(arches)
}

func list_software(w rest.ResponseWriter, r *rest.Request) {
	arch := r.PathParam("arch")

	// Use caching to reduce calls to the Dropbox API
	cache_path := arch
	data, found := cache_instance.Get(cache_path)
	if found {
		if cached, ok := data.([]string); ok {
			w.WriteJson(cached)
			return
		} else {
			log.Println("Error: Unable to retrieve from cache")
		}
	}

	softwares := make(map[string]string)
	files := get_files(cache_instance, db, arch)
	for _, file := range files {
		archive := new(Archive)
		archive = archive.Init(file.Path)
		softwares[archive.Software] = ""
	}
	keys := make([]string, 0, len(softwares))
	for k := range softwares {
		keys = append(keys, k)
	}
	cache_instance.Set(cache_path, keys, 0)
	w.WriteJson(keys)
}

func list_versions(w rest.ResponseWriter, r *rest.Request) {
	w.WriteJson([]string{"nightly", "beta", "stable"})
}

func list_targets(w rest.ResponseWriter, r *rest.Request) {
	arch := r.PathParam("arch")
	software := r.PathParam("software")
	version := r.PathParam("version")

	// Doesn't need to be cached, as its calls are already cached.

	targets := []BasicResult{BasicResult{"latest", ""}}
	target_path := get_target_path(arch, version)
	files := get_files(cache_instance, db, target_path)
	for _, file := range files {
		archive := new(Archive)
		archive = archive.Init(file.Path)
		if archive.Software == software {
			targets = append(targets, BasicResult{archive.Tag, archive.Date})
		}
	}
	w.WriteJson(targets)
}

func link_target(w rest.ResponseWriter, r *rest.Request) {
	arch := r.PathParam("arch")
	software := r.PathParam("software")
	version := r.PathParam("version")
	target := r.PathParam("target")

	target_file, found := get_target(arch, software, version, target)
	if found {
		target_link := get_link(cache_instance, db, target_file.Path)
		w.Header().Set("Location", target_link)
		w.WriteHeader(302)
	} else {
		w.WriteHeader(404)
		w.WriteJson(map[string]string{"error": "Target Not Found"})
	}
}

/*
	Get only a slice of the directories at a path
*/
func get_directories(cache *cache.Cache, db *dropbox.Dropbox, path string) []dropbox.Entry {
	return get(cache, db, path, true)
}

/*
	Get only a slice of the directories at a path
*/
func get_files(cache *cache.Cache, db *dropbox.Dropbox, path string) []dropbox.Entry {
	return get(cache, db, path, false)
}

/*
	Actually get a list of directories or files from the dropbox connection
*/
func get(cache *cache.Cache, db *dropbox.Dropbox, path string, directories bool) []dropbox.Entry {
	// Use caching to reduce calls to the Dropbox API
	var cache_descriptor string
	if directories {
		cache_descriptor = "dirs:"
	} else {
		cache_descriptor = "files:"
	}
	s := []string{}
	s = append(s, cache_descriptor)
	s = append(s, path)
	cache_path := strings.Join(s, "")

	data, found := cache.Get(cache_path)
	if found {
		if cached_paths, ok := data.([]dropbox.Entry); ok {
			return (cached_paths)
		} else {
			log.Println("Error: Unable to retrieve from cache")
		}
	}

	entry, err := db.Metadata(path, true, false, "", "", 500)
	if err != nil {
		log.Println(err)
		return []dropbox.Entry{}
	}
	paths := make([]dropbox.Entry, 0)
	for i := 0; i < len(entry.Contents); i++ {
		entry := entry.Contents[i]
		if directories {
			if entry.IsDir {
				paths = append(paths, entry)
			}
		} else {
			if !entry.IsDir {
				include := true
				for _, lookup := range do_not_include {
					if strings.Contains(entry.Path, lookup) {
						include = false
					}
				}
				if include {
					paths = append(paths, entry)
				}
			}
		}
	}
	cache.Set(cache_path, paths, 0)
	return paths
}

/*
	Divine the correct target path from the provided info
*/
func get_target_path(arch string, version string) string {
	var target_path string
	if version == "nightly" {
		target_path = arch
	} else {
		directories := get_directories(cache_instance, db, arch)
		mTime := time.Time(dropbox.DBTime{})
		var latest_directory dropbox.Entry
		for _, dir := range directories {
			if strings.Contains(dir.Path, version) {
				if time.Time(dir.Modified).After(mTime) {
					mTime = time.Time(dir.Modified)
					latest_directory = dir
				}
			}
		}
		target_path = latest_directory.Path
	}
	return target_path
}

/*
	Returns a shared link to dropbox file
*/
func get_link(cache *cache.Cache, db *dropbox.Dropbox, path string) string {

	// Use caching to reduce calls to the Dropbox API
	cache_path := strings.Join([]string{"link", path}, ":")
	data, found := cache.Get(cache_path)
	if found {
		if cached, ok := data.(string); ok {
			return cached
		} else {
			log.Println("Error: Unable to retrieve from cache")
		}
	}

	link, err := db.Shares(path, false)
	if err != nil {
		log.Println(err)
		return ""
	}
	cache.Set(cache_path, link.URL, 0)
	return link.URL
}

func get_target(arch string, software string, version string, target string) (dropbox.Entry, bool) {
	if target == "latest" {
		return get_latest(arch, software, version)
	} else {
		target_path := get_target_path(arch, version)
		files := get_files(cache_instance, db, target_path)
		for _, file := range files {
			archive := new(Archive)
			archive = archive.Init(file.Path)
			if archive.Software == software {
				if archive.Tag == target {
					return file, true
				}
			}
		}
	}
	return dropbox.Entry{}, false
}

/*
	Use the arch, software and version to find the latest
*/
func get_latest(arch string, software string, version string) (dropbox.Entry, bool) {
	target_path := get_target_path(arch, version)

	s := []string{}
	s = append(s, software)
	s = append(s, "-")
	search := strings.Join(s, "")

	mTime := time.Time(dropbox.DBTime{})
	var latest_file dropbox.Entry
	files := get_files(cache_instance, db, target_path)
	for _, file := range files {
		if strings.Contains(file.Path, search) {
			if time.Time(file.Modified).After(mTime) {
				mTime = time.Time(file.Modified)
				latest_file = file
			}
		}
	}
	if latest_file.Path == "" {
		return latest_file, false
	} else {
		return latest_file, true
	}
}
<commit_msg>Added simple sorting for human users.<commit_after>package main

import (
	"fmt"
	"github.com/ant0ine/go-json-rest/rest"
	"github.com/pmylund/go-cache"
	"github.com/stacktic/dropbox"
	"gopkg.in/yaml.v2"
	"io/ioutil"
	"log"
	"net/http"
	"path/filepath"
	"sort"
    "strings"
    "time"
)

type Config struct {
	Api_secret    string
	Api_key       string
	Client_token  string
	Context_root  string
	Server_listen string
	Server_port   string
}

type Archive struct {
	Software string
	Date     string
	Version  string
	Tag      string
}

type BasicResult struct {
	Tag  string
	Date string
}

// Allow BasicResults to be sorted
func (slice BasicResults) Len() int {
    return len(slice)
}

func (slice BasicResults) Less(i, j int) bool {
    return slice[i].Date > slice[j].Date
}

func (slice BasicResults) Swap(i,j int) {
    slice[i], slice[j] = slice[j], slice[i]
}

type BasicResults []BasicResult

func (a Archive) Init(path string) *Archive {
	_, filename := filepath.Split(path)
	parts := strings.Split(filename, "-")
	a.Software = parts[0]
	a.Version = strings.Join([]string{parts[1], parts[2]}, " ")
	a.Date = strings.Join([]string{parts[3], parts[4], parts[5]}, "-")
	a.Tag = parts[6]
	return &a
}

// The struct to store the configuration data
var config Config

// 12 hour caching that cleans up every 15 minutes
var cache_instance *cache.Cache

//Link to dropbox
var db *dropbox.Dropbox

var do_not_include []string

func main() {
	var err error

	config = Config{}
	cache_instance = cache.New(12*time.Hour, 15*time.Minute)
	db = dropbox.NewDropbox()

	do_not_include = []string{}
	do_not_include = append(do_not_include, ".txt")

	data, read_err := ioutil.ReadFile("config.yml")
	if read_err != nil {
		log.Fatal(read_err)
	}

	yaml_err := yaml.Unmarshal(data, &config)
	if yaml_err != nil {
		log.Fatal("error: %v", yaml_err)
	}
	fmt.Printf("--- config_file dump:\n%s\n\n", data)
	fmt.Printf("--- config dump:\n%s\n\n", config)

	db.SetAppInfo(config.Api_key, config.Api_secret)
	if len(config.Client_token) >= 1 {
		db.SetAccessToken(config.Client_token)
	} else {
		if err = db.Auth(); err != nil {
			fmt.Println(err)
			return
		}
		config.Client_token = db.AccessToken()
		db.SetAccessToken(config.Client_token)
		d, marshal_err := yaml.Marshal(&config)
		if marshal_err != nil {
			log.Fatal("error: %v", marshal_err)
		}
		ioutil.WriteFile("config.yml", []byte(d), 0644)
	}

	// root_paths := get_directories(cache, db, "")
	// fmt.Printf("--- paths dump:\n%s\n\n", root_paths)

	// nightly_files := get_files(cache, db, "ARMv7")
	// fmt.Printf("--- paths dump:\n%s\n\n", nightly_files)

	// setup server to link
	api := rest.NewApi()
	api.Use(rest.DefaultDevStack...)
	router, err := rest.MakeRouter(
		rest.Get(strings.Join([]string{config.Context_root, ""}, "/"), list_arch),
		rest.Get(strings.Join([]string{config.Context_root, "#arch"}, "/"), list_software),
		rest.Get(strings.Join([]string{config.Context_root, "#arch/#software"}, "/"), list_versions),
		rest.Get(strings.Join([]string{config.Context_root, "#arch/#software/#version"}, "/"), list_targets),
		rest.Get(strings.Join([]string{config.Context_root, "#arch/#software/#version/#target"}, "/"), link_target),
	)
	if err != nil {
		log.Fatal(err)
	}
	api.SetApp(router)
	s := []string{}
	s = append(s, config.Server_listen)
	s = append(s, config.Server_port)
	server_listen := strings.Join(s, ":")
	log.Fatal(http.ListenAndServe(server_listen, api.MakeHandler()))
}

func list_arch(w rest.ResponseWriter, r *rest.Request) {
	// Use caching to reduce calls to the Dropbox API
	cache_path := "arches"
	data, found := cache_instance.Get(cache_path)
	if found {
		if cached, ok := data.([]string); ok {
			w.WriteJson(cached)
			return
		} else {
			log.Println("Error: Unable to retrieve from cache")
		}
	}

	arches := []string{}
	directories := get_directories(cache_instance, db, "/")
	for _, arch := range directories {
		arches = append(arches, strings.Replace(arch.Path, "/", "", -1))
	}
	cache_instance.Set(cache_path, arches, 0)
	w.WriteJson(arches)
}

func list_software(w rest.ResponseWriter, r *rest.Request) {
	arch := r.PathParam("arch")

	// Use caching to reduce calls to the Dropbox API
	cache_path := arch
	data, found := cache_instance.Get(cache_path)
	if found {
		if cached, ok := data.([]string); ok {
			w.WriteJson(cached)
			return
		} else {
			log.Println("Error: Unable to retrieve from cache")
		}
	}

	softwares := make(map[string]string)
	files := get_files(cache_instance, db, arch)
	for _, file := range files {
		archive := new(Archive)
		archive = archive.Init(file.Path)
		softwares[archive.Software] = ""
	}
	keys := make([]string, 0, len(softwares))
	for k := range softwares {
		keys = append(keys, k)
	}
	cache_instance.Set(cache_path, keys, 0)
	w.WriteJson(keys)
}

func list_versions(w rest.ResponseWriter, r *rest.Request) {
	w.WriteJson([]string{"nightly", "beta", "stable"})
}

func list_targets(w rest.ResponseWriter, r *rest.Request) {
	arch := r.PathParam("arch")
	software := r.PathParam("software")
	version := r.PathParam("version")

	// Doesn't need to be cached, as its calls are already cached.

	targets := BasicResults{}
    latest_date := time.Time{}
    target_path := get_target_path(arch, version)
	files := get_files(cache_instance, db, target_path)
	for _, file := range files {
		archive := new(Archive)
		archive = archive.Init(file.Path)
		if archive.Software == software {
            parsed_time, err := time.Parse("2006-01-02", archive.Date)
	        if err != nil {
		        log.Println(err)
		        parsed_time = time.Time{}
	        }
            if parsed_time.After(latest_date) {
                latest_date = parsed_time
            }
			targets = append(targets, BasicResult{archive.Tag, archive.Date})
		}
	}
	targets = append(targets, BasicResult{"latest", latest_date.Format("2006-01-02")})

    // Sort the targets by date descending.
    sort.Sort(targets)

    w.WriteJson(targets)
}

func link_target(w rest.ResponseWriter, r *rest.Request) {
	arch := r.PathParam("arch")
	software := r.PathParam("software")
	version := r.PathParam("version")
	target := r.PathParam("target")

	target_file, found := get_target(arch, software, version, target)
	if found {
		target_link := get_link(cache_instance, db, target_file.Path)
		w.Header().Set("Location", target_link)
		w.WriteHeader(302)
	} else {
		w.WriteHeader(404)
		w.WriteJson(map[string]string{"error": "Target Not Found"})
	}
}

/*
	Get only a slice of the directories at a path
*/
func get_directories(cache *cache.Cache, db *dropbox.Dropbox, path string) []dropbox.Entry {
	return get(cache, db, path, true)
}

/*
	Get only a slice of the directories at a path
*/
func get_files(cache *cache.Cache, db *dropbox.Dropbox, path string) []dropbox.Entry {
	return get(cache, db, path, false)
}

/*
	Actually get a list of directories or files from the dropbox connection
*/
func get(cache *cache.Cache, db *dropbox.Dropbox, path string, directories bool) []dropbox.Entry {
	// Use caching to reduce calls to the Dropbox API
	var cache_descriptor string
	if directories {
		cache_descriptor = "dirs:"
	} else {
		cache_descriptor = "files:"
	}
	s := []string{}
	s = append(s, cache_descriptor)
	s = append(s, path)
	cache_path := strings.Join(s, "")

	data, found := cache.Get(cache_path)
	if found {
		if cached_paths, ok := data.([]dropbox.Entry); ok {
			return (cached_paths)
		} else {
			log.Println("Error: Unable to retrieve from cache")
		}
	}

	entry, err := db.Metadata(path, true, false, "", "", 500)
	if err != nil {
		log.Println(err)
		return []dropbox.Entry{}
	}
	paths := make([]dropbox.Entry, 0)
	for i := 0; i < len(entry.Contents); i++ {
		entry := entry.Contents[i]
		if directories {
			if entry.IsDir {
				paths = append(paths, entry)
			}
		} else {
			if !entry.IsDir {
				include := true
				for _, lookup := range do_not_include {
					if strings.Contains(entry.Path, lookup) {
						include = false
					}
				}
				if include {
					paths = append(paths, entry)
				}
			}
		}
	}
	cache.Set(cache_path, paths, 0)
	return paths
}

/*
	Divine the correct target path from the provided info
*/
func get_target_path(arch string, version string) string {
	var target_path string
	if version == "nightly" {
		target_path = arch
	} else {
		directories := get_directories(cache_instance, db, arch)
		mTime := time.Time(dropbox.DBTime{})
		var latest_directory dropbox.Entry
		for _, dir := range directories {
			if strings.Contains(dir.Path, version) {
				if time.Time(dir.Modified).After(mTime) {
					mTime = time.Time(dir.Modified)
					latest_directory = dir
				}
			}
		}
		target_path = latest_directory.Path
	}
	return target_path
}

/*
	Returns a shared link to dropbox file
*/
func get_link(cache *cache.Cache, db *dropbox.Dropbox, path string) string {

	// Use caching to reduce calls to the Dropbox API
	cache_path := strings.Join([]string{"link", path}, ":")
	data, found := cache.Get(cache_path)
	if found {
		if cached, ok := data.(string); ok {
			return cached
		} else {
			log.Println("Error: Unable to retrieve from cache")
		}
	}

	link, err := db.Shares(path, false)
	if err != nil {
		log.Println(err)
		return ""
	}
	cache.Set(cache_path, link.URL, 0)
	return link.URL
}

func get_target(arch string, software string, version string, target string) (dropbox.Entry, bool) {
	if target == "latest" {
		return get_latest(arch, software, version)
	} else {
		target_path := get_target_path(arch, version)
		files := get_files(cache_instance, db, target_path)
		for _, file := range files {
			archive := new(Archive)
			archive = archive.Init(file.Path)
			if archive.Software == software {
				if archive.Tag == target {
					return file, true
				}
			}
		}
	}
	return dropbox.Entry{}, false
}

/*
	Use the arch, software and version to find the latest
*/
func get_latest(arch string, software string, version string) (dropbox.Entry, bool) {
	target_path := get_target_path(arch, version)

	s := []string{}
	s = append(s, software)
	s = append(s, "-")
	search := strings.Join(s, "")

	mTime := time.Time(dropbox.DBTime{})
	var latest_file dropbox.Entry
	files := get_files(cache_instance, db, target_path)
	for _, file := range files {
		if strings.Contains(file.Path, search) {
			if time.Time(file.Modified).After(mTime) {
				mTime = time.Time(file.Modified)
				latest_file = file
			}
		}
	}
	if latest_file.Path == "" {
		return latest_file, false
	} else {
		return latest_file, true
	}
}
<|endoftext|>


<commit_before>/*
Copyright IBM Corp. All Rights Reserved.

SPDX-License-Identifier: Apache-2.0
*/

package inproccontroller

import (
	"context"
	"fmt"

	"github.com/hyperledger/fabric/common/flogging"
	"github.com/hyperledger/fabric/core/chaincode/shim"
	"github.com/hyperledger/fabric/core/container"
	"github.com/hyperledger/fabric/core/container/ccintf"
	pb "github.com/hyperledger/fabric/protos/peer"
)

// ContainerType is the string which the inproc container type
// is registered with the container.VMController
const ContainerType = "SYSTEM"

type inprocContainer struct {
	ChaincodeSupport ccintf.CCSupport
	chaincode        shim.Chaincode
	running          bool
	args             []string
	env              []string
	stopChan         chan struct{}
}

var (
	inprocLogger = flogging.MustGetLogger("inproccontroller")

	// TODO this is a very hacky way to do testing, we should find other ways
	// to test, or not statically inject these depenencies.
	_shimStartInProc    = shim.StartInProc
	_inprocLoggerErrorf = inprocLogger.Errorf
)

// errors

//SysCCRegisteredErr registered error
type SysCCRegisteredErr string

func (s SysCCRegisteredErr) Error() string {
	return fmt.Sprintf("%s already registered", string(s))
}

// Registry stores registered system chaincodes.
// It implements container.VMProvider and scc.Registrar
type Registry struct {
	typeRegistry     map[string]*inprocContainer
	instRegistry     map[string]*inprocContainer
	ChaincodeSupport ccintf.CCSupport
}

// NewRegistry creates an initialized registry, ready to register system chaincodes.
// The returned *Registry is _not_ ready to use as is.  You must set the ChaincodeSupport
// as soon as one is available, before any chaincode invocations occur.  This is because
// the chaincode support used to be a latent dependency, snuck in on the context, but now
// it is being made an explicit part of the startup.
func NewRegistry() *Registry {
	return &Registry{
		typeRegistry: make(map[string]*inprocContainer),
		instRegistry: make(map[string]*inprocContainer),
	}
}

// NewVM creates an inproc VM instance
func (r *Registry) NewVM() container.VM {
	return NewInprocVM(r)
}

// Register registers system chaincode with given path. The deploy should be called to initialize
func (r *Registry) Register(ccid *ccintf.CCID, cc shim.Chaincode) error {
	name := ccid.GetName()
	inprocLogger.Debugf("Registering chaincode instance: %s", name)
	tmp := r.typeRegistry[name]
	if tmp != nil {
		return SysCCRegisteredErr(name)
	}

	r.typeRegistry[name] = &inprocContainer{chaincode: cc}
	return nil
}

// InprocVM is a vm. It is identified by a executable name
type InprocVM struct {
	registry *Registry
}

// NewInprocVM creates a new InprocVM
func NewInprocVM(r *Registry) *InprocVM {
	return &InprocVM{
		registry: r,
	}
}

func (vm *InprocVM) getInstance(ipctemplate *inprocContainer, instName string, args []string, env []string) (*inprocContainer, error) {
	ipc := vm.registry.instRegistry[instName]
	if ipc != nil {
		inprocLogger.Warningf("chaincode instance exists for %s", instName)
		return ipc, nil
	}
	ipc = &inprocContainer{
		ChaincodeSupport: vm.registry.ChaincodeSupport,
		args:             args,
		env:              env,
		chaincode:        ipctemplate.chaincode,
		stopChan:         make(chan struct{}),
	}
	vm.registry.instRegistry[instName] = ipc
	inprocLogger.Debugf("chaincode instance created for %s", instName)
	return ipc, nil
}

func (ipc *inprocContainer) launchInProc(id string, args []string, env []string) error {
	if ipc.ChaincodeSupport == nil {
		inprocLogger.Panicf("Chaincode support is nil, most likely you forgot to set it immediately after calling inproccontroller.NewRegsitry()")
	}

	peerRcvCCSend := make(chan *pb.ChaincodeMessage)
	ccRcvPeerSend := make(chan *pb.ChaincodeMessage)
	var err error
	ccchan := make(chan struct{}, 1)
	ccsupportchan := make(chan struct{}, 1)
	shimStartInProc := _shimStartInProc // shadow to avoid race in test
	go func() {
		defer close(ccchan)
		inprocLogger.Debugf("chaincode started for %s", id)
		if args == nil {
			args = ipc.args
		}
		if env == nil {
			env = ipc.env
		}
		err := shimStartInProc(env, args, ipc.chaincode, ccRcvPeerSend, peerRcvCCSend)
		if err != nil {
			err = fmt.Errorf("chaincode-support ended with err: %s", err)
			_inprocLoggerErrorf("%s", err)
		}
		inprocLogger.Debugf("chaincode ended for %s with err: %s", id, err)
	}()

	// shadow function to avoid data race
	inprocLoggerErrorf := _inprocLoggerErrorf
	go func() {
		defer close(ccsupportchan)
		inprocStream := newInProcStream(peerRcvCCSend, ccRcvPeerSend)
		inprocLogger.Debugf("chaincode-support started for  %s", id)
		err := ipc.ChaincodeSupport.HandleChaincodeStream(inprocStream)
		if err != nil {
			err = fmt.Errorf("chaincode ended with err: %s", err)
			inprocLoggerErrorf("%s", err)
		}
		inprocLogger.Debugf("chaincode-support ended for %s with err: %s", id, err)
	}()

	select {
	case <-ccchan:
		close(peerRcvCCSend)
		inprocLogger.Debugf("chaincode %s quit", id)
	case <-ccsupportchan:
		close(ccRcvPeerSend)
		inprocLogger.Debugf("chaincode support %s quit", id)
	case <-ipc.stopChan:
		close(ccRcvPeerSend)
		close(peerRcvCCSend)
		inprocLogger.Debugf("chaincode %s stopped", id)
	}
	return err
}

//Start starts a previously registered system codechain
func (vm *InprocVM) Start(ccid ccintf.CCID, args []string, env []string, filesToUpload map[string][]byte, builder container.Builder) error {
	path := ccid.GetName()

	ipctemplate := vm.registry.typeRegistry[path]

	if ipctemplate == nil {
		return fmt.Errorf(fmt.Sprintf("%s not registered", path))
	}

	instName := vm.GetVMName(ccid)

	ipc, err := vm.getInstance(ipctemplate, instName, args, env)

	if err != nil {
		return fmt.Errorf(fmt.Sprintf("could not create instance for %s", instName))
	}

	if ipc.running {
		return fmt.Errorf(fmt.Sprintf("chaincode running %s", path))
	}

	ipc.running = true

	go func() {
		defer func() {
			if r := recover(); r != nil {
				inprocLogger.Criticalf("caught panic from chaincode  %s", instName)
			}
		}()
		ipc.launchInProc(instName, args, env)
	}()

	return nil
}

//Stop stops a system codechain
func (vm *InprocVM) Stop(ccid ccintf.CCID, timeout uint, dontkill bool, dontremove bool) error {
	path := ccid.GetName()

	ipctemplate := vm.registry.typeRegistry[path]
	if ipctemplate == nil {
		return fmt.Errorf("%s not registered", path)
	}

	instName := vm.GetVMName(ccid)

	ipc := vm.registry.instRegistry[instName]

	if ipc == nil {
		return fmt.Errorf("%s not found", instName)
	}

	if !ipc.running {
		return fmt.Errorf("%s not running", instName)
	}

	ipc.stopChan <- struct{}{}

	delete(vm.registry.instRegistry, instName)
	//TODO stop
	return nil
}

// HealthCheck is provided in order to implement the VMProvider interface.
// It always returns nil..
func (vm *InprocVM) HealthCheck(ctx context.Context) error {
	return nil
}

// GetVMName ignores the peer and network name as it just needs to be unique in
// process.  It accepts a format function parameter to allow different
// formatting based on the desired use of the name.
func (vm *InprocVM) GetVMName(ccid ccintf.CCID) string {
	return ccid.GetName()
}
<commit_msg>[FAB-14257] mutex around inproccontroller maps<commit_after>/*
Copyright IBM Corp. All Rights Reserved.

SPDX-License-Identifier: Apache-2.0
*/

package inproccontroller

import (
	"context"
	"fmt"
	"sync"

	"github.com/hyperledger/fabric/common/flogging"
	"github.com/hyperledger/fabric/core/chaincode/shim"
	"github.com/hyperledger/fabric/core/container"
	"github.com/hyperledger/fabric/core/container/ccintf"
	pb "github.com/hyperledger/fabric/protos/peer"
)

// ContainerType is the string which the inproc container type
// is registered with the container.VMController
const ContainerType = "SYSTEM"

type inprocContainer struct {
	ChaincodeSupport ccintf.CCSupport
	chaincode        shim.Chaincode
	running          bool
	args             []string
	env              []string
	stopChan         chan struct{}
}

var (
	inprocLogger = flogging.MustGetLogger("inproccontroller")

	// TODO this is a very hacky way to do testing, we should find other ways
	// to test, or not statically inject these depenencies.
	_shimStartInProc    = shim.StartInProc
	_inprocLoggerErrorf = inprocLogger.Errorf
)

// errors

//SysCCRegisteredErr registered error
type SysCCRegisteredErr string

func (s SysCCRegisteredErr) Error() string {
	return fmt.Sprintf("%s already registered", string(s))
}

// Registry stores registered system chaincodes.
// It implements container.VMProvider and scc.Registrar
type Registry struct {
	mutex        sync.Mutex
	typeRegistry map[string]*inprocContainer
	instRegistry map[string]*inprocContainer

	ChaincodeSupport ccintf.CCSupport
}

// NewRegistry creates an initialized registry, ready to register system chaincodes.
// The returned *Registry is _not_ ready to use as is.  You must set the ChaincodeSupport
// as soon as one is available, before any chaincode invocations occur.  This is because
// the chaincode support used to be a latent dependency, snuck in on the context, but now
// it is being made an explicit part of the startup.
func NewRegistry() *Registry {
	return &Registry{
		typeRegistry: make(map[string]*inprocContainer),
		instRegistry: make(map[string]*inprocContainer),
	}
}

// NewVM creates an inproc VM instance
func (r *Registry) NewVM() container.VM {
	return NewInprocVM(r)
}

// Register registers system chaincode with given path. The deploy should be called to initialize
func (r *Registry) Register(ccid *ccintf.CCID, cc shim.Chaincode) error {
	r.mutex.Lock()
	defer r.mutex.Unlock()

	name := ccid.GetName()
	inprocLogger.Debugf("Registering chaincode instance: %s", name)
	tmp := r.typeRegistry[name]
	if tmp != nil {
		return SysCCRegisteredErr(name)
	}

	r.typeRegistry[name] = &inprocContainer{chaincode: cc}
	return nil
}

func (r *Registry) getType(name string) *inprocContainer {
	r.mutex.Lock()
	defer r.mutex.Unlock()
	return r.typeRegistry[name]
}

func (r *Registry) getInstance(name string) *inprocContainer {
	r.mutex.Lock()
	defer r.mutex.Unlock()
	return r.instRegistry[name]
}

func (r *Registry) setInstance(name string, inst *inprocContainer) {
	r.mutex.Lock()
	defer r.mutex.Unlock()
	r.instRegistry[name] = inst
}

func (r *Registry) removeInstance(name string) {
	r.mutex.Lock()
	defer r.mutex.Unlock()
	delete(r.instRegistry, name)
}

// InprocVM is a vm. It is identified by a executable name
type InprocVM struct {
	registry *Registry
}

// NewInprocVM creates a new InprocVM
func NewInprocVM(r *Registry) *InprocVM {
	return &InprocVM{
		registry: r,
	}
}

func (vm *InprocVM) getInstance(ipctemplate *inprocContainer, instName string, args []string, env []string) (*inprocContainer, error) {
	ipc := vm.registry.getInstance(instName)
	if ipc != nil {
		inprocLogger.Warningf("chaincode instance exists for %s", instName)
		return ipc, nil
	}
	ipc = &inprocContainer{
		ChaincodeSupport: vm.registry.ChaincodeSupport,
		args:             args,
		env:              env,
		chaincode:        ipctemplate.chaincode,
		stopChan:         make(chan struct{}),
	}
	vm.registry.setInstance(instName, ipc)
	inprocLogger.Debugf("chaincode instance created for %s", instName)
	return ipc, nil
}

func (ipc *inprocContainer) launchInProc(id string, args []string, env []string) error {
	if ipc.ChaincodeSupport == nil {
		inprocLogger.Panicf("Chaincode support is nil, most likely you forgot to set it immediately after calling inproccontroller.NewRegsitry()")
	}

	peerRcvCCSend := make(chan *pb.ChaincodeMessage)
	ccRcvPeerSend := make(chan *pb.ChaincodeMessage)
	var err error
	ccchan := make(chan struct{}, 1)
	ccsupportchan := make(chan struct{}, 1)
	shimStartInProc := _shimStartInProc // shadow to avoid race in test
	go func() {
		defer close(ccchan)
		inprocLogger.Debugf("chaincode started for %s", id)
		if args == nil {
			args = ipc.args
		}
		if env == nil {
			env = ipc.env
		}
		err := shimStartInProc(env, args, ipc.chaincode, ccRcvPeerSend, peerRcvCCSend)
		if err != nil {
			err = fmt.Errorf("chaincode-support ended with err: %s", err)
			_inprocLoggerErrorf("%s", err)
		}
		inprocLogger.Debugf("chaincode ended for %s with err: %s", id, err)
	}()

	// shadow function to avoid data race
	inprocLoggerErrorf := _inprocLoggerErrorf
	go func() {
		defer close(ccsupportchan)
		inprocStream := newInProcStream(peerRcvCCSend, ccRcvPeerSend)
		inprocLogger.Debugf("chaincode-support started for  %s", id)
		err := ipc.ChaincodeSupport.HandleChaincodeStream(inprocStream)
		if err != nil {
			err = fmt.Errorf("chaincode ended with err: %s", err)
			inprocLoggerErrorf("%s", err)
		}
		inprocLogger.Debugf("chaincode-support ended for %s with err: %s", id, err)
	}()

	select {
	case <-ccchan:
		close(peerRcvCCSend)
		inprocLogger.Debugf("chaincode %s quit", id)
	case <-ccsupportchan:
		close(ccRcvPeerSend)
		inprocLogger.Debugf("chaincode support %s quit", id)
	case <-ipc.stopChan:
		close(ccRcvPeerSend)
		close(peerRcvCCSend)
		inprocLogger.Debugf("chaincode %s stopped", id)
	}
	return err
}

//Start starts a previously registered system codechain
func (vm *InprocVM) Start(ccid ccintf.CCID, args []string, env []string, filesToUpload map[string][]byte, builder container.Builder) error {
	path := ccid.GetName()

	ipctemplate := vm.registry.getType(path)
	if ipctemplate == nil {
		return fmt.Errorf(fmt.Sprintf("%s not registered", path))
	}

	instName := vm.GetVMName(ccid)

	ipc, err := vm.getInstance(ipctemplate, instName, args, env)

	if err != nil {
		return fmt.Errorf(fmt.Sprintf("could not create instance for %s", instName))
	}

	if ipc.running {
		return fmt.Errorf(fmt.Sprintf("chaincode running %s", path))
	}

	ipc.running = true

	go func() {
		defer func() {
			if r := recover(); r != nil {
				inprocLogger.Criticalf("caught panic from chaincode  %s", instName)
			}
		}()
		ipc.launchInProc(instName, args, env)
	}()

	return nil
}

//Stop stops a system codechain
func (vm *InprocVM) Stop(ccid ccintf.CCID, timeout uint, dontkill bool, dontremove bool) error {
	path := ccid.GetName()

	ipctemplate := vm.registry.getType(path)
	if ipctemplate == nil {
		return fmt.Errorf("%s not registered", path)
	}

	instName := vm.GetVMName(ccid)

	ipc := vm.registry.getInstance(instName)
	if ipc == nil {
		return fmt.Errorf("%s not found", instName)
	}

	if !ipc.running {
		return fmt.Errorf("%s not running", instName)
	}

	ipc.stopChan <- struct{}{}
	vm.registry.removeInstance(instName)
	//TODO stop
	return nil
}

// HealthCheck is provided in order to implement the VMProvider interface.
// It always returns nil..
func (vm *InprocVM) HealthCheck(ctx context.Context) error {
	return nil
}

// GetVMName ignores the peer and network name as it just needs to be unique in
// process.  It accepts a format function parameter to allow different
// formatting based on the desired use of the name.
func (vm *InprocVM) GetVMName(ccid ccintf.CCID) string {
	return ccid.GetName()
}
<|endoftext|>


<commit_before>package httptreemux

import (
	"bufio"
	"html/template"
	"net/http"
	"os"
	"runtime"
	"strings"
)

// SimplePanicHandler just returns error 500.
func SimplePanicHandler(w http.ResponseWriter, r *http.Request, err interface{}) {
	w.WriteHeader(http.StatusInternalServerError)
}

// ShowErrorsPanicHandler prints a nice representation of an error to the browser.
// This was taken from github.com/gocraft/web, which adapted it from the Traffic project.
func ShowErrorsPanicHandler(w http.ResponseWriter, r *http.Request, err interface{}) {
	const size = 4096
	stack := make([]byte, size)
	stack = stack[:runtime.Stack(stack, false)]
	renderPrettyError(w, r, err, stack)
}

func renderPrettyError(rw http.ResponseWriter, req *http.Request, err interface{}, stack []byte) {
	_, filePath, line, _ := runtime.Caller(5)

	data := map[string]interface{}{
		"Error":    err,
		"Stack":    string(stack),
		"Params":   req.URL.Query(),
		"Method":   req.Method,
		"FilePath": filePath,
		"Line":     line,
		"Lines":    readErrorFileLines(filePath, line),
	}

	rw.Header().Set("Content-Type", "text/html")
	rw.WriteHeader(http.StatusInternalServerError)

	tpl := template.Must(template.New("ErrorPage").Parse(panicPageTpl))
	tpl.Execute(rw, data)
}

func readErrorFileLines(filePath string, errorLine int) map[int]string {
	lines := make(map[int]string)

	file, err := os.Open(filePath)
	if err != nil {
		return lines
	}

	defer file.Close()

	reader := bufio.NewReader(file)
	currentLine := 0
	for {
		line, err := reader.ReadString('\n')
		if err != nil || currentLine > errorLine+5 {
			break
		}

		currentLine++

		if currentLine >= errorLine-5 {
			lines[currentLine] = strings.Replace(line, "\n", "", -1)
		}
	}

	return lines
}

const panicPageTpl string = `

      <title>Panic</title>
      <meta http-equiv="Content-Type" content="text/html; charset=utf-8">

Error

{{ .Error }}

        In {{ .FilePath }}:{{ .Line }}

{{ range $lineNumber, $line :=  .Lines }}{{ $lineNumber }}{{ end }}

{{ range $lineNumber, $line :=  .Lines }}{{ $line }}
{{ end }}

Stack

{{ .Stack }}

Request

Method: {{ .Method }}

Paramenters:

        {{ range $key, $value := .Params }}

{{ $key }}: {{ $value }}

        {{ end }}

  `
<commit_msg>Unwind stack 4 levels in panic handler<commit_after>package httptreemux

import (
	"bufio"
	"html/template"
	"net/http"
	"os"
	"runtime"
	"strings"
)

// SimplePanicHandler just returns error 500.
func SimplePanicHandler(w http.ResponseWriter, r *http.Request, err interface{}) {
	w.WriteHeader(http.StatusInternalServerError)
}

// ShowErrorsPanicHandler prints a nice representation of an error to the browser.
// This was taken from github.com/gocraft/web, which adapted it from the Traffic project.
func ShowErrorsPanicHandler(w http.ResponseWriter, r *http.Request, err interface{}) {
	const size = 4096
	stack := make([]byte, size)
	stack = stack[:runtime.Stack(stack, false)]
	renderPrettyError(w, r, err, stack)
}

func renderPrettyError(rw http.ResponseWriter, req *http.Request, err interface{}, stack []byte) {
	_, filePath, line, _ := runtime.Caller(4)

	data := map[string]interface{}{
		"Error":    err,
		"Stack":    string(stack),
		"Params":   req.URL.Query(),
		"Method":   req.Method,
		"FilePath": filePath,
		"Line":     line,
		"Lines":    readErrorFileLines(filePath, line),
	}

	rw.Header().Set("Content-Type", "text/html")
	rw.WriteHeader(http.StatusInternalServerError)

	tpl := template.Must(template.New("ErrorPage").Parse(panicPageTpl))
	tpl.Execute(rw, data)
}

func readErrorFileLines(filePath string, errorLine int) map[int]string {
	lines := make(map[int]string)

	file, err := os.Open(filePath)
	if err != nil {
		return lines
	}

	defer file.Close()

	reader := bufio.NewReader(file)
	currentLine := 0
	for {
		line, err := reader.ReadString('\n')
		if err != nil || currentLine > errorLine+5 {
			break
		}

		currentLine++

		if currentLine >= errorLine-5 {
			lines[currentLine] = strings.Replace(line, "\n", "", -1)
		}
	}

	return lines
}

const panicPageTpl string = `

      <title>Panic</title>
      <meta http-equiv="Content-Type" content="text/html; charset=utf-8">

Error

{{ .Error }}

        In {{ .FilePath }}:{{ .Line }}

{{ range $lineNumber, $line :=  .Lines }}{{ $lineNumber }}{{ end }}

{{ range $lineNumber, $line :=  .Lines }}{{ $line }}
{{ end }}

Stack

{{ .Stack }}

Request

Method: {{ .Method }}

Paramenters:

        {{ range $key, $value := .Params }}

{{ $key }}: {{ $value }}

        {{ end }}

  `
<|endoftext|>


<commit_before>package worker

import (
	"bytes"
	"crypto/tls"
	"fmt"
	"io/ioutil"
	"net/http"
	"net/url"
	"time"

	gocontext "context"

	"github.com/travis-ci/worker/config"
	"github.com/travis-ci/worker/metrics"
)

// A BuildScriptGeneratorError is sometimes used by the Generate method on a
// BuildScriptGenerator to return more metadata about an error.
type BuildScriptGeneratorError struct {
	error

	// true when this error can be recovered by retrying later
	Recover bool
}

// A BuildScriptGenerator generates a build script for a given job payload.
type BuildScriptGenerator interface {
	Generate(gocontext.Context, Job) ([]byte, error)
}

type webBuildScriptGenerator struct {
	URL               string
	aptCacheHost      string
	npmCacheHost      string
	paranoid          bool
	fixResolvConf     bool
	fixEtcHosts       bool
	cacheType         string
	cacheFetchTimeout int
	cachePushTimeout  int
	s3CacheOptions    s3BuildCacheOptions

	httpClient *http.Client
}

type s3BuildCacheOptions struct {
	scheme          string
	region          string
	bucket          string
	accessKeyID     string
	secretAccessKey string
}

// NewBuildScriptGenerator creates a generator backed by an HTTP API.
func NewBuildScriptGenerator(cfg *config.Config) BuildScriptGenerator {
	return &webBuildScriptGenerator{
		URL:               cfg.BuildAPIURI,
		aptCacheHost:      cfg.BuildAptCache,
		npmCacheHost:      cfg.BuildNpmCache,
		paranoid:          cfg.BuildParanoid,
		fixResolvConf:     cfg.BuildFixResolvConf,
		fixEtcHosts:       cfg.BuildFixEtcHosts,
		cacheType:         cfg.BuildCacheType,
		cacheFetchTimeout: int(cfg.BuildCacheFetchTimeout.Seconds()),
		cachePushTimeout:  int(cfg.BuildCachePushTimeout.Seconds()),
		s3CacheOptions: s3BuildCacheOptions{
			scheme:          cfg.BuildCacheS3Scheme,
			region:          cfg.BuildCacheS3Region,
			bucket:          cfg.BuildCacheS3Bucket,
			accessKeyID:     cfg.BuildCacheS3AccessKeyID,
			secretAccessKey: cfg.BuildCacheS3SecretAccessKey,
		},
		httpClient: &http.Client{
			Transport: &http.Transport{
				TLSClientConfig: &tls.Config{
					InsecureSkipVerify: cfg.BuildAPIInsecureSkipVerify,
				},
			},
		},
	}
}

func (g *webBuildScriptGenerator) Generate(ctx gocontext.Context, job Job) ([]byte, error) {
	payload := job.RawPayload()

	if g.aptCacheHost != "" {
		payload.SetPath([]string{"hosts", "apt_cache"}, g.aptCacheHost)
	}
	if g.npmCacheHost != "" {
		payload.SetPath([]string{"hosts", "npm_cache"}, g.npmCacheHost)
	}

	payload.Set("paranoid", g.paranoid)
	payload.Set("fix_resolv_conf", g.fixResolvConf)
	payload.Set("fix_etc_hosts", g.fixEtcHosts)

	if g.cacheType != "" {
		payload.SetPath([]string{"cache_options", "type"}, g.cacheType)
		payload.SetPath([]string{"cache_options", "fetch_timeout"}, g.cacheFetchTimeout)
		payload.SetPath([]string{"cache_options", "push_timeout"}, g.cachePushTimeout)
		payload.SetPath([]string{"cache_options", "s3", "scheme"}, g.s3CacheOptions.scheme)
		payload.SetPath([]string{"cache_options", "s3", "region"}, g.s3CacheOptions.region)
		payload.SetPath([]string{"cache_options", "s3", "bucket"}, g.s3CacheOptions.bucket)
		payload.SetPath([]string{"cache_options", "s3", "access_key_id"}, g.s3CacheOptions.accessKeyID)
		payload.SetPath([]string{"cache_options", "s3", "secret_access_key"}, g.s3CacheOptions.secretAccessKey)
	}

	b, err := payload.Encode()
	if err != nil {
		return nil, err
	}

	var token string
	u, err := url.Parse(g.URL)
	if err != nil {
		return nil, err
	}
	if u.User != nil {
		token = u.User.Username()
		u.User = nil
	}

	buf := bytes.NewBuffer(b)
	req, err := http.NewRequest("POST", u.String(), buf)
	if err != nil {
		return nil, err
	}
	if token != "" {
		req.Header.Set("Authorization", "token "+token)
	}
	req.Header.Set("User-Agent", fmt.Sprintf("worker-go v=%v rev=%v d=%v", VersionString, RevisionString, GeneratedString))
	req.Header.Set("Content-Type", "application/json")

	startRequest := time.Now()

	resp, err := g.httpClient.Do(req)
	if err != nil {
		return nil, err
	}
	defer resp.Body.Close()
	metrics.TimeSince("worker.job.script.api", startRequest)

	body, err := ioutil.ReadAll(resp.Body)
	if err != nil {
		return nil, err
	}

	if resp.StatusCode >= 500 {
		return nil, BuildScriptGeneratorError{error: fmt.Errorf("server error: %q", string(body)), Recover: true}
	} else if resp.StatusCode >= 400 {
		return nil, BuildScriptGeneratorError{error: fmt.Errorf("client error: %q", string(body)), Recover: false}
	}

	return body, nil
}
<commit_msg>Add job_id to build script request query string for correlation<commit_after>package worker

import (
	"bytes"
	"crypto/tls"
	"fmt"
	"io/ioutil"
	"net/http"
	"net/url"
	"time"

	gocontext "context"

	"github.com/travis-ci/worker/config"
	"github.com/travis-ci/worker/metrics"
)

// A BuildScriptGeneratorError is sometimes used by the Generate method on a
// BuildScriptGenerator to return more metadata about an error.
type BuildScriptGeneratorError struct {
	error

	// true when this error can be recovered by retrying later
	Recover bool
}

// A BuildScriptGenerator generates a build script for a given job payload.
type BuildScriptGenerator interface {
	Generate(gocontext.Context, Job) ([]byte, error)
}

type webBuildScriptGenerator struct {
	URL               string
	aptCacheHost      string
	npmCacheHost      string
	paranoid          bool
	fixResolvConf     bool
	fixEtcHosts       bool
	cacheType         string
	cacheFetchTimeout int
	cachePushTimeout  int
	s3CacheOptions    s3BuildCacheOptions

	httpClient *http.Client
}

type s3BuildCacheOptions struct {
	scheme          string
	region          string
	bucket          string
	accessKeyID     string
	secretAccessKey string
}

// NewBuildScriptGenerator creates a generator backed by an HTTP API.
func NewBuildScriptGenerator(cfg *config.Config) BuildScriptGenerator {
	return &webBuildScriptGenerator{
		URL:               cfg.BuildAPIURI,
		aptCacheHost:      cfg.BuildAptCache,
		npmCacheHost:      cfg.BuildNpmCache,
		paranoid:          cfg.BuildParanoid,
		fixResolvConf:     cfg.BuildFixResolvConf,
		fixEtcHosts:       cfg.BuildFixEtcHosts,
		cacheType:         cfg.BuildCacheType,
		cacheFetchTimeout: int(cfg.BuildCacheFetchTimeout.Seconds()),
		cachePushTimeout:  int(cfg.BuildCachePushTimeout.Seconds()),
		s3CacheOptions: s3BuildCacheOptions{
			scheme:          cfg.BuildCacheS3Scheme,
			region:          cfg.BuildCacheS3Region,
			bucket:          cfg.BuildCacheS3Bucket,
			accessKeyID:     cfg.BuildCacheS3AccessKeyID,
			secretAccessKey: cfg.BuildCacheS3SecretAccessKey,
		},
		httpClient: &http.Client{
			Transport: &http.Transport{
				TLSClientConfig: &tls.Config{
					InsecureSkipVerify: cfg.BuildAPIInsecureSkipVerify,
				},
			},
		},
	}
}

func (g *webBuildScriptGenerator) Generate(ctx gocontext.Context, job Job) ([]byte, error) {
	payload := job.RawPayload()

	if g.aptCacheHost != "" {
		payload.SetPath([]string{"hosts", "apt_cache"}, g.aptCacheHost)
	}
	if g.npmCacheHost != "" {
		payload.SetPath([]string{"hosts", "npm_cache"}, g.npmCacheHost)
	}

	payload.Set("paranoid", g.paranoid)
	payload.Set("fix_resolv_conf", g.fixResolvConf)
	payload.Set("fix_etc_hosts", g.fixEtcHosts)

	if g.cacheType != "" {
		payload.SetPath([]string{"cache_options", "type"}, g.cacheType)
		payload.SetPath([]string{"cache_options", "fetch_timeout"}, g.cacheFetchTimeout)
		payload.SetPath([]string{"cache_options", "push_timeout"}, g.cachePushTimeout)
		payload.SetPath([]string{"cache_options", "s3", "scheme"}, g.s3CacheOptions.scheme)
		payload.SetPath([]string{"cache_options", "s3", "region"}, g.s3CacheOptions.region)
		payload.SetPath([]string{"cache_options", "s3", "bucket"}, g.s3CacheOptions.bucket)
		payload.SetPath([]string{"cache_options", "s3", "access_key_id"}, g.s3CacheOptions.accessKeyID)
		payload.SetPath([]string{"cache_options", "s3", "secret_access_key"}, g.s3CacheOptions.secretAccessKey)
	}

	b, err := payload.Encode()
	if err != nil {
		return nil, err
	}

	var token string
	u, err := url.Parse(g.URL)
	if err != nil {
		return nil, err
	}
	if u.User != nil {
		token = u.User.Username()
		u.User = nil
	}

	jp := job.Payload()
	if jp != nil {
		q := u.Query()
		q.Add("job_id", jp.Job.ID)
		u.RawQuery = q.Encode()
	}

	buf := bytes.NewBuffer(b)
	req, err := http.NewRequest("POST", u.String(), buf)
	if err != nil {
		return nil, err
	}
	if token != "" {
		req.Header.Set("Authorization", "token "+token)
	}
	req.Header.Set("User-Agent", fmt.Sprintf("worker-go v=%v rev=%v d=%v", VersionString, RevisionString, GeneratedString))
	req.Header.Set("Content-Type", "application/json")

	startRequest := time.Now()

	resp, err := g.httpClient.Do(req)
	if err != nil {
		return nil, err
	}
	defer resp.Body.Close()
	metrics.TimeSince("worker.job.script.api", startRequest)

	body, err := ioutil.ReadAll(resp.Body)
	if err != nil {
		return nil, err
	}

	if resp.StatusCode >= 500 {
		return nil, BuildScriptGeneratorError{error: fmt.Errorf("server error: %q", string(body)), Recover: true}
	} else if resp.StatusCode >= 400 {
		return nil, BuildScriptGeneratorError{error: fmt.Errorf("client error: %q", string(body)), Recover: false}
	}

	return body, nil
}
<|endoftext|>


<commit_before>// Copyright 2016 Palantir Technologies, 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 git

import (
	"os/exec"
	"regexp"
	"strings"

	"github.com/pkg/errors"
)

const Unspecified = "Unspecified"

type ProjectInfo struct {
	Version  string
	Branch   string
	Revision string
}

const snapshotRegexp = `.+g[-+.]?[a-fA-F0-9]{3,}$`

func IsSnapshotVersion(version string) bool {
	return regexp.MustCompile(snapshotRegexp).MatchString(version)
}

func NewProjectInfo(gitDir string) (ProjectInfo, error) {
	version, err := ProjectVersion(gitDir)
	if err != nil {
		return ProjectInfo{}, err
	}

	branch, err := ProjectBranch(gitDir)
	if err != nil {
		return ProjectInfo{}, err
	}

	revision, err := ProjectRevision(gitDir)
	if err != nil {
		return ProjectInfo{}, err
	}

	return ProjectInfo{
		Version:  version,
		Branch:   branch,
		Revision: revision,
	}, nil
}

// ProjectVersion returns the version string for the git repository that the provided directory is in. The output is the output
// of "git describe --tags" followed by ".dirty" if the repository currently has any uncommitted changes. Returns
// an error if the provided path is not in a git root or if the git repository has no commits or no tags.
func ProjectVersion(gitDir string) (string, error) {
	tags, err := tags(gitDir)
	if err != nil {
		return "", err
	}

	// if no tags exist, return Unspecified as the version
	if tags == "" {
		return Unspecified, nil
	}

	result, err := trimmedCombinedGitCmdOutput(gitDir, "describe", "--tags", "--first-parent")
	if err != nil {
		return "", err
	}

	// trim "v" prefix in tags
	if strings.HasPrefix(result, "v") {
		result = result[1:]
	}

	// handle untracked files as well as "actual" dirtiness
	dirtyFiles, err := trimmedCombinedGitCmdOutput(gitDir, "status", "--porcelain")
	if err != nil {
		return "", err
	}
	if dirtyFiles != "" {
		result += ".dirty"
	}
	return result, nil
}

func ProjectBranch(gitDir string) (string, error) {
	tags, err := tags(gitDir)
	if err != nil {
		return "", err
	}

	// if no tags exist, return Unspecified as the branch
	if tags == "" {
		return Unspecified, nil
	}

	branch, err := branch(gitDir)
	if err != nil {
		return "", err
	}

	if strings.HasPrefix(branch, "v") {
		branch = branch[1:]
	}

	return branch, nil
}

func ProjectRevision(gitDir string) (string, error) {
	tags, err := tags(gitDir)
	if err != nil {
		return "", err
	}

	// if no tags exist, return revision count from first commit
	if tags == "" {
		return trimmedCombinedGitCmdOutput(gitDir, "rev-list", "HEAD", "--count")
	}

	branch, err := branch(gitDir)
	if err != nil {
		return "", err
	}
	return trimmedCombinedGitCmdOutput(gitDir, "rev-list", branch+"..HEAD", "--count")
}

func tags(gitDir string) (string, error) {
	return trimmedCombinedGitCmdOutput(gitDir, "tag", "-l")
}

func branch(gitDir string) (string, error) {
	return trimmedCombinedGitCmdOutput(gitDir, "describe", "--abbrev=0", "--tags", "--first-parent")
}

func trimmedCombinedGitCmdOutput(gitDir string, args ...string) (string, error) {
	cmd := exec.Command("git", args...)
	cmd.Dir = gitDir
	out, err := cmd.CombinedOutput()
	if err != nil {
		return "", errors.Wrapf(err, "Command %v failed. Output: %v", cmd.Args, string(out))
	}
	return strings.TrimSpace(string(out)), err
}
<commit_msg>Fix regression in value of 'unspecified' string (#198)<commit_after>// Copyright 2016 Palantir Technologies, 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 git

import (
	"os/exec"
	"regexp"
	"strings"

	"github.com/pkg/errors"
)

const Unspecified = "unspecified"

type ProjectInfo struct {
	Version  string
	Branch   string
	Revision string
}

const snapshotRegexp = `.+g[-+.]?[a-fA-F0-9]{3,}$`

func IsSnapshotVersion(version string) bool {
	return regexp.MustCompile(snapshotRegexp).MatchString(version)
}

func NewProjectInfo(gitDir string) (ProjectInfo, error) {
	version, err := ProjectVersion(gitDir)
	if err != nil {
		return ProjectInfo{}, err
	}

	branch, err := ProjectBranch(gitDir)
	if err != nil {
		return ProjectInfo{}, err
	}

	revision, err := ProjectRevision(gitDir)
	if err != nil {
		return ProjectInfo{}, err
	}

	return ProjectInfo{
		Version:  version,
		Branch:   branch,
		Revision: revision,
	}, nil
}

// ProjectVersion returns the version string for the git repository that the provided directory is in. The output is the output
// of "git describe --tags" followed by ".dirty" if the repository currently has any uncommitted changes. Returns
// an error if the provided path is not in a git root or if the git repository has no commits or no tags.
func ProjectVersion(gitDir string) (string, error) {
	tags, err := tags(gitDir)
	if err != nil {
		return "", err
	}

	// if no tags exist, return Unspecified as the version
	if tags == "" {
		return Unspecified, nil
	}

	result, err := trimmedCombinedGitCmdOutput(gitDir, "describe", "--tags", "--first-parent")
	if err != nil {
		return "", err
	}

	// trim "v" prefix in tags
	if strings.HasPrefix(result, "v") {
		result = result[1:]
	}

	// handle untracked files as well as "actual" dirtiness
	dirtyFiles, err := trimmedCombinedGitCmdOutput(gitDir, "status", "--porcelain")
	if err != nil {
		return "", err
	}
	if dirtyFiles != "" {
		result += ".dirty"
	}
	return result, nil
}

func ProjectBranch(gitDir string) (string, error) {
	tags, err := tags(gitDir)
	if err != nil {
		return "", err
	}

	// if no tags exist, return Unspecified as the branch
	if tags == "" {
		return Unspecified, nil
	}

	branch, err := branch(gitDir)
	if err != nil {
		return "", err
	}

	if strings.HasPrefix(branch, "v") {
		branch = branch[1:]
	}

	return branch, nil
}

func ProjectRevision(gitDir string) (string, error) {
	tags, err := tags(gitDir)
	if err != nil {
		return "", err
	}

	// if no tags exist, return revision count from first commit
	if tags == "" {
		return trimmedCombinedGitCmdOutput(gitDir, "rev-list", "HEAD", "--count")
	}

	branch, err := branch(gitDir)
	if err != nil {
		return "", err
	}
	return trimmedCombinedGitCmdOutput(gitDir, "rev-list", branch+"..HEAD", "--count")
}

func tags(gitDir string) (string, error) {
	return trimmedCombinedGitCmdOutput(gitDir, "tag", "-l")
}

func branch(gitDir string) (string, error) {
	return trimmedCombinedGitCmdOutput(gitDir, "describe", "--abbrev=0", "--tags", "--first-parent")
}

func trimmedCombinedGitCmdOutput(gitDir string, args ...string) (string, error) {
	cmd := exec.Command("git", args...)
	cmd.Dir = gitDir
	out, err := cmd.CombinedOutput()
	if err != nil {
		return "", errors.Wrapf(err, "Command %v failed. Output: %v", cmd.Args, string(out))
	}
	return strings.TrimSpace(string(out)), err
}
<|endoftext|>


<commit_before>// Copyright (c) 2013-2014 by Michael Dvorkin. All Rights Reserved.
// Use of this source code is governed by a MIT-style license that can
// be found in the LICENSE file.

package donna

func (e *Evaluation) analyzePieces() {
	p := e.position
	var white, black [4]Score

	if Settings.Trace {
		defer func() {
			var his, her Score
			e.checkpoint(`+Pieces`,  Total{*his.add(white[0]).add(white[1]).add(white[2]).add(white[3]),
				*her.add(black[0]).add(black[1]).add(black[2]).add(black[3])})
			e.checkpoint(`-Knights`, Total{white[0], black[0]})
			e.checkpoint(`-Bishops`, Total{white[1], black[1]})
			e.checkpoint(`-Rooks`,   Total{white[2], black[2]})
			e.checkpoint(`-Queens`,  Total{white[3], black[3]})
		}()
	}

	// Mobility mask for both sides excludes a) squares attacked by enemy's
	// pawns and b) squares occupied by own pawns and king.
	maskMobile := [2]Bitmask{
		^(e.attacks[BlackPawn] | p.outposts[Pawn] | p.outposts[King]),
		^(e.attacks[Pawn] | p.outposts[BlackPawn] | p.outposts[BlackKing]),
	}

	// Initialize king fort bitmasks only when we need them.
	// TODO: rearrange e.material.flags and invoke e.enemyKingThreat()
	//       only when necessary.
	if e.material.flags & whiteKingSafety != 0 {
		e.safety[White].fort = e.setupFort(White)
	}
	if e.material.flags & blackKingSafety != 0 {
		e.safety[Black].fort = e.setupFort(Black)
	}

	// Evaluate white pieces except queen.
	if p.count[Knight] > 0 {
		white[0] = e.knights(White, maskMobile[White])
	}
	if p.count[Bishop] > 0 {
		white[1] = e.bishops(White, maskMobile[White])
	}
	if p.count[Rook] > 0 {
		white[2] = e.rooks(White, maskMobile[White])
	}

	// Evaluate black pieces except queen.
	if p.count[BlackKnight] > 0 {
		black[0] = e.knights(Black, maskMobile[Black])
	}
	if p.count[BlackBishop] > 0 {
		black[1] = e.bishops(Black, maskMobile[Black])
	}
	if p.count[BlackRook] > 0 {
		black[2] = e.rooks(Black, maskMobile[Black])
	}

	// Now that we've built all attack bitmasks we can adjust mobility to
	// exclude attacks by enemy's knights, bishops, and rooks and evaluate
	// the queens.
	if p.count[Queen] > 0 {
		maskMobile[White] &= ^(e.attacks[BlackKnight] | e.attacks[BlackBishop] | e.attacks[BlackRook])
		white[3] = e.queens(White, maskMobile[White])
	}
	if p.count[BlackQueen] > 0 {
		maskMobile[Black] &= ^(e.attacks[Knight] | e.attacks[Bishop] | e.attacks[Rook])
		black[3] = e.queens(Black, maskMobile[Black])
	}

	// Update attack bitmasks for both sides.
	e.attacks[White] |= e.attacks[Knight] | e.attacks[Bishop] | e.attacks[Rook] | e.attacks[Queen]
	e.attacks[Black] |= e.attacks[BlackKnight] | e.attacks[BlackBishop] | e.attacks[BlackRook] | e.attacks[BlackQueen]

	// Update cumulative score based on white vs. black delta.
	e.score.add(white[0]).add(white[1]).add(white[2]).add(white[3])
	e.score.subtract(black[0]).subtract(black[1]).subtract(black[2]).subtract(black[3])
}

func (e *Evaluation) knights(color int, maskMobile Bitmask) (score Score) {
	p := e.position
	outposts := p.outposts[knight(color)]

	for outposts != 0 {
		square := outposts.pop()
		attacks := p.attacks(square)

		// Bonus for knight's mobility.
		score.add(mobilityKnight[(attacks & maskMobile).count()])

		// Penalty if knight is attacked by enemy's pawn.
		if maskPawn[color^1][square] & p.outposts[pawn(color^1)] != 0 {
			score.subtract(penaltyPawnThreat[Knight/2])
		}

		// Bonus if knight is behind friendly pawn.
		if RelRow(color, square) < 4 && p.outposts[pawn(color)].isSet(square + eight[color]) {
			score.add(behindPawn)
		}

		// Extra bonus if knight is in the center. Increase the extra
		// bonus if the knight is supported by a pawn and can't be
		// exchanged.
		flip := Flip(color, square)
		if extra := extraKnight[flip]; extra > 0 {
			if p.pawnAttacks(color).isSet(square) {
				if p.count[knight(color^1)] == 0 {
					extra *= 2 // No knights to exchange.
				}
				extra += extra / 2 // Supported by a pawn.
			}
			score.adjust(extra)
		}

		// Track if knight attacks squares around enemy's king.
		e.enemyKingThreat(knight(color), attacks)
	}
	return
}

func (e *Evaluation) bishops(color int, maskMobile Bitmask) (score Score) {
	p := e.position
	outposts := p.outposts[bishop(color)]

	for outposts != 0 {
		square := outposts.pop()
		attacks := p.xrayAttacks(square)

		// Bonus for bishop's mobility
		score.add(mobilityBishop[(attacks & maskMobile).count()])

		// Penalty for light/dark square bishop and matching pawns.
		if count := (Same(square) & p.outposts[pawn(color)]).count(); count > 0 {
			score.subtract(bishopPawn)
		}

		// Penalty if bishop is attacked by enemy's pawn.
		if maskPawn[color^1][square] & p.outposts[pawn(color^1)] != 0 {
			score.subtract(penaltyPawnThreat[Bishop/2])
		}

		// Bonus if bishop is behind friendly pawn.
		if RelRow(color, square) < 4 && p.outposts[pawn(color)].isSet(square + eight[color]) {
			score.add(behindPawn)
		}

		// Middle game penalty for boxed bishop.
		if e.material.phase > 160 {
			if color == White {
				if (square == C1 && p.pieces[D2].isPawn() && p.pieces[D3] != 0) ||
				   (square == F1 && p.pieces[E2].isPawn() && p.pieces[E3] != 0) {
					score.midgame -= bishopBoxed.midgame
				}
			} else {
				if (square == C8 && p.pieces[D7].isPawn() && p.pieces[D6] != 0) ||
				   (square == F8 && p.pieces[E7].isPawn() && p.pieces[E6] != 0) {
					score.midgame -= bishopBoxed.midgame
				}
			}
		}

		// Extra bonus if bishop is in the center. Increase the extra
		// bonus if the bishop is supported by a pawn and can't be
		// exchanged.
		flip := Flip(color, square)
		if extra := extraBishop[flip]; extra > 0 {
			if p.pawnAttacks(color).isSet(square) {
				if p.count[bishop(color^1)] == 0 {
					extra *= 2 // No bishops to exchange.
				}
				extra += extra / 2 // Supported by a pawn.
			}
			score.adjust(extra)
		}

		// Track if bishop attacks squares around enemy's king.
		e.enemyKingThreat(bishop(color), attacks)
	}
	return
}

func (e *Evaluation) rooks(color int, maskMobile Bitmask) (score Score) {
	p := e.position
	hisPawns := p.outposts[pawn(color)]
	herPawns := p.outposts[pawn(color^1)]
	outposts := p.outposts[rook(color)]

	// Bonus if rook is on 7th rank and enemy's king trapped on 8th.
	if count := (outposts & mask7th[color]).count(); count > 0 && p.outposts[king(color^1)] & mask8th[color] != 0 {
		score.add(rookOn7th.times(count))
	}
	for outposts != 0 {
		square := outposts.pop()
		attacks := p.xrayAttacks(square)

		// Bonus for rook's mobility
		mobility := (attacks & maskMobile).count()
		score.add(mobilityRook[mobility])

		// Penalty if rook is attacked by enemy's pawn.
		if maskPawn[color^1][square] & p.outposts[pawn(color^1)] != 0 {
			score.subtract(penaltyPawnThreat[Rook/2])
		}

		// Bonus if rook is attacking enemy's pawns.
		if count := (attacks & p.outposts[pawn(color^1)]).count(); count > 0 {
			score.add(rookOnPawn.times(count))
		}

		// Bonuses if rook is on open or semi-open file.
		column := Col(square)
		isFileAjar := (hisPawns & maskFile[column] == 0)
		if isFileAjar {
			if herPawns & maskFile[column] == 0 {
				score.add(rookOnOpen)
			} else {
				score.add(rookOnSemiOpen)
			}
		}

		// Middle game penalty if a rook is boxed. Extra penalty if castle
		// rights have been lost.
		if mobility <= 3 || !isFileAjar {
			kingSquare := p.king[color]
			kingColumn := Col(kingSquare)

			// Queenside box: king on D/C/B vs. rook on A/B/C files. Double the
			// the penalty since no castle is possible.
			if column < kingColumn && rookBoxA[color].isSet(square) && kingBoxA[color].isSet(kingSquare) {
				score.midgame -= rookBoxed.midgame * 2
			}

			// Kingside box: king on E/F/G vs. rook on H/G/F files.
			if column > kingColumn && rookBoxH[color].isSet(square) && kingBoxH[color].isSet(kingSquare) {
				score.midgame -= rookBoxed.midgame
				if p.castles & castleKingside[color] == 0 {
					score.midgame -= rookBoxed.midgame
				}
			}
		}

		// Track if rook attacks squares around enemy's king.
		e.enemyKingThreat(rook(color), attacks)
	}
	return
}

func (e *Evaluation) queens(color int, maskMobile Bitmask) (score Score) {
	p := e.position
	outposts := p.outposts[queen(color)]

	// Bonus if queen is on 7th rank and enemy's king trapped on 8th.
	if count := (outposts & mask7th[color]).count(); count > 0 && p.outposts[king(color^1)] & mask8th[color] != 0 {
		score.add(queenOn7th.times(count))
	}
	for outposts != 0 {
		square := outposts.pop()
		attacks := p.attacks(square)

		// Bonus for queen's mobility.
		score.add(mobilityQueen[Min(15, (attacks & maskMobile).count())])

		// Penalty if queen is attacked by enemy's pawn.
		if maskPawn[color^1][square] & p.outposts[pawn(color^1)] != 0 {
			score.subtract(penaltyPawnThreat[Queen/2])
		}

		// Bonus if queen is out and attacking enemy's pawns.
		if count := (attacks & p.outposts[pawn(color^1)]).count(); count > 0 && RelRow(color, square) > 3 {
			score.add(queenOnPawn.times(count))
		}

		// Track if queen attacks squares around enemy's king.
		e.enemyKingThreat(queen(color), attacks)
	}
	return
}

func (e *Evaluation) enemyKingThreat(piece Piece, attacks Bitmask) {
	color := piece.color() ^ 1

	if attacks & e.safety[color].fort != 0 {
		e.safety[color].attackers++
		e.safety[color].threats += bonusKingThreat[piece.kind()/2]
		if bits := attacks & e.attacks[king(color)]; bits != 0 {
			e.safety[color].attacks += bits.count()
		}
	}

	// Update attack bitmask for the given piece.
	e.attacks[piece] |= attacks
}
<commit_msg>Compute king safety bits only when king safety evaluation is pending<commit_after>// Copyright (c) 2013-2014 by Michael Dvorkin. All Rights Reserved.
// Use of this source code is governed by a MIT-style license that can
// be found in the LICENSE file.

package donna

func (e *Evaluation) analyzePieces() {
	p := e.position
	var white, black [4]Score

	if Settings.Trace {
		defer func() {
			var his, her Score
			e.checkpoint(`+Pieces`,  Total{*his.add(white[0]).add(white[1]).add(white[2]).add(white[3]),
				*her.add(black[0]).add(black[1]).add(black[2]).add(black[3])})
			e.checkpoint(`-Knights`, Total{white[0], black[0]})
			e.checkpoint(`-Bishops`, Total{white[1], black[1]})
			e.checkpoint(`-Rooks`,   Total{white[2], black[2]})
			e.checkpoint(`-Queens`,  Total{white[3], black[3]})
		}()
	}

	// Mobility mask for both sides excludes a) squares attacked by enemy's
	// pawns and b) squares occupied by own pawns and king.
	maskMobile := [2]Bitmask{
		^(e.attacks[BlackPawn] | p.outposts[Pawn] | p.outposts[King]),
		^(e.attacks[Pawn] | p.outposts[BlackPawn] | p.outposts[BlackKing]),
	}

	// Initialize king fort bitmasks only when we need them.
	whiteSafety := (e.material.flags & whiteKingSafety != 0)
	blackSafety := (e.material.flags & blackKingSafety != 0)
	if whiteSafety {
		e.safety[White].fort = e.setupFort(White)
	}
	if blackSafety {
		e.safety[Black].fort = e.setupFort(Black)
	}

	// Evaluate white pieces except queen.
	if p.count[Knight] > 0 {
		white[0] = e.knights(White, maskMobile[White], blackSafety)
	}
	if p.count[Bishop] > 0 {
		white[1] = e.bishops(White, maskMobile[White], blackSafety)
	}
	if p.count[Rook] > 0 {
		white[2] = e.rooks(White, maskMobile[White], blackSafety)
	}

	// Evaluate black pieces except queen.
	if p.count[BlackKnight] > 0 {
		black[0] = e.knights(Black, maskMobile[Black], whiteSafety)
	}
	if p.count[BlackBishop] > 0 {
		black[1] = e.bishops(Black, maskMobile[Black], whiteSafety)
	}
	if p.count[BlackRook] > 0 {
		black[2] = e.rooks(Black, maskMobile[Black], whiteSafety)
	}

	// Now that we've built all attack bitmasks we can adjust mobility to
	// exclude attacks by enemy's knights, bishops, and rooks and evaluate
	// the queens.
	if p.count[Queen] > 0 {
		maskMobile[White] &= ^(e.attacks[BlackKnight] | e.attacks[BlackBishop] | e.attacks[BlackRook])
		white[3] = e.queens(White, maskMobile[White], blackSafety)
	}
	if p.count[BlackQueen] > 0 {
		maskMobile[Black] &= ^(e.attacks[Knight] | e.attacks[Bishop] | e.attacks[Rook])
		black[3] = e.queens(Black, maskMobile[Black], whiteSafety)
	}

	// Update attack bitmasks for both sides.
	e.attacks[White] |= e.attacks[Knight] | e.attacks[Bishop] | e.attacks[Rook] | e.attacks[Queen]
	e.attacks[Black] |= e.attacks[BlackKnight] | e.attacks[BlackBishop] | e.attacks[BlackRook] | e.attacks[BlackQueen]

	// Update cumulative score based on white vs. black delta.
	e.score.add(white[0]).add(white[1]).add(white[2]).add(white[3])
	e.score.subtract(black[0]).subtract(black[1]).subtract(black[2]).subtract(black[3])
}

func (e *Evaluation) knights(color int, maskMobile Bitmask, kingSafety bool) (score Score) {
	p := e.position
	outposts := p.outposts[knight(color)]

	for outposts != 0 {
		square := outposts.pop()
		attacks := p.attacks(square)

		// Bonus for knight's mobility.
		score.add(mobilityKnight[(attacks & maskMobile).count()])

		// Penalty if knight is attacked by enemy's pawn.
		if maskPawn[color^1][square] & p.outposts[pawn(color^1)] != 0 {
			score.subtract(penaltyPawnThreat[Knight/2])
		}

		// Bonus if knight is behind friendly pawn.
		if RelRow(color, square) < 4 && p.outposts[pawn(color)].isSet(square + eight[color]) {
			score.add(behindPawn)
		}

		// Extra bonus if knight is in the center. Increase the extra
		// bonus if the knight is supported by a pawn and can't be
		// exchanged.
		flip := Flip(color, square)
		if extra := extraKnight[flip]; extra > 0 {
			if p.pawnAttacks(color).isSet(square) {
				if p.count[knight(color^1)] == 0 {
					extra *= 2 // No knights to exchange.
				}
				extra += extra / 2 // Supported by a pawn.
			}
			score.adjust(extra)
		}

		// Track if knight attacks squares around enemy's king.
		if kingSafety {
			e.enemyKingThreat(knight(color), attacks)
		}

		// Update attack bitmask for the knight.
		e.attacks[knight(color)] |= attacks
	}
	return
}

func (e *Evaluation) bishops(color int, maskMobile Bitmask, kingSafety bool) (score Score) {
	p := e.position
	outposts := p.outposts[bishop(color)]

	for outposts != 0 {
		square := outposts.pop()
		attacks := p.xrayAttacks(square)

		// Bonus for bishop's mobility
		score.add(mobilityBishop[(attacks & maskMobile).count()])

		// Penalty for light/dark square bishop and matching pawns.
		if count := (Same(square) & p.outposts[pawn(color)]).count(); count > 0 {
			score.subtract(bishopPawn)
		}

		// Penalty if bishop is attacked by enemy's pawn.
		if maskPawn[color^1][square] & p.outposts[pawn(color^1)] != 0 {
			score.subtract(penaltyPawnThreat[Bishop/2])
		}

		// Bonus if bishop is behind friendly pawn.
		if RelRow(color, square) < 4 && p.outposts[pawn(color)].isSet(square + eight[color]) {
			score.add(behindPawn)
		}

		// Middle game penalty for boxed bishop.
		if e.material.phase > 160 {
			if color == White {
				if (square == C1 && p.pieces[D2].isPawn() && p.pieces[D3] != 0) ||
				   (square == F1 && p.pieces[E2].isPawn() && p.pieces[E3] != 0) {
					score.midgame -= bishopBoxed.midgame
				}
			} else {
				if (square == C8 && p.pieces[D7].isPawn() && p.pieces[D6] != 0) ||
				   (square == F8 && p.pieces[E7].isPawn() && p.pieces[E6] != 0) {
					score.midgame -= bishopBoxed.midgame
				}
			}
		}

		// Extra bonus if bishop is in the center. Increase the extra
		// bonus if the bishop is supported by a pawn and can't be
		// exchanged.
		flip := Flip(color, square)
		if extra := extraBishop[flip]; extra > 0 {
			if p.pawnAttacks(color).isSet(square) {
				if p.count[bishop(color^1)] == 0 {
					extra *= 2 // No bishops to exchange.
				}
				extra += extra / 2 // Supported by a pawn.
			}
			score.adjust(extra)
		}

		// Track if bishop attacks squares around enemy's king.
		if kingSafety {
			e.enemyKingThreat(bishop(color), attacks)
		}

		// Update attack bitmask for the bishop.
		e.attacks[bishop(color)] |= attacks
	}
	return
}

func (e *Evaluation) rooks(color int, maskMobile Bitmask, kingSafety bool) (score Score) {
	p := e.position
	hisPawns := p.outposts[pawn(color)]
	herPawns := p.outposts[pawn(color^1)]
	outposts := p.outposts[rook(color)]

	// Bonus if rook is on 7th rank and enemy's king trapped on 8th.
	if count := (outposts & mask7th[color]).count(); count > 0 && p.outposts[king(color^1)] & mask8th[color] != 0 {
		score.add(rookOn7th.times(count))
	}
	for outposts != 0 {
		square := outposts.pop()
		attacks := p.xrayAttacks(square)

		// Bonus for rook's mobility
		mobility := (attacks & maskMobile).count()
		score.add(mobilityRook[mobility])

		// Penalty if rook is attacked by enemy's pawn.
		if maskPawn[color^1][square] & p.outposts[pawn(color^1)] != 0 {
			score.subtract(penaltyPawnThreat[Rook/2])
		}

		// Bonus if rook is attacking enemy's pawns.
		if count := (attacks & p.outposts[pawn(color^1)]).count(); count > 0 {
			score.add(rookOnPawn.times(count))
		}

		// Bonuses if rook is on open or semi-open file.
		column := Col(square)
		isFileAjar := (hisPawns & maskFile[column] == 0)
		if isFileAjar {
			if herPawns & maskFile[column] == 0 {
				score.add(rookOnOpen)
			} else {
				score.add(rookOnSemiOpen)
			}
		}

		// Middle game penalty if a rook is boxed. Extra penalty if castle
		// rights have been lost.
		if mobility <= 3 || !isFileAjar {
			kingSquare := p.king[color]
			kingColumn := Col(kingSquare)

			// Queenside box: king on D/C/B vs. rook on A/B/C files. Double the
			// the penalty since no castle is possible.
			if column < kingColumn && rookBoxA[color].isSet(square) && kingBoxA[color].isSet(kingSquare) {
				score.midgame -= rookBoxed.midgame * 2
			}

			// Kingside box: king on E/F/G vs. rook on H/G/F files.
			if column > kingColumn && rookBoxH[color].isSet(square) && kingBoxH[color].isSet(kingSquare) {
				score.midgame -= rookBoxed.midgame
				if p.castles & castleKingside[color] == 0 {
					score.midgame -= rookBoxed.midgame
				}
			}
		}

		// Track if rook attacks squares around enemy's king.
		if kingSafety {
			e.enemyKingThreat(rook(color), attacks)
		}

		// Update attack bitmask for the rook.
		e.attacks[rook(color)] |= attacks
	}
	return
}

func (e *Evaluation) queens(color int, maskMobile Bitmask, kingSafety bool) (score Score) {
	p := e.position
	outposts := p.outposts[queen(color)]

	// Bonus if queen is on 7th rank and enemy's king trapped on 8th.
	if count := (outposts & mask7th[color]).count(); count > 0 && p.outposts[king(color^1)] & mask8th[color] != 0 {
		score.add(queenOn7th.times(count))
	}
	for outposts != 0 {
		square := outposts.pop()
		attacks := p.attacks(square)

		// Bonus for queen's mobility.
		score.add(mobilityQueen[Min(15, (attacks & maskMobile).count())])

		// Penalty if queen is attacked by enemy's pawn.
		if maskPawn[color^1][square] & p.outposts[pawn(color^1)] != 0 {
			score.subtract(penaltyPawnThreat[Queen/2])
		}

		// Bonus if queen is out and attacking enemy's pawns.
		if count := (attacks & p.outposts[pawn(color^1)]).count(); count > 0 && RelRow(color, square) > 3 {
			score.add(queenOnPawn.times(count))
		}

		// Track if queen attacks squares around enemy's king.
		if kingSafety {
			e.enemyKingThreat(queen(color), attacks)
		}

		// Update attack bitmask for the queen.
		e.attacks[queen(color)] |= attacks
	}
	return
}

func (e *Evaluation) enemyKingThreat(piece Piece, attacks Bitmask) {
	color := piece.color() ^ 1

	if attacks & e.safety[color].fort != 0 {
		e.safety[color].attackers++
		e.safety[color].threats += bonusKingThreat[piece.kind()/2]
		if bits := attacks & e.attacks[king(color)]; bits != 0 {
			e.safety[color].attacks += bits.count()
		}
	}
}
<|endoftext|>