<commit_before>package factorlib

type bitVec []uint

// wordsize == {32,64} when uint == uint{32,64}
const wordsize = 32 << (^uint(0) >> 63)

func newBitVec(n uint) bitVec {
	return make([]uint, (n+wordsize-1)/wordsize)
}

func (b bitVec) getBit(i uint) bool {
	return b[i/wordsize]>>uint(i%wordsize) & 1 != 0
}

func (b bitVec) setBit(i uint) {
	b[i/wordsize] |= uint(1) << uint(i%wordsize)
}

func (b bitVec) toggleBit(i uint) {
	b[i/wordsize] ^= uint(1) << uint(i%wordsize)
}

func (b bitVec) empty() bool {
	for _, x := range b {
		if x != 0 {
			return false
		}
	}
	return true
}

func (b bitVec) firstBit() uint {
	for i, x := range b {
		if x != 0 {
			for j := uint(0); j < wordsize; j++ {
				if x & 1 != 0 {
					return uint(i)*wordsize + j
				}
				x >>= 1
			}
		}
	}
	panic("firstBit not defined on empty bitVec")
}

// b ^= c.  b and c must be constructed using the same length.
func (b bitVec) xor(c bitVec) {
	for i, x := range c {
		b[i] ^= x
	}
}
<commit_msg>Remove unneeded uint casts.<commit_after>package factorlib

type bitVec []uint

// wordsize == {32,64} when uint == uint{32,64}
const wordsize = 32 << (^uint(0) >> 63)

func newBitVec(n uint) bitVec {
	return make([]uint, (n+wordsize-1)/wordsize)
}

func (b bitVec) getBit(i uint) bool {
	return b[i/wordsize]>>(i%wordsize) & 1 != 0
}

func (b bitVec) setBit(i uint) {
	b[i/wordsize] |= uint(1) << (i%wordsize)
}

func (b bitVec) toggleBit(i uint) {
	b[i/wordsize] ^= uint(1) << (i%wordsize)
}

func (b bitVec) empty() bool {
	for _, x := range b {
		if x != 0 {
			return false
		}
	}
	return true
}

func (b bitVec) firstBit() uint {
	for i, x := range b {
		if x != 0 {
			for j := uint(0); j < wordsize; j++ {
				if x & 1 != 0 {
					return uint(i)*wordsize + j
				}
				x >>= 1
			}
		}
	}
	panic("firstBit not defined on empty bitVec")
}

// b ^= c.  b and c must be constructed using the same length.
func (b bitVec) xor(c bitVec) {
	for i, x := range c {
		b[i] ^= x
	}
}
<|endoftext|>


<commit_before>/*

Series:

	╔══════Series List═════╗
	║ ┌───────────────────┐║
	║ │     Term List     │║
	║ ├───────────────────┤║
	║ │    Series Data    │║
	║ ├───────────────────┤║
	║ │      Trailer      │║
	║ └───────────────────┘║
	╚══════════════════════╝

	╔══════════Term List═══════════╗
	║ ┌──────────────────────────┐ ║
	║ │   Term Count <uint32>    │ ║
	║ └──────────────────────────┘ ║
	║ ┏━━━━━━━━━━━━━━━━━━━━━━━━━━┓ ║
	║ ┃ ┌──────────────────────┐ ┃ ║
	║ ┃ │  len(Term) <varint>  │ ┃ ║
	║ ┃ ├──────────────────────┤ ┃ ║
	║ ┃ │    Term <byte...>    │ ┃ ║
	║ ┃ └──────────────────────┘ ┃ ║
	║ ┗━━━━━━━━━━━━━━━━━━━━━━━━━━┛ ║
	║ ┏━━━━━━━━━━━━━━━━━━━━━━━━━━┓ ║
	║ ┃ ┌──────────────────────┐ ┃ ║
	║ ┃ │  len(Term) <varint>  │ ┃ ║
	║ ┃ ├──────────────────────┤ ┃ ║
	║ ┃ │    Term <byte...>    │ ┃ ║
	║ ┃ └──────────────────────┘ ┃ ║
	║ ┗━━━━━━━━━━━━━━━━━━━━━━━━━━┛ ║
	╚══════════════════════════════╝

	╔═════════Series Data══════════╗
	║ ┌──────────────────────────┐ ║
	║ │  Series Count <uint32>   │ ║
	║ └──────────────────────────┘ ║
	║ ┏━━━━━━━━━━━━━━━━━━━━━━━━━━┓ ║
	║ ┃ ┌──────────────────────┐ ┃ ║
	║ ┃ │     Flag <uint8>     │ ┃ ║
	║ ┃ ├──────────────────────┤ ┃ ║
	║ ┃ │ len(Series) <varint> │ ┃ ║
	║ ┃ ├──────────────────────┤ ┃ ║
	║ ┃ │   Series <byte...>   │ ┃ ║
	║ ┃ └──────────────────────┘ ┃ ║
	║ ┗━━━━━━━━━━━━━━━━━━━━━━━━━━┛ ║
	║             ...              ║
	╚══════════════════════════════╝

	╔════════════Trailer══════════════╗
	║ ┌─────────────────────────────┐ ║
	║ │  Term List Offset <uint64>  │ ║
	║ ├─────────────────────────────┤ ║
	║ │   Term List Size <uint64>   │ ║
	║ ├─────────────────────────────┤ ║
	║ │ Series Data Offset <uint64> │ ║
	║ ├─────────────────────────────┤ ║
	║ │  Series Data Pos <uint64>   │ ║
	║ ├─────────────────────────────┤ ║
	║ │    Sketch Offset <uint64>   │ ║
	║ ├─────────────────────────────┤ ║
	║ │     Sketch Size <uint64>    │ ║
	║ ├─────────────────────────────┤ ║
	║ │ Tomb Sketch Offset <uint64> │ ║
	║ ├─────────────────────────────┤ ║
	║ │  Tomb Sketch Size <uint64>  │ ║
	║ └─────────────────────────────┘ ║
	╚═════════════════════════════════╝

Tag Block:

	╔═══════Tag Block════════╗
	║┌──────────────────────┐║
	║│   Tag Values Block   │║
	║├──────────────────────┤║
	║│         ...          │║
	║├──────────────────────┤║
	║│    Tag Keys Block    │║
	║├──────────────────────┤║
	║│       Trailer        │║
	║└──────────────────────┘║
	╚════════════════════════╝

	╔═══════Tag Values Block═══════╗
	║                              ║
	║ ┏━━━━━━━━Value List━━━━━━━━┓ ║
	║ ┃                          ┃ ║
	║ ┃┏━━━━━━━━━Value━━━━━━━━━━┓┃ ║
	║ ┃┃┌──────────────────────┐┃┃ ║
	║ ┃┃│     Flag <uint8>     │┃┃ ║
	║ ┃┃├──────────────────────┤┃┃ ║
	║ ┃┃│ len(Value) <varint>  │┃┃ ║
	║ ┃┃├──────────────────────┤┃┃ ║
	║ ┃┃│   Value <byte...>    │┃┃ ║
	║ ┃┃├──────────────────────┤┃┃ ║
	║ ┃┃│ len(Series) <varint> │┃┃ ║
	║ ┃┃├──────────────────────┤┃┃ ║
	║ ┃┃│SeriesIDs <uint32...> │┃┃ ║
	║ ┃┃└──────────────────────┘┃┃ ║
	║ ┃┗━━━━━━━━━━━━━━━━━━━━━━━━┛┃ ║
	║ ┃           ...            ┃ ║
	║ ┗━━━━━━━━━━━━━━━━━━━━━━━━━━┛ ║
	║ ┏━━━━━━━━Hash Index━━━━━━━━┓ ║
	║ ┃ ┌──────────────────────┐ ┃ ║
	║ ┃ │ len(Values) <uint32> │ ┃ ║
	║ ┃ ├──────────────────────┤ ┃ ║
	║ ┃ │Value Offset <uint64> │ ┃ ║
	║ ┃ ├──────────────────────┤ ┃ ║
	║ ┃ │         ...          │ ┃ ║
	║ ┃ └──────────────────────┘ ┃ ║
	║ ┗━━━━━━━━━━━━━━━━━━━━━━━━━━┛ ║
	╚══════════════════════════════╝

	╔════════Tag Key Block═════════╗
	║                              ║
	║ ┏━━━━━━━━━Key List━━━━━━━━━┓ ║
	║ ┃                          ┃ ║
	║ ┃┏━━━━━━━━━━Key━━━━━━━━━━━┓┃ ║
	║ ┃┃┌──────────────────────┐┃┃ ║
	║ ┃┃│     Flag <uint8>     │┃┃ ║
	║ ┃┃├──────────────────────┤┃┃ ║
	║ ┃┃│Value Offset <uint64> │┃┃ ║
	║ ┃┃├──────────────────────┤┃┃ ║
	║ ┃┃│  len(Key) <varint>   │┃┃ ║
	║ ┃┃├──────────────────────┤┃┃ ║
	║ ┃┃│    Key <byte...>     │┃┃ ║
	║ ┃┃└──────────────────────┘┃┃ ║
	║ ┃┗━━━━━━━━━━━━━━━━━━━━━━━━┛┃ ║
	║ ┃           ...            ┃ ║
	║ ┗━━━━━━━━━━━━━━━━━━━━━━━━━━┛ ║
	║ ┏━━━━━━━━Hash Index━━━━━━━━┓ ║
	║ ┃ ┌──────────────────────┐ ┃ ║
	║ ┃ │  len(Keys) <uint32>  │ ┃ ║
	║ ┃ ├──────────────────────┤ ┃ ║
	║ ┃ │ Key Offset <uint64>  │ ┃ ║
	║ ┃ ├──────────────────────┤ ┃ ║
	║ ┃ │         ...          │ ┃ ║
	║ ┃ └──────────────────────┘ ┃ ║
	║ ┗━━━━━━━━━━━━━━━━━━━━━━━━━━┛ ║
	╚══════════════════════════════╝

	╔════════════Trailer══════════════╗
	║ ┌─────────────────────────────┐ ║
	║ │  Hash Index Offset <uint64> │ ║
	║ ├─────────────────────────────┤ ║
	║ │    Tag Set Size <uint64>    │ ║
	║ ├─────────────────────────────┤ ║
	║ │   Tag Set Version <uint16>  │ ║
	║ └─────────────────────────────┘ ║
	╚═════════════════════════════════╝

Measurements

	╔══════════Measurements Block═══════════╗
	║                                       ║
	║ ┏━━━━━━━━━Measurement List━━━━━━━━━━┓ ║
	║ ┃                                   ┃ ║
	║ ┃┏━━━━━━━━━━Measurement━━━━━━━━━━━┓ ┃ ║
	║ ┃┃┌─────────────────────────────┐ ┃ ┃ ║
	║ ┃┃│        Flag <uint8>         │ ┃ ┃ ║
	║ ┃┃├─────────────────────────────┤ ┃ ┃ ║
	║ ┃┃│  Tag Block Offset <uint64>  │ ┃ ┃ ║
	║ ┃┃├─────────────────────────────┤ ┃ ┃ ║
	║ ┃┃│     len(Name) <varint>      │ ┃ ┃ ║
	║ ┃┃├─────────────────────────────┤ ┃ ┃ ║
	║ ┃┃│       Name <byte...>        │ ┃ ┃ ║
	║ ┃┃├─────────────────────────────┤ ┃ ┃ ║
	║ ┃┃│    len(Series) <uint32>     │ ┃ ┃ ║
	║ ┃┃├─────────────────────────────┤ ┃ ┃ ║
	║ ┃┃│    SeriesIDs <uint32...>    │ ┃ ┃ ║
	║ ┃┃└─────────────────────────────┘ ┃ ┃ ║
	║ ┃┗━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━┛ ┃ ║
	║ ┃                ...                ┃ ║
	║ ┗━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━┛ ║
	║ ┏━━━━━━━━━━━━Hash Index━━━━━━━━━━━━━┓ ║
	║ ┃ ┌───────────────────────────────┐ ┃ ║
	║ ┃ │  len(Measurements) <uint32>   │ ┃ ║
	║ ┃ ├───────────────────────────────┤ ┃ ║
	║ ┃ │  Measurement Offset <uint64>  │ ┃ ║
	║ ┃ ├───────────────────────────────┤ ┃ ║
	║ ┃ │              ...              │ ┃ ║
	║ ┃ └───────────────────────────────┘ ┃ ║
	║ ┗━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━┛ ║
	║ ┏━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━┓ ║
	║ ┃              Sketch               ┃ ║
	║ ┗━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━┛ ║
	║ ┏━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━┓ ║
	║ ┃          Tombstone Sketch         ┃ ║
	║ ┗━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━┛ ║
	║ ┏━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━┓ ║
	║ ┃              Trailer              ┃ ║
	║ ┗━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━┛ ║
	╚═══════════════════════════════════════╝

	╔════════════Trailer══════════════╗
	║ ┌─────────────────────────────┐ ║
	║ │     Block Offset <uint64>   │ ║
	║ ├─────────────────────────────┤ ║
	║ │     Block Size <uint64>     │ ║
	║ ├─────────────────────────────┤ ║
	║ │  Hash Index Offset <uint64> │ ║
	║ ├─────────────────────────────┤ ║
	║ │   Hash Index Size <uint64>  │ ║
	║ ├─────────────────────────────┤ ║
	║ │    Sketch Offset <uint64>   │ ║
	║ ├─────────────────────────────┤ ║
	║ │     Sketch Size <uint64>    │ ║
	║ ├─────────────────────────────┤ ║
	║ │ Tomb Sketch Offset <uint64> │ ║
	║ ├─────────────────────────────┤ ║
	║ │  Tomb Sketch Size <uint64>  │ ║
	║ ├─────────────────────────────┤ ║
	║ │    Block Version <uint16>   │ ║
	║ └─────────────────────────────┘ ║
	╚═════════════════════════════════╝

WAL

	╔═════════════WAL══════════════╗
	║                              ║
	║ ┏━━━━━━━━━━Entry━━━━━━━━━━━┓ ║
	║ ┃ ┌──────────────────────┐ ┃ ║
	║ ┃ │     Flag <uint8>     │ ┃ ║
	║ ┃ ├──────────────────────┤ ┃ ║
	║ ┃ │  len(Name) <varint>  │ ┃ ║
	║ ┃ ├──────────────────────┤ ┃ ║
	║ ┃ │    Name <byte...>    │ ┃ ║
	║ ┃ ├──────────────────────┤ ┃ ║
	║ ┃ │  len(Tags) <varint>  │ ┃ ║
	║ ┃ ├──────────────────────┤ ┃ ║
	║ ┃ │  len(Key0) <varint>  │ ┃ ║
	║ ┃ ├──────────────────────┤ ┃ ║
	║ ┃ │    Key0 <byte...>    │ ┃ ║
	║ ┃ ├──────────────────────┤ ┃ ║
	║ ┃ │ len(Value0) <varint> │ ┃ ║
	║ ┃ ├──────────────────────┤ ┃ ║
	║ ┃ │   Value0 <byte...>   │ ┃ ║
	║ ┃ ├──────────────────────┤ ┃ ║
	║ ┃ │         ...          │ ┃ ║
	║ ┃ ├──────────────────────┤ ┃ ║
	║ ┃ │  Checksum <uint32>   │ ┃ ║
	║ ┃ └──────────────────────┘ ┃ ║
	║ ┗━━━━━━━━━━━━━━━━━━━━━━━━━━┛ ║
	║             ...              ║
	╚══════════════════════════════╝

*/
package tsi1
<commit_msg>Add TSI documentation.<commit_after>/*

Package tsi1 provides a memory-mapped index implementation that supports
high cardinality series.

Overview

The top-level object in tsi1 is the Index. It is the primary access point from
the rest of the system. The Index is composed of LogFile and IndexFile objects.

Log files are small write-ahead log files that record new series immediately
in the order that they are received. The data within the file is indexed
in-memory so it can be quickly accessed. When the system is restarted, this log
file is replayed and the in-memory representation is rebuilt.

Index files also contain series information, however, they are highly indexed
so that reads can be performed quickly. Index files are built through a process
called compaction where a log file or multiple index files are merged together.

Operations

The index can perform many tasks related to series, measurement, & tag data.
All data is inserted by adding a series to the index. When adding a series,
the measurement, tag keys, and tag values are all extracted and indexed
separately.

Once a series has been added, it can be removed in several ways. First, the
individual series can be removed. Second, it can be removed as part of a bulk
operation by deleting the entire measurement.

The query engine needs to be able to look up series in a variety of ways such
as by measurement name, by tag value, or by using regular expressions. The
index provides an API to iterate over subsets of series and perform set
operations such as unions and intersections.

Log File Layout

The write-ahead file that series initially are inserted into simply appends
all new operations sequentially. It is simply composed of a series of log
entries. An entry contains a flag to specify the operation type, the measurement
name, the tag set, and a checksum.

	┏━━━━━━━━━LogEntry━━━━━━━━━┓
	┃ ┌──────────────────────┐ ┃
	┃ │         Flag         │ ┃
	┃ ├──────────────────────┤ ┃
	┃ │     Measurement      │ ┃
	┃ ├──────────────────────┤ ┃
	┃ │      Key/Value       │ ┃
	┃ ├──────────────────────┤ ┃
	┃ │      Key/Value       │ ┃
	┃ ├──────────────────────┤ ┃
	┃ │      Key/Value       │ ┃
	┃ ├──────────────────────┤ ┃
	┃ │       Checksum       │ ┃
	┃ └──────────────────────┘ ┃
	┗━━━━━━━━━━━━━━━━━━━━━━━━━━┛

When the log file is replayed, if the checksum is incorrect or the entry is
incomplete (because of a partially failed write) then the log is truncated.

Index File Layout

The index file is composed of 3 main block types: one series block, one or more
tag blocks, and one measurement block. At the end of the index file is a
trailer that records metadata such as the offsets to these blocks.

Series Block Layout

The series block stores raw series keys in sorted order. It also provides hash
indexes so that series can be looked up quickly. Hash indexes are inserted
periodically so that memory size is limited at write time. Once all the series
and hash indexes have been written then a list of index entries are written
so that hash indexes can be looked up via binary search.

The end of the block contains two HyperLogLog++ sketches which track the
estimated number of created series and deleted series. After the sketches is
a trailer which contains metadata about the block.

	┏━━━━━━━SeriesBlock━━━━━━━━┓
	┃ ┌──────────────────────┐ ┃
	┃ │      Series Key      │ ┃
	┃ ├──────────────────────┤ ┃
	┃ │      Series Key      │ ┃
	┃ ├──────────────────────┤ ┃
	┃ │      Series Key      │ ┃
	┃ ├──────────────────────┤ ┃
	┃ │                      │ ┃
	┃ │      Hash Index      │ ┃
	┃ │                      │ ┃
	┃ ├──────────────────────┤ ┃
	┃ │      Series Key      │ ┃
	┃ ├──────────────────────┤ ┃
	┃ │      Series Key      │ ┃
	┃ ├──────────────────────┤ ┃
	┃ │      Series Key      │ ┃
	┃ ├──────────────────────┤ ┃
	┃ │                      │ ┃
	┃ │      Hash Index      │ ┃
	┃ │                      │ ┃
	┃ ├──────────────────────┤ ┃
	┃ │    Index Entries     │ ┃
	┃ ├──────────────────────┤ ┃
	┃ │     HLL Sketches     │ ┃
	┃ ├──────────────────────┤ ┃
	┃ │       Trailer        │ ┃
	┃ └──────────────────────┘ ┃
	┗━━━━━━━━━━━━━━━━━━━━━━━━━━┛

Tag Block Layout

After the series block is one or more tag blocks. One of these blocks exists
for every measurement in the index file. The block is structured as a sorted
list of values for each key and then a sorted list of keys. Each of these lists
has their own hash index for fast direct lookups.

	┏━━━━━━━━Tag Block━━━━━━━━━┓
	┃ ┌──────────────────────┐ ┃
	┃ │        Value         │ ┃
	┃ ├──────────────────────┤ ┃
	┃ │        Value         │ ┃
	┃ ├──────────────────────┤ ┃
	┃ │        Value         │ ┃
	┃ ├──────────────────────┤ ┃
	┃ │                      │ ┃
	┃ │      Hash Index      │ ┃
	┃ │                      │ ┃
	┃ └──────────────────────┘ ┃
	┃ ┌──────────────────────┐ ┃
	┃ │        Value         │ ┃
	┃ ├──────────────────────┤ ┃
	┃ │        Value         │ ┃
	┃ ├──────────────────────┤ ┃
	┃ │                      │ ┃
	┃ │      Hash Index      │ ┃
	┃ │                      │ ┃
	┃ └──────────────────────┘ ┃
	┃ ┌──────────────────────┐ ┃
	┃ │         Key          │ ┃
	┃ ├──────────────────────┤ ┃
	┃ │         Key          │ ┃
	┃ ├──────────────────────┤ ┃
	┃ │                      │ ┃
	┃ │      Hash Index      │ ┃
	┃ │                      │ ┃
	┃ └──────────────────────┘ ┃
	┃ ┌──────────────────────┐ ┃
	┃ │       Trailer        │ ┃
	┃ └──────────────────────┘ ┃
	┗━━━━━━━━━━━━━━━━━━━━━━━━━━┛

Each entry for values contains a sorted list of offsets for series keys that use
that value. Series iterators can be built around a single tag key value or
multiple iterators can be merged with set operators such as union or
intersection.

Measurement block

The measurement block stores a sorted list of measurements, their associated
series offsets, and the offset to their tag block. This allows all series for
a measurement to be traversed quickly and it allows fast direct lookups of
measurements and their tags.

This block also contains HyperLogLog++ sketches for new and deleted
measurements.

	┏━━━━Measurement Block━━━━━┓
	┃ ┌──────────────────────┐ ┃
	┃ │     Measurement      │ ┃
	┃ ├──────────────────────┤ ┃
	┃ │     Measurement      │ ┃
	┃ ├──────────────────────┤ ┃
	┃ │     Measurement      │ ┃
	┃ ├──────────────────────┤ ┃
	┃ │                      │ ┃
	┃ │      Hash Index      │ ┃
	┃ │                      │ ┃
	┃ ├──────────────────────┤ ┃
	┃ │     HLL Sketches     │ ┃
	┃ ├──────────────────────┤ ┃
	┃ │       Trailer        │ ┃
	┃ └──────────────────────┘ ┃
	┗━━━━━━━━━━━━━━━━━━━━━━━━━━┛

Manifest file

The index is simply an ordered set of log and index files. These files can be
merged together or rewritten but their order must always be the same. This is
because series, measurements, & tags can be marked as deleted (aka tombstoned)
and this action needs to be tracked in time order.

Whenever the set of active files is changed, a manifest file is written to
track the set. The manifest specifies the ordering of files and, on startup,
all files not in the manifest are removed from the index directory.

Compacting index files

Compaction is the process of taking files and merging them together into a
single file. There are two stages of compaction within TSI.

First, once log files exceed a size threshold then they are compacted into an
index file. This threshold is relatively small because log files must maintain
their index in the heap which TSI tries to avoid. Small log files are also very
quick to convert into an index file so this is done aggressively.

Second, once a contiguous set of index files exceed a factor (e.g. 10x) then
they are all merged together into a single index file and the old files are
discarded. Because all blocks are written in sorted order, the new index file
can be streamed and minimize memory use.

Concurrency

Index files are immutable so they do not require fine grained locks, however,
compactions require that we track which files are in use so they are not
discarded too soon. This is done by using reference counting with file sets.

A file set is simply an ordered list of index files. When the current file set
is obtained from the index, a counter is incremented to track its usage. Once
the user is done with the file set, it is released and the counter is
decremented. A file cannot be removed from the file system until this counter
returns to zero.

Besides the reference counting, there are no other locking mechanisms when
reading or writing index files. Log files, however, do require a lock whenever
they are accessed. This is another reason to minimize log file size.

*/
package tsi1
<|endoftext|>


<commit_before>package bitswap

import (
	"context"
	"errors"

	bsmsg "github.com/ipfs/go-ipfs/exchange/bitswap/message"
	bsnet "github.com/ipfs/go-ipfs/exchange/bitswap/network"
	mockrouting "github.com/ipfs/go-ipfs/routing/mock"
	delay "github.com/ipfs/go-ipfs/thirdparty/delay"

	routing "gx/ipfs/QmPCGUjMRuBcPybZFpjhzpifwPP9wPRoiy5geTQKU4vqWA/go-libp2p-routing"
	ifconnmgr "gx/ipfs/QmSAJm4QdTJ3EGF2cvgNcQyXTEbxqWSW1x4kCVV1aJQUQr/go-libp2p-interface-connmgr"
	logging "gx/ipfs/QmSpJByNKFX1sCsHBEp3R73FL4NF6FnQTEGyNAXHm2GS52/go-log"
	peer "gx/ipfs/QmWNY7dV54ZDYmTA1ykVdwNCqC11mpU4zSUp6XDpLTH9eG/go-libp2p-peer"
	testutil "gx/ipfs/QmeDA8gNhvRTsbrjEieay5wezupJDiky8xvCzDABbsGzmp/go-testutil"
	cid "gx/ipfs/QmeSrf6pzut73u6zLQkRFQ3ygt3k6XFT2kjdYP8Tnkwwyg/go-cid"
)

var log = logging.Logger("bstestnet")

func VirtualNetwork(rs mockrouting.Server, d delay.D) Network {
	return &network{
		clients:       make(map[peer.ID]bsnet.Receiver),
		delay:         d,
		routingserver: rs,
		conns:         make(map[string]struct{}),
	}
}

type network struct {
	clients       map[peer.ID]bsnet.Receiver
	routingserver mockrouting.Server
	delay         delay.D
	conns         map[string]struct{}
}

func (n *network) Adapter(p testutil.Identity) bsnet.BitSwapNetwork {
	client := &networkClient{
		local:   p.ID(),
		network: n,
		routing: n.routingserver.Client(p),
	}
	n.clients[p.ID()] = client
	return client
}

func (n *network) HasPeer(p peer.ID) bool {
	_, found := n.clients[p]
	return found
}

// TODO should this be completely asynchronous?
// TODO what does the network layer do with errors received from services?
func (n *network) SendMessage(
	ctx context.Context,
	from peer.ID,
	to peer.ID,
	message bsmsg.BitSwapMessage) error {

	receiver, ok := n.clients[to]
	if !ok {
		return errors.New("Cannot locate peer on network")
	}

	// nb: terminate the context since the context wouldn't actually be passed
	// over the network in a real scenario

	go n.deliver(receiver, from, message)

	return nil
}

func (n *network) deliver(
	r bsnet.Receiver, from peer.ID, message bsmsg.BitSwapMessage) error {
	if message == nil || from == "" {
		return errors.New("Invalid input")
	}

	n.delay.Wait()

	r.ReceiveMessage(context.TODO(), from, message)
	return nil
}

type networkClient struct {
	local peer.ID
	bsnet.Receiver
	network *network
	routing routing.IpfsRouting
}

func (nc *networkClient) SendMessage(
	ctx context.Context,
	to peer.ID,
	message bsmsg.BitSwapMessage) error {
	return nc.network.SendMessage(ctx, nc.local, to, message)
}

// FindProvidersAsync returns a channel of providers for the given key
func (nc *networkClient) FindProvidersAsync(ctx context.Context, k *cid.Cid, max int) <-chan peer.ID {

	// NB: this function duplicates the PeerInfo -> ID transformation in the
	// bitswap network adapter. Not to worry. This network client will be
	// deprecated once the ipfsnet.Mock is added. The code below is only
	// temporary.

	out := make(chan peer.ID)
	go func() {
		defer close(out)
		providers := nc.routing.FindProvidersAsync(ctx, k, max)
		for info := range providers {
			select {
			case <-ctx.Done():
			case out <- info.ID:
			}
		}
	}()
	return out
}

func (nc *networkClient) ConnectionManager() ifconnmgr.ConnManager {
	return &ifconnmgr.NullConnMgr{}
}

type messagePasser struct {
	net    *network
	target peer.ID
	local  peer.ID
	ctx    context.Context
}

func (mp *messagePasser) SendMsg(ctx context.Context, m bsmsg.BitSwapMessage) error {
	return mp.net.SendMessage(ctx, mp.local, mp.target, m)
}

func (mp *messagePasser) Close() error {
	return nil
}

func (mp *messagePasser) Reset() error {
	return nil
}

func (n *networkClient) NewMessageSender(ctx context.Context, p peer.ID) (bsnet.MessageSender, error) {
	return &messagePasser{
		net:    n.network,
		target: p,
		local:  n.local,
		ctx:    ctx,
	}, nil
}

// Provide provides the key to the network
func (nc *networkClient) Provide(ctx context.Context, k *cid.Cid) error {
	return nc.routing.Provide(ctx, k, true)
}

func (nc *networkClient) SetDelegate(r bsnet.Receiver) {
	nc.Receiver = r
}

func (nc *networkClient) ConnectTo(_ context.Context, p peer.ID) error {
	if !nc.network.HasPeer(p) {
		return errors.New("no such peer in network")
	}
	tag := tagForPeers(nc.local, p)
	if _, ok := nc.network.conns[tag]; ok {
		log.Warning("ALREADY CONNECTED TO PEER (is this a reconnect? test lib needs fixing)")
		return nil
	}
	nc.network.conns[tag] = struct{}{}
	// TODO: add handling for disconnects

	nc.network.clients[p].PeerConnected(nc.local)
	nc.Receiver.PeerConnected(p)
	return nil
}

func tagForPeers(a, b peer.ID) string {
	if a < b {
		return string(a + b)
	}
	return string(b + a)
}
<commit_msg>fix races in testnet<commit_after>package bitswap

import (
	"context"
	"errors"
	"sync"

	bsmsg "github.com/ipfs/go-ipfs/exchange/bitswap/message"
	bsnet "github.com/ipfs/go-ipfs/exchange/bitswap/network"
	mockrouting "github.com/ipfs/go-ipfs/routing/mock"
	delay "github.com/ipfs/go-ipfs/thirdparty/delay"

	routing "gx/ipfs/QmPCGUjMRuBcPybZFpjhzpifwPP9wPRoiy5geTQKU4vqWA/go-libp2p-routing"
	ifconnmgr "gx/ipfs/QmSAJm4QdTJ3EGF2cvgNcQyXTEbxqWSW1x4kCVV1aJQUQr/go-libp2p-interface-connmgr"
	logging "gx/ipfs/QmSpJByNKFX1sCsHBEp3R73FL4NF6FnQTEGyNAXHm2GS52/go-log"
	peer "gx/ipfs/QmWNY7dV54ZDYmTA1ykVdwNCqC11mpU4zSUp6XDpLTH9eG/go-libp2p-peer"
	testutil "gx/ipfs/QmeDA8gNhvRTsbrjEieay5wezupJDiky8xvCzDABbsGzmp/go-testutil"
	cid "gx/ipfs/QmeSrf6pzut73u6zLQkRFQ3ygt3k6XFT2kjdYP8Tnkwwyg/go-cid"
)

var log = logging.Logger("bstestnet")

func VirtualNetwork(rs mockrouting.Server, d delay.D) Network {
	return &network{
		clients:       make(map[peer.ID]bsnet.Receiver),
		delay:         d,
		routingserver: rs,
		conns:         make(map[string]struct{}),
	}
}

type network struct {
	mu            sync.Mutex
	clients       map[peer.ID]bsnet.Receiver
	routingserver mockrouting.Server
	delay         delay.D
	conns         map[string]struct{}
}

func (n *network) Adapter(p testutil.Identity) bsnet.BitSwapNetwork {
	n.mu.Lock()
	defer n.mu.Unlock()

	client := &networkClient{
		local:   p.ID(),
		network: n,
		routing: n.routingserver.Client(p),
	}
	n.clients[p.ID()] = client
	return client
}

func (n *network) HasPeer(p peer.ID) bool {
	n.mu.Lock()
	defer n.mu.Unlock()

	_, found := n.clients[p]
	return found
}

// TODO should this be completely asynchronous?
// TODO what does the network layer do with errors received from services?
func (n *network) SendMessage(
	ctx context.Context,
	from peer.ID,
	to peer.ID,
	message bsmsg.BitSwapMessage) error {

	n.mu.Lock()
	defer n.mu.Unlock()

	receiver, ok := n.clients[to]
	if !ok {
		return errors.New("Cannot locate peer on network")
	}

	// nb: terminate the context since the context wouldn't actually be passed
	// over the network in a real scenario

	go n.deliver(receiver, from, message)

	return nil
}

func (n *network) deliver(
	r bsnet.Receiver, from peer.ID, message bsmsg.BitSwapMessage) error {
	if message == nil || from == "" {
		return errors.New("Invalid input")
	}

	n.delay.Wait()

	r.ReceiveMessage(context.TODO(), from, message)
	return nil
}

type networkClient struct {
	local peer.ID
	bsnet.Receiver
	network *network
	routing routing.IpfsRouting
}

func (nc *networkClient) SendMessage(
	ctx context.Context,
	to peer.ID,
	message bsmsg.BitSwapMessage) error {
	return nc.network.SendMessage(ctx, nc.local, to, message)
}

// FindProvidersAsync returns a channel of providers for the given key
func (nc *networkClient) FindProvidersAsync(ctx context.Context, k *cid.Cid, max int) <-chan peer.ID {

	// NB: this function duplicates the PeerInfo -> ID transformation in the
	// bitswap network adapter. Not to worry. This network client will be
	// deprecated once the ipfsnet.Mock is added. The code below is only
	// temporary.

	out := make(chan peer.ID)
	go func() {
		defer close(out)
		providers := nc.routing.FindProvidersAsync(ctx, k, max)
		for info := range providers {
			select {
			case <-ctx.Done():
			case out <- info.ID:
			}
		}
	}()
	return out
}

func (nc *networkClient) ConnectionManager() ifconnmgr.ConnManager {
	return &ifconnmgr.NullConnMgr{}
}

type messagePasser struct {
	net    *network
	target peer.ID
	local  peer.ID
	ctx    context.Context
}

func (mp *messagePasser) SendMsg(ctx context.Context, m bsmsg.BitSwapMessage) error {
	return mp.net.SendMessage(ctx, mp.local, mp.target, m)
}

func (mp *messagePasser) Close() error {
	return nil
}

func (mp *messagePasser) Reset() error {
	return nil
}

func (n *networkClient) NewMessageSender(ctx context.Context, p peer.ID) (bsnet.MessageSender, error) {
	return &messagePasser{
		net:    n.network,
		target: p,
		local:  n.local,
		ctx:    ctx,
	}, nil
}

// Provide provides the key to the network
func (nc *networkClient) Provide(ctx context.Context, k *cid.Cid) error {
	return nc.routing.Provide(ctx, k, true)
}

func (nc *networkClient) SetDelegate(r bsnet.Receiver) {
	nc.Receiver = r
}

func (nc *networkClient) ConnectTo(_ context.Context, p peer.ID) error {
	nc.network.mu.Lock()

	otherClient, ok := nc.network.clients[p]
	if !ok {
		nc.network.mu.Unlock()
		return errors.New("no such peer in network")
	}

	tag := tagForPeers(nc.local, p)
	if _, ok := nc.network.conns[tag]; ok {
		nc.network.mu.Unlock()
		log.Warning("ALREADY CONNECTED TO PEER (is this a reconnect? test lib needs fixing)")
		return nil
	}
	nc.network.conns[tag] = struct{}{}
	nc.network.mu.Unlock()

	// TODO: add handling for disconnects

	otherClient.PeerConnected(nc.local)
	nc.Receiver.PeerConnected(p)
	return nil
}

func tagForPeers(a, b peer.ID) string {
	if a < b {
		return string(a + b)
	}
	return string(b + a)
}
<|endoftext|>


<commit_before>package bitswap

import (
	"context"
	"errors"
	"sync"

	bsmsg "github.com/ipfs/go-ipfs/exchange/bitswap/message"
	bsnet "github.com/ipfs/go-ipfs/exchange/bitswap/network"
	mockrouting "github.com/ipfs/go-ipfs/routing/mock"
	delay "github.com/ipfs/go-ipfs/thirdparty/delay"

	logging "gx/ipfs/QmRb5jh8z2E8hMGN2tkvs1yHynUanqnZ3UeKwgN1i9P1F8/go-log"
	routing "gx/ipfs/QmTiWLZ6Fo5j4KcTVutZJ5KWRRJrbxzmxA4td8NfEdrPh7/go-libp2p-routing"
	testutil "gx/ipfs/QmVvkK7s5imCiq3JVbL3pGfnhcCnf3LrFJPF4GE2sAoGZf/go-testutil"
	peer "gx/ipfs/QmZoWKhxUmZ2seW4BzX6fJkNR8hh9PsGModr7q171yq2SS/go-libp2p-peer"
	ifconnmgr "gx/ipfs/Qmax8X1Kfahf5WfSB68EWDG3d3qyS3Sqs1v412fjPTfRwx/go-libp2p-interface-connmgr"
	cid "gx/ipfs/QmcZfnkapfECQGcLZaf9B79NRg7cRa9EnZh4LSbkCzwNvY/go-cid"
)

var log = logging.Logger("bstestnet")

func VirtualNetwork(rs mockrouting.Server, d delay.D) Network {
	return &network{
		clients:       make(map[peer.ID]bsnet.Receiver),
		delay:         d,
		routingserver: rs,
		conns:         make(map[string]struct{}),
	}
}

type network struct {
	mu            sync.Mutex
	clients       map[peer.ID]bsnet.Receiver
	routingserver mockrouting.Server
	delay         delay.D
	conns         map[string]struct{}
}

func (n *network) Adapter(p testutil.Identity) bsnet.BitSwapNetwork {
	n.mu.Lock()
	defer n.mu.Unlock()

	client := &networkClient{
		local:   p.ID(),
		network: n,
		routing: n.routingserver.Client(p),
	}
	n.clients[p.ID()] = client
	return client
}

func (n *network) HasPeer(p peer.ID) bool {
	n.mu.Lock()
	defer n.mu.Unlock()

	_, found := n.clients[p]
	return found
}

// TODO should this be completely asynchronous?
// TODO what does the network layer do with errors received from services?
func (n *network) SendMessage(
	ctx context.Context,
	from peer.ID,
	to peer.ID,
	message bsmsg.BitSwapMessage) error {

	n.mu.Lock()
	defer n.mu.Unlock()

	receiver, ok := n.clients[to]
	if !ok {
		return errors.New("Cannot locate peer on network")
	}

	// nb: terminate the context since the context wouldn't actually be passed
	// over the network in a real scenario

	go n.deliver(receiver, from, message)

	return nil
}

func (n *network) deliver(
	r bsnet.Receiver, from peer.ID, message bsmsg.BitSwapMessage) error {
	if message == nil || from == "" {
		return errors.New("Invalid input")
	}

	n.delay.Wait()

	r.ReceiveMessage(context.TODO(), from, message)
	return nil
}

type networkClient struct {
	local peer.ID
	bsnet.Receiver
	network *network
	routing routing.IpfsRouting
}

func (nc *networkClient) SendMessage(
	ctx context.Context,
	to peer.ID,
	message bsmsg.BitSwapMessage) error {
	return nc.network.SendMessage(ctx, nc.local, to, message)
}

// FindProvidersAsync returns a channel of providers for the given key
func (nc *networkClient) FindProvidersAsync(ctx context.Context, k *cid.Cid, max int) <-chan peer.ID {

	// NB: this function duplicates the PeerInfo -> ID transformation in the
	// bitswap network adapter. Not to worry. This network client will be
	// deprecated once the ipfsnet.Mock is added. The code below is only
	// temporary.

	out := make(chan peer.ID)
	go func() {
		defer close(out)
		providers := nc.routing.FindProvidersAsync(ctx, k, max)
		for info := range providers {
			select {
			case <-ctx.Done():
			case out <- info.ID:
			}
		}
	}()
	return out
}

func (nc *networkClient) ConnectionManager() ifconnmgr.ConnManager {
	return &ifconnmgr.NullConnMgr{}
}

type messagePasser struct {
	net    *network
	target peer.ID
	local  peer.ID
	ctx    context.Context
}

func (mp *messagePasser) SendMsg(ctx context.Context, m bsmsg.BitSwapMessage) error {
	return mp.net.SendMessage(ctx, mp.local, mp.target, m)
}

func (mp *messagePasser) Close() error {
	return nil
}

func (mp *messagePasser) Reset() error {
	return nil
}

func (n *networkClient) NewMessageSender(ctx context.Context, p peer.ID) (bsnet.MessageSender, error) {
	return &messagePasser{
		net:    n.network,
		target: p,
		local:  n.local,
		ctx:    ctx,
	}, nil
}

// Provide provides the key to the network
func (nc *networkClient) Provide(ctx context.Context, k *cid.Cid) error {
	return nc.routing.Provide(ctx, k, true)
}

func (nc *networkClient) SetDelegate(r bsnet.Receiver) {
	nc.Receiver = r
}

func (nc *networkClient) ConnectTo(_ context.Context, p peer.ID) error {
	nc.network.mu.Lock()

	otherClient, ok := nc.network.clients[p]
	if !ok {
		nc.network.mu.Unlock()
		return errors.New("no such peer in network")
	}

	tag := tagForPeers(nc.local, p)
	if _, ok := nc.network.conns[tag]; ok {
		nc.network.mu.Unlock()
		log.Warning("ALREADY CONNECTED TO PEER (is this a reconnect? test lib needs fixing)")
		return nil
	}
	nc.network.conns[tag] = struct{}{}
	nc.network.mu.Unlock()

	// TODO: add handling for disconnects

	otherClient.PeerConnected(nc.local)
	nc.Receiver.PeerConnected(p)
	return nil
}

func tagForPeers(a, b peer.ID) string {
	if a < b {
		return string(a + b)
	}
	return string(b + a)
}
<commit_msg>bitswap virtual test net code should send messages in order<commit_after>package bitswap

import (
	"context"
	"errors"
	"sync"
	"time"

	bsmsg "github.com/ipfs/go-ipfs/exchange/bitswap/message"
	bsnet "github.com/ipfs/go-ipfs/exchange/bitswap/network"
	mockrouting "github.com/ipfs/go-ipfs/routing/mock"
	delay "github.com/ipfs/go-ipfs/thirdparty/delay"

	logging "gx/ipfs/QmRb5jh8z2E8hMGN2tkvs1yHynUanqnZ3UeKwgN1i9P1F8/go-log"
	routing "gx/ipfs/QmTiWLZ6Fo5j4KcTVutZJ5KWRRJrbxzmxA4td8NfEdrPh7/go-libp2p-routing"
	testutil "gx/ipfs/QmVvkK7s5imCiq3JVbL3pGfnhcCnf3LrFJPF4GE2sAoGZf/go-testutil"
	peer "gx/ipfs/QmZoWKhxUmZ2seW4BzX6fJkNR8hh9PsGModr7q171yq2SS/go-libp2p-peer"
	ifconnmgr "gx/ipfs/Qmax8X1Kfahf5WfSB68EWDG3d3qyS3Sqs1v412fjPTfRwx/go-libp2p-interface-connmgr"
	cid "gx/ipfs/QmcZfnkapfECQGcLZaf9B79NRg7cRa9EnZh4LSbkCzwNvY/go-cid"
)

var log = logging.Logger("bstestnet")

func VirtualNetwork(rs mockrouting.Server, d delay.D) Network {
	return &network{
		clients:       make(map[peer.ID]*receiverQueue),
		delay:         d,
		routingserver: rs,
		conns:         make(map[string]struct{}),
	}
}

type network struct {
	mu            sync.Mutex
	clients       map[peer.ID]*receiverQueue
	routingserver mockrouting.Server
	delay         delay.D
	conns         map[string]struct{}
}

type message struct {
	from       peer.ID
	msg        bsmsg.BitSwapMessage
	shouldSend time.Time
}

// receiverQueue queues up a set of messages to be sent, and sends them *in
// order* with their delays respected as much as sending them in order allows
// for
type receiverQueue struct {
	receiver bsnet.Receiver
	queue    []*message
	active   bool
	lk       sync.Mutex
}

func (n *network) Adapter(p testutil.Identity) bsnet.BitSwapNetwork {
	n.mu.Lock()
	defer n.mu.Unlock()

	client := &networkClient{
		local:   p.ID(),
		network: n,
		routing: n.routingserver.Client(p),
	}
	n.clients[p.ID()] = &receiverQueue{receiver: client}
	return client
}

func (n *network) HasPeer(p peer.ID) bool {
	n.mu.Lock()
	defer n.mu.Unlock()

	_, found := n.clients[p]
	return found
}

// TODO should this be completely asynchronous?
// TODO what does the network layer do with errors received from services?
func (n *network) SendMessage(
	ctx context.Context,
	from peer.ID,
	to peer.ID,
	mes bsmsg.BitSwapMessage) error {

	n.mu.Lock()
	defer n.mu.Unlock()

	receiver, ok := n.clients[to]
	if !ok {
		return errors.New("Cannot locate peer on network")
	}

	// nb: terminate the context since the context wouldn't actually be passed
	// over the network in a real scenario

	msg := &message{
		from:       from,
		msg:        mes,
		shouldSend: time.Now().Add(n.delay.Get()),
	}
	receiver.enqueue(msg)

	return nil
}

func (n *network) deliver(
	r bsnet.Receiver, from peer.ID, message bsmsg.BitSwapMessage) error {
	if message == nil || from == "" {
		return errors.New("Invalid input")
	}

	n.delay.Wait()

	r.ReceiveMessage(context.TODO(), from, message)
	return nil
}

type networkClient struct {
	local peer.ID
	bsnet.Receiver
	network *network
	routing routing.IpfsRouting
}

func (nc *networkClient) SendMessage(
	ctx context.Context,
	to peer.ID,
	message bsmsg.BitSwapMessage) error {
	return nc.network.SendMessage(ctx, nc.local, to, message)
}

// FindProvidersAsync returns a channel of providers for the given key
func (nc *networkClient) FindProvidersAsync(ctx context.Context, k *cid.Cid, max int) <-chan peer.ID {

	// NB: this function duplicates the PeerInfo -> ID transformation in the
	// bitswap network adapter. Not to worry. This network client will be
	// deprecated once the ipfsnet.Mock is added. The code below is only
	// temporary.

	out := make(chan peer.ID)
	go func() {
		defer close(out)
		providers := nc.routing.FindProvidersAsync(ctx, k, max)
		for info := range providers {
			select {
			case <-ctx.Done():
			case out <- info.ID:
			}
		}
	}()
	return out
}

func (nc *networkClient) ConnectionManager() ifconnmgr.ConnManager {
	return &ifconnmgr.NullConnMgr{}
}

type messagePasser struct {
	net    *network
	target peer.ID
	local  peer.ID
	ctx    context.Context
}

func (mp *messagePasser) SendMsg(ctx context.Context, m bsmsg.BitSwapMessage) error {
	return mp.net.SendMessage(ctx, mp.local, mp.target, m)
}

func (mp *messagePasser) Close() error {
	return nil
}

func (mp *messagePasser) Reset() error {
	return nil
}

func (n *networkClient) NewMessageSender(ctx context.Context, p peer.ID) (bsnet.MessageSender, error) {
	return &messagePasser{
		net:    n.network,
		target: p,
		local:  n.local,
		ctx:    ctx,
	}, nil
}

// Provide provides the key to the network
func (nc *networkClient) Provide(ctx context.Context, k *cid.Cid) error {
	return nc.routing.Provide(ctx, k, true)
}

func (nc *networkClient) SetDelegate(r bsnet.Receiver) {
	nc.Receiver = r
}

func (nc *networkClient) ConnectTo(_ context.Context, p peer.ID) error {
	nc.network.mu.Lock()

	otherClient, ok := nc.network.clients[p]
	if !ok {
		nc.network.mu.Unlock()
		return errors.New("no such peer in network")
	}

	tag := tagForPeers(nc.local, p)
	if _, ok := nc.network.conns[tag]; ok {
		nc.network.mu.Unlock()
		log.Warning("ALREADY CONNECTED TO PEER (is this a reconnect? test lib needs fixing)")
		return nil
	}
	nc.network.conns[tag] = struct{}{}
	nc.network.mu.Unlock()

	// TODO: add handling for disconnects

	otherClient.receiver.PeerConnected(nc.local)
	nc.Receiver.PeerConnected(p)
	return nil
}

func (rq *receiverQueue) enqueue(m *message) {
	rq.lk.Lock()
	defer rq.lk.Unlock()
	rq.queue = append(rq.queue, m)
	if !rq.active {
		rq.active = true
		go rq.process()
	}
}

func (rq *receiverQueue) process() {
	for {
		rq.lk.Lock()
		if len(rq.queue) == 0 {
			rq.active = false
			rq.lk.Unlock()
			return
		}
		m := rq.queue[0]
		rq.queue = rq.queue[1:]
		rq.lk.Unlock()

		time.Sleep(time.Until(m.shouldSend))
		rq.receiver.ReceiveMessage(context.TODO(), m.from, m.msg)
	}
}

func tagForPeers(a, b peer.ID) string {
	if a < b {
		return string(a + b)
	}
	return string(b + a)
}
<|endoftext|>


<commit_before>// This Source Code Form is subject to the terms of the Mozilla Public
// License, v. 2.0. If a copy of the MPL was not distributed with this
// file, You can obtain one at http://mozilla.org/MPL/2.0/.
//
// Contributor: Julien Vehent <ulfr@mozilla.com>

package aws

import (
	"crypto/rand"
	"fmt"
	"log"
	"strings"

	"github.com/aws/aws-sdk-go/aws"
	"github.com/aws/aws-sdk-go/aws/awsutil"
	awscred "github.com/aws/aws-sdk-go/aws/credentials"
	"github.com/aws/aws-sdk-go/aws/session"
	"github.com/aws/aws-sdk-go/service/iam"
	"github.com/mozilla-services/userplex/modules"
)

func init() {
	modules.Register("aws", new(module))
}

type module struct {
}

func (m *module) NewRun(c modules.Configuration) modules.Runner {
	r := new(run)
	r.Conf = c
	return r
}

type run struct {
	Conf modules.Configuration
	p    parameters
	c    credentials
	cli  *iam.IAM
}

type parameters struct {
	IamGroups      []string
	NotifyNewUsers bool
	SmtpRelay      string
	SmtpFrom       string
	AccountName    string
}

type credentials struct {
	AccessKey string
	SecretKey string
}

func (r *run) Run() (err error) {
	err = r.Conf.GetParameters(&r.p)
	if err != nil {
		return
	}
	err = r.Conf.GetCredentials(&r.c)
	if err != nil {
		return
	}
	r.cli = r.initIamClient()
	if r.cli == nil {
		return fmt.Errorf("failed to connect to aws using access key %q", r.c.AccessKey)
	}
	// Retrieve a list of ldap users from the groups configured
	ldapers := r.getLdapers()
	// create or add the users to groups.
	if r.Conf.Create {
		for uid, haspubkey := range ldapers {
			resp, err := r.cli.GetUser(&iam.GetUserInput{
				UserName: aws.String(uid),
			})
			if err != nil || resp == nil {
				log.Printf("[info] aws: user %q not found, needs to be created", uid)
				if !haspubkey && r.Conf.NotifyUsers {
					log.Printf("[warning] aws: %q has no PGP fingerprint in LDAP, skipping creation", uid)
					continue
				}
				r.createIamUser(uid)
			} else {
				r.updateUserGroups(uid)
			}
		}
	}
	// find which users are no longer in ldap and needs to be removed from aws
	if r.Conf.Delete {
		iamers := r.getIamers()
		for uid, _ := range iamers {
			if _, ok := ldapers[uid]; !ok {
				log.Printf("[info] aws: %q found in IAM group but not in LDAP, needs deletion", uid)
				r.removeIamUser(uid)
			}
		}
	}
	return
}

func (r *run) getLdapers() (lgm map[string]bool) {
	lgm = make(map[string]bool)
	users, err := r.Conf.LdapCli.GetEnabledUsersInGroups(r.Conf.LdapGroups)
	if err != nil {
		return
	}
	for _, user := range users {
		shortdn := strings.Split(user, ",")[0]
		uid, err := r.Conf.LdapCli.GetUserId(shortdn)
		if err != nil {
			log.Printf("[warning] aws: ldap query failed with error %v", err)
			continue
		}
		// apply the uid map: only store the translated uid in the ldapuid
		for _, mapping := range r.Conf.UidMap {
			if mapping.LdapUid == uid {
				uid = mapping.LocalUID
			}
		}
		lgm[uid] = true
		if r.Conf.NotifyUsers {
			// make sure we can find a PGP public key for the user to encrypt the notification
			// if no pubkey is found, log an error and set the user's entry to False
			_, err = r.Conf.LdapCli.GetUserPGPKey(shortdn)
			if err != nil {
				log.Printf("[warning] aws: no pgp public key could be found for %s: %v", shortdn, err)
				lgm[uid] = false
			}
		}
	}
	return
}

func (r *run) getIamers() (igm map[string]bool) {
	igm = make(map[string]bool)
	for _, group := range r.p.IamGroups {
		resp, err := r.cli.GetGroup(&iam.GetGroupInput{
			GroupName: aws.String(group),
		})
		if err != nil || resp == nil {
			log.Printf("[error] failed to retrieve users from IAM group %q: %v", group, err)
			continue
		}
		for _, user := range resp.Users {
			iamuser := strings.Replace(awsutil.Prettify(user.UserName), `"`, ``, -1)
			igm[iamuser] = true
		}
	}
	return
}

func (r *run) initIamClient() *iam.IAM {
	var awsconf aws.Config
	if r.c.AccessKey != "" && r.c.SecretKey != "" {
		awscreds := awscred.NewStaticCredentials(r.c.AccessKey, r.c.SecretKey, "")
		awsconf.Credentials = awscreds
	}
	return iam.New(session.New(), &awsconf)
}

// create a user in aws, assign temporary credentials and force password change, add the
// user to the necessary groups, and send it an email
func (r *run) createIamUser(uid string) {
	var (
		cuo  *iam.CreateUserOutput
		clpo *iam.CreateLoginProfileOutput
	)
	password := "P" + randToken() + "%"
	mail, err := r.Conf.LdapCli.GetUserEmailByUid(uid)
	if err != nil {
		log.Printf("[error] aws: couldn't find email of user %q in ldap, notification not sent: %v", uid, err)
		return
	}
	if !r.Conf.ApplyChanges {
		log.Printf("[dryrun] aws: would have created AWS IAM user %q with password %q", uid, password)
		goto notify
	}
	cuo, err = r.cli.CreateUser(&iam.CreateUserInput{
		UserName: aws.String(uid),
	})
	if err != nil || cuo == nil {
		log.Printf("[error] aws: failed to create user %q: %v", uid, err)
		return
	}
	clpo, err = r.cli.CreateLoginProfile(&iam.CreateLoginProfileInput{
		Password:              aws.String(password),
		UserName:              aws.String(uid),
		PasswordResetRequired: aws.Bool(true),
	})
	if err != nil || clpo == nil {
		log.Printf("[error] aws: failed to create user %q: %v", uid, err)
		return
	}
	for _, group := range r.p.IamGroups {
		r.addUserToIamGroup(uid, group)
	}
notify:
	// notify user
	rcpt := r.Conf.Notify.Recipient
	if rcpt == "{ldap:mail}" {
		rcpt = mail
	}
	r.Conf.Notify.Channel <- modules.Notification{
		Module:      "aws",
		Recipient:   rcpt,
		Mode:        r.Conf.Notify.Mode,
		MustEncrypt: true,
		Body: []byte(fmt.Sprintf(`New AWS account:
login: %s
pass:  %s (change at first login)
url:   %s`, uid, password, fmt.Sprintf("https://%s.signin.aws.amazon.com/console", r.p.AccountName))),
	}
	return
}

func (r *run) updateUserGroups(uid string) {
	gresp, err := r.cli.ListGroupsForUser(&iam.ListGroupsForUserInput{
		UserName: aws.String(uid),
	})
	if err != nil || gresp == nil {
		log.Printf("[info] aws: groups of user %q not found, needs to be added", uid)
	}
	// iterate through the groups and find the missing ones
	for _, iamgroup := range r.p.IamGroups {
		found := false
		for _, group := range gresp.Groups {
			gname := strings.Replace(awsutil.Prettify(group.GroupName), `"`, ``, -1)
			if iamgroup == gname {
				found = true
			}
		}
		if !found {
			r.addUserToIamGroup(uid, iamgroup)
		}
	}
	return
}

func (r *run) addUserToIamGroup(uid, group string) {
	if !r.Conf.Create {
		return
	}
	if !r.Conf.ApplyChanges {
		log.Printf("[dryrun] aws: would have added AWS IAM user %q to group %q", uid, group)
		return
	}
	resp, err := r.cli.AddUserToGroup(&iam.AddUserToGroupInput{
		GroupName: aws.String(group),
		UserName:  aws.String(uid),
	})
	if err != nil || resp == nil {
		log.Printf("[error] aws: failed to add user %q to group %q: %v", uid, group, err)
	}
	return
}

func (r *run) removeIamUser(uid string) {
	var (
		err  error
		lgfu *iam.ListGroupsForUserOutput
		dlpo *iam.DeleteLoginProfileOutput
		duo  *iam.DeleteUserOutput
	)
	if !r.Conf.Delete {
		return
	}
	if !r.Conf.ApplyChanges {
		log.Printf("[dryrun] aws: would have deleted AWS IAM user %q", uid)
		goto notify
	}
	lgfu, err = r.cli.ListGroupsForUser(&iam.ListGroupsForUserInput{
		UserName: aws.String(uid),
	})
	if err != nil || lgfu == nil {
		log.Printf("[error] aws: failed to list groups for user %q: %v", uid, err)
		return
	}
	// iterate through the groups and find the missing ones
	for _, iamgroup := range r.p.IamGroups {
		resp, err := r.cli.RemoveUserFromGroup(&iam.RemoveUserFromGroupInput{
			GroupName: &iamgroup,
			UserName:  aws.String(uid),
		})
		gname := strings.Replace(awsutil.Prettify(iamgroup), `"`, ``, -1)
		if err != nil || resp == nil {
			log.Printf("[error] aws: failed to remove user %q from group %q: %v", uid, gname, err)
		} else {
			log.Printf("[info] aws: removed user %q from group %q", uid, gname)
		}
	}
	dlpo, err = r.cli.DeleteLoginProfile(&iam.DeleteLoginProfileInput{
		UserName: aws.String(uid),
	})
	if err != nil || dlpo == nil {
		log.Printf("[error] aws: failed to delete aws login profile for user %q: %v", uid, err)
		return
	}
	duo, err = r.cli.DeleteUser(&iam.DeleteUserInput{
		UserName: aws.String(uid),
	})
	if err != nil || duo == nil {
		log.Printf("[error] aws: failed to delete aws user %q: %v", uid, err)
		return
	}
notify:
	// notify user
	rcpt := r.Conf.Notify.Recipient
	if rcpt == "{ldap:mail}" {
		rcpt, err = r.Conf.LdapCli.GetUserEmailByUid(uid)
		if err != nil {
			log.Printf("[error] aws: couldn't find email of user %q in ldap, notification not sent: %v", uid, err)
			return
		}
	}
	r.Conf.Notify.Channel <- modules.Notification{
		Module:      "aws",
		Recipient:   rcpt,
		Mode:        r.Conf.Notify.Mode,
		MustEncrypt: false,
		Body: []byte(fmt.Sprintf(`Deleted AWS account:
The account %q has been removed from %q.`, uid, r.p.AccountName)),
	}
	return
}

func randToken() string {
	b := make([]byte, 8)
	rand.Read(b)
	return fmt.Sprintf("%x", b)
}
<commit_msg>aws: add intermediate steps to removeIamUser<commit_after>// This Source Code Form is subject to the terms of the Mozilla Public
// License, v. 2.0. If a copy of the MPL was not distributed with this
// file, You can obtain one at http://mozilla.org/MPL/2.0/.
//
// Contributor: Julien Vehent <ulfr@mozilla.com>

package aws

import (
	"crypto/rand"
	"fmt"
	"log"
	"strings"

	"github.com/aws/aws-sdk-go/aws"
	"github.com/aws/aws-sdk-go/aws/awsutil"
	awscred "github.com/aws/aws-sdk-go/aws/credentials"
	"github.com/aws/aws-sdk-go/aws/session"
	"github.com/aws/aws-sdk-go/service/iam"
	"github.com/mozilla-services/userplex/modules"
)

func init() {
	modules.Register("aws", new(module))
}

type module struct {
}

func (m *module) NewRun(c modules.Configuration) modules.Runner {
	r := new(run)
	r.Conf = c
	return r
}

type run struct {
	Conf modules.Configuration
	p    parameters
	c    credentials
	cli  *iam.IAM
}

type parameters struct {
	IamGroups      []string
	NotifyNewUsers bool
	SmtpRelay      string
	SmtpFrom       string
	AccountName    string
}

type credentials struct {
	AccessKey string
	SecretKey string
}

func (r *run) Run() (err error) {
	err = r.Conf.GetParameters(&r.p)
	if err != nil {
		return
	}
	err = r.Conf.GetCredentials(&r.c)
	if err != nil {
		return
	}
	r.cli = r.initIamClient()
	if r.cli == nil {
		return fmt.Errorf("failed to connect to aws using access key %q", r.c.AccessKey)
	}
	// Retrieve a list of ldap users from the groups configured
	ldapers := r.getLdapers()
	// create or add the users to groups.
	if r.Conf.Create {
		for uid, haspubkey := range ldapers {
			resp, err := r.cli.GetUser(&iam.GetUserInput{
				UserName: aws.String(uid),
			})
			if err != nil || resp == nil {
				log.Printf("[info] aws: user %q not found, needs to be created", uid)
				if !haspubkey && r.Conf.NotifyUsers {
					log.Printf("[warning] aws: %q has no PGP fingerprint in LDAP, skipping creation", uid)
					continue
				}
				r.createIamUser(uid)
			} else {
				r.updateUserGroups(uid)
			}
		}
	}
	// find which users are no longer in ldap and needs to be removed from aws
	if r.Conf.Delete {
		iamers := r.getIamers()
		for uid, _ := range iamers {
			if _, ok := ldapers[uid]; !ok {
				log.Printf("[info] aws: %q found in IAM group but not in LDAP, needs deletion", uid)
				r.removeIamUser(uid)
			}
		}
	}
	return
}

func (r *run) getLdapers() (lgm map[string]bool) {
	lgm = make(map[string]bool)
	users, err := r.Conf.LdapCli.GetEnabledUsersInGroups(r.Conf.LdapGroups)
	if err != nil {
		return
	}
	for _, user := range users {
		shortdn := strings.Split(user, ",")[0]
		uid, err := r.Conf.LdapCli.GetUserId(shortdn)
		if err != nil {
			log.Printf("[warning] aws: ldap query failed with error %v", err)
			continue
		}
		// apply the uid map: only store the translated uid in the ldapuid
		for _, mapping := range r.Conf.UidMap {
			if mapping.LdapUid == uid {
				uid = mapping.LocalUID
			}
		}
		lgm[uid] = true
		if r.Conf.NotifyUsers {
			// make sure we can find a PGP public key for the user to encrypt the notification
			// if no pubkey is found, log an error and set the user's entry to False
			_, err = r.Conf.LdapCli.GetUserPGPKey(shortdn)
			if err != nil {
				log.Printf("[warning] aws: no pgp public key could be found for %s: %v", shortdn, err)
				lgm[uid] = false
			}
		}
	}
	return
}

func (r *run) getIamers() (igm map[string]bool) {
	igm = make(map[string]bool)
	for _, group := range r.p.IamGroups {
		resp, err := r.cli.GetGroup(&iam.GetGroupInput{
			GroupName: aws.String(group),
		})
		if err != nil || resp == nil {
			log.Printf("[error] failed to retrieve users from IAM group %q: %v", group, err)
			continue
		}
		for _, user := range resp.Users {
			iamuser := strings.Replace(awsutil.Prettify(user.UserName), `"`, ``, -1)
			igm[iamuser] = true
		}
	}
	return
}

func (r *run) initIamClient() *iam.IAM {
	var awsconf aws.Config
	if r.c.AccessKey != "" && r.c.SecretKey != "" {
		awscreds := awscred.NewStaticCredentials(r.c.AccessKey, r.c.SecretKey, "")
		awsconf.Credentials = awscreds
	}
	return iam.New(session.New(), &awsconf)
}

// create a user in aws, assign temporary credentials and force password change, add the
// user to the necessary groups, and send it an email
func (r *run) createIamUser(uid string) {
	var (
		cuo  *iam.CreateUserOutput
		clpo *iam.CreateLoginProfileOutput
	)
	password := "P" + randToken() + "%"
	mail, err := r.Conf.LdapCli.GetUserEmailByUid(uid)
	if err != nil {
		log.Printf("[error] aws: couldn't find email of user %q in ldap, notification not sent: %v", uid, err)
		return
	}
	if !r.Conf.ApplyChanges {
		log.Printf("[dryrun] aws: would have created AWS IAM user %q with password %q", uid, password)
		goto notify
	}
	cuo, err = r.cli.CreateUser(&iam.CreateUserInput{
		UserName: aws.String(uid),
	})
	if err != nil || cuo == nil {
		log.Printf("[error] aws: failed to create user %q: %v", uid, err)
		return
	}
	clpo, err = r.cli.CreateLoginProfile(&iam.CreateLoginProfileInput{
		Password:              aws.String(password),
		UserName:              aws.String(uid),
		PasswordResetRequired: aws.Bool(true),
	})
	if err != nil || clpo == nil {
		log.Printf("[error] aws: failed to create user %q: %v", uid, err)
		return
	}
	for _, group := range r.p.IamGroups {
		r.addUserToIamGroup(uid, group)
	}
notify:
	// notify user
	rcpt := r.Conf.Notify.Recipient
	if rcpt == "{ldap:mail}" {
		rcpt = mail
	}
	r.Conf.Notify.Channel <- modules.Notification{
		Module:      "aws",
		Recipient:   rcpt,
		Mode:        r.Conf.Notify.Mode,
		MustEncrypt: true,
		Body: []byte(fmt.Sprintf(`New AWS account:
login: %s
pass:  %s (change at first login)
url:   %s`, uid, password, fmt.Sprintf("https://%s.signin.aws.amazon.com/console", r.p.AccountName))),
	}
	return
}

func (r *run) updateUserGroups(uid string) {
	gresp, err := r.cli.ListGroupsForUser(&iam.ListGroupsForUserInput{
		UserName: aws.String(uid),
	})
	if err != nil || gresp == nil {
		log.Printf("[info] aws: groups of user %q not found, needs to be added", uid)
	}
	// iterate through the groups and find the missing ones
	for _, iamgroup := range r.p.IamGroups {
		found := false
		for _, group := range gresp.Groups {
			gname := strings.Replace(awsutil.Prettify(group.GroupName), `"`, ``, -1)
			if iamgroup == gname {
				found = true
			}
		}
		if !found {
			r.addUserToIamGroup(uid, iamgroup)
		}
	}
	return
}

func (r *run) addUserToIamGroup(uid, group string) {
	if !r.Conf.Create {
		return
	}
	if !r.Conf.ApplyChanges {
		log.Printf("[dryrun] aws: would have added AWS IAM user %q to group %q", uid, group)
		return
	}
	resp, err := r.cli.AddUserToGroup(&iam.AddUserToGroupInput{
		GroupName: aws.String(group),
		UserName:  aws.String(uid),
	})
	if err != nil || resp == nil {
		log.Printf("[error] aws: failed to add user %q to group %q: %v", uid, group, err)
	}
	return
}

func (r *run) removeIamUser(uid string) {
	var (
		err  error
		lgfu *iam.ListGroupsForUserOutput
		dlpo *iam.DeleteLoginProfileOutput
		duo  *iam.DeleteUserOutput
	)
	if !r.Conf.Delete {
		return
	}
	// remove all user's access keys
	lakfu, err := r.cli.ListAccessKeys(&iam.ListAccessKeysInput{
		UserName: aws.String(uid),
	})
	if err != nil || lakfu == nil {
		log.Printf("[error] aws: failed to list access keys for user %q: %v", uid, err)
		return
	}
	for _, accesskey := range lakfu.AccessKeyMetadata {
		keyid := strings.Replace(awsutil.Prettify(accesskey.AccessKeyId), `"`, ``, -1)
		if !r.Conf.ApplyChanges {
			log.Printf("[dryrun] aws: would have removed access key id %q of user %q", keyid, uid)
			continue
		}
		daki := &iam.DeleteAccessKeyInput{
			AccessKeyId: accesskey.AccessKeyId,
			UserName:    aws.String(uid),
		}
		resp, err := r.cli.DeleteAccessKey(daki)
		if err != nil || resp == nil {
			log.Printf("[error] aws: failed to delete access key %q of user %q: %v. request was %q.",
				keyid, uid, err, daki.String())
		} else {
			log.Printf("[info] aws: deleted access key %q of user %q", keyid, uid)
		}

	}
	// remove the user from all IAM groups
	lgfu, err = r.cli.ListGroupsForUser(&iam.ListGroupsForUserInput{
		UserName: aws.String(uid),
	})
	if err != nil || lgfu == nil {
		log.Printf("[error] aws: failed to list groups for user %q: %v", uid, err)
		return
	}
	// iterate through the groups and find the missing ones
	for _, iamgroup := range lgfu.Groups {
		gname := strings.Replace(awsutil.Prettify(iamgroup.GroupName), `"`, ``, -1)
		if !r.Conf.ApplyChanges {
			log.Printf("[dryrun] aws: would have removed user %q from group %q", uid, gname)
			continue
		}
		rufgi := &iam.RemoveUserFromGroupInput{
			GroupName: iamgroup.GroupName,
			UserName:  aws.String(uid),
		}
		resp, err := r.cli.RemoveUserFromGroup(rufgi)
		if err != nil || resp == nil {
			log.Printf("[error] aws: failed to remove user %q from group %q: %v. request was %q.",
				uid, gname, err, rufgi.String())
		} else {
			log.Printf("[info] aws: removed user %q from group %q", uid, gname)
		}
	}
	if !r.Conf.ApplyChanges {
		log.Printf("[dryrun] aws: would have deleted AWS IAM user %q", uid)
		goto notify
	}
	dlpo, err = r.cli.DeleteLoginProfile(&iam.DeleteLoginProfileInput{
		UserName: aws.String(uid),
	})
	if err != nil || dlpo == nil {
		log.Printf("[info] aws: user %q did not have an aws login profile to delete", uid)
	}
	duo, err = r.cli.DeleteUser(&iam.DeleteUserInput{
		UserName: aws.String(uid),
	})
	if err != nil || duo == nil {
		log.Printf("[error] aws: failed to delete aws user %q: %v", uid, err)
		return
	} else {
		log.Printf("[info] aws: deleted user %q", uid)
	}
notify:
	// notify user
	rcpt := r.Conf.Notify.Recipient
	if rcpt == "{ldap:mail}" {
		rcpt, err = r.Conf.LdapCli.GetUserEmailByUid(uid)
		if err != nil {
			log.Printf("[error] aws: couldn't find email of user %q in ldap, notification not sent: %v", uid, err)
			return
		}
	}
	r.Conf.Notify.Channel <- modules.Notification{
		Module:      "aws",
		Recipient:   rcpt,
		Mode:        r.Conf.Notify.Mode,
		MustEncrypt: false,
		Body: []byte(fmt.Sprintf(`Deleted AWS account:
The account %q has been removed from %q.`, uid, r.p.AccountName)),
	}
	return
}

func randToken() string {
	b := make([]byte, 8)
	rand.Read(b)
	return fmt.Sprintf("%x", b)
}
<|endoftext|>


<commit_before>package utilities

import (
	"sort"
)

// DoubleArray is a Double Array implementation of trie on sequences of strings.
type DoubleArray struct {
	// Encoding keeps an encoding from string to int
	Encoding map[string]int
	// Base is the base array of Double Array
	Base []int
	// Check is the check array of Double Array
	Check []int
}

// NewDoubleArray builds a DoubleArray from a set of sequences of strings.
func NewDoubleArray(seqs [][]string) *DoubleArray {
	da := &DoubleArray{Encoding: make(map[string]int)}
	if len(seqs) == 0 {
		return da
	}

	encoded := registerTokens(da, seqs)
	sort.Sort(byLex(encoded))

	root := node{row: -1, col: -1, left: 0, right: len(encoded)}
	addSeqs(da, encoded, 0, root)

	for i := len(da.Base); i > 0; i-- {
		if da.Check[i-1] != 0 {
			da.Base = da.Base[:i]
			da.Check = da.Check[:i]
			break
		}
	}
	return da
}

func registerTokens(da *DoubleArray, seqs [][]string) [][]int {
	var result [][]int
	for _, seq := range seqs {
		var encoded []int
		for _, token := range seq {
			if _, ok := da.Encoding[token]; !ok {
				da.Encoding[token] = len(da.Encoding)
			}
			encoded = append(encoded, da.Encoding[token])
		}
		result = append(result, encoded)
	}
	for i := range result {
		result[i] = append(result[i], len(da.Encoding))
	}
	return result
}

type node struct {
	row, col    int
	left, right int
}

func (n node) value(seqs [][]int) int {
	return seqs[n.row][n.col]
}

func (n node) children(seqs [][]int) []*node {
	var result []*node
	lastVal := int(-1)
	last := new(node)
	for i := n.left; i < n.right; i++ {
		if lastVal == seqs[i][n.col+1] {
			continue
		}
		last.right = i
		last = &node{
			row:  i,
			col:  n.col + 1,
			left: i,
		}
		result = append(result, last)
	}
	last.right = n.right
	return result
}

func addSeqs(da *DoubleArray, seqs [][]int, pos int, n node) {
	ensureSize(da, pos)

	children := n.children(seqs)
	var i int
	for i = 1; ; i++ {
		ok := func() bool {
			for _, child := range children {
				code := child.value(seqs)
				j := i + code
				ensureSize(da, j)
				if da.Check[j] != 0 {
					return false
				}
			}
			return true
		}()
		if ok {
			break
		}
	}
	da.Base[pos] = i
	for _, child := range children {
		code := child.value(seqs)
		j := i + code
		da.Check[j] = pos + 1
	}
	terminator := len(da.Encoding)
	for _, child := range children {
		code := child.value(seqs)
		if code == terminator {
			continue
		}
		j := i + code
		addSeqs(da, seqs, j, *child)
	}
}

func ensureSize(da *DoubleArray, i int) {
	for i >= len(da.Base) {
		da.Base = append(da.Base, make([]int, len(da.Base)+1)...)
		da.Check = append(da.Check, make([]int, len(da.Check)+1)...)
	}
}

type byLex [][]int

func (l byLex) Len() int      { return len(l) }
func (l byLex) Swap(i, j int) { l[i], l[j] = l[j], l[i] }
func (l byLex) Less(i, j int) bool {
	si := l[i]
	sj := l[j]
	var k int
	for k = 0; k < len(si) && k < len(sj); k++ {
		if si[k] < sj[k] {
			return true
		}
		if si[k] > sj[k] {
			return false
		}
	}
	return k < len(sj)
}

// HasCommonPrefix determines if any sequence in the DoubleArray is a prefix of the given sequence.
func (da *DoubleArray) HasCommonPrefix(seq []string) bool {
	if len(da.Base) == 0 {
		return false
	}

	var i int
	for _, t := range seq {
		code, ok := da.Encoding[t]
		if !ok {
			break
		}
		j := da.Base[i] + code
		if len(da.Check) <= j || da.Check[j] != i+1 {
			break
		}
		i = j
	}
	j := da.Base[i] + len(da.Encoding)
	if len(da.Check) <= j || da.Check[j] != i+1 {
		return false
	}
	return true
}
<commit_msg>utilities: Add preallocate for encoded (#3024)<commit_after>package utilities

import (
	"sort"
)

// DoubleArray is a Double Array implementation of trie on sequences of strings.
type DoubleArray struct {
	// Encoding keeps an encoding from string to int
	Encoding map[string]int
	// Base is the base array of Double Array
	Base []int
	// Check is the check array of Double Array
	Check []int
}

// NewDoubleArray builds a DoubleArray from a set of sequences of strings.
func NewDoubleArray(seqs [][]string) *DoubleArray {
	da := &DoubleArray{Encoding: make(map[string]int)}
	if len(seqs) == 0 {
		return da
	}

	encoded := registerTokens(da, seqs)
	sort.Sort(byLex(encoded))

	root := node{row: -1, col: -1, left: 0, right: len(encoded)}
	addSeqs(da, encoded, 0, root)

	for i := len(da.Base); i > 0; i-- {
		if da.Check[i-1] != 0 {
			da.Base = da.Base[:i]
			da.Check = da.Check[:i]
			break
		}
	}
	return da
}

func registerTokens(da *DoubleArray, seqs [][]string) [][]int {
	var result [][]int
	for _, seq := range seqs {
		encoded := make([]int, 0, len(seq))
		for _, token := range seq {
			if _, ok := da.Encoding[token]; !ok {
				da.Encoding[token] = len(da.Encoding)
			}
			encoded = append(encoded, da.Encoding[token])
		}
		result = append(result, encoded)
	}
	for i := range result {
		result[i] = append(result[i], len(da.Encoding))
	}
	return result
}

type node struct {
	row, col    int
	left, right int
}

func (n node) value(seqs [][]int) int {
	return seqs[n.row][n.col]
}

func (n node) children(seqs [][]int) []*node {
	var result []*node
	lastVal := int(-1)
	last := new(node)
	for i := n.left; i < n.right; i++ {
		if lastVal == seqs[i][n.col+1] {
			continue
		}
		last.right = i
		last = &node{
			row:  i,
			col:  n.col + 1,
			left: i,
		}
		result = append(result, last)
	}
	last.right = n.right
	return result
}

func addSeqs(da *DoubleArray, seqs [][]int, pos int, n node) {
	ensureSize(da, pos)

	children := n.children(seqs)
	var i int
	for i = 1; ; i++ {
		ok := func() bool {
			for _, child := range children {
				code := child.value(seqs)
				j := i + code
				ensureSize(da, j)
				if da.Check[j] != 0 {
					return false
				}
			}
			return true
		}()
		if ok {
			break
		}
	}
	da.Base[pos] = i
	for _, child := range children {
		code := child.value(seqs)
		j := i + code
		da.Check[j] = pos + 1
	}
	terminator := len(da.Encoding)
	for _, child := range children {
		code := child.value(seqs)
		if code == terminator {
			continue
		}
		j := i + code
		addSeqs(da, seqs, j, *child)
	}
}

func ensureSize(da *DoubleArray, i int) {
	for i >= len(da.Base) {
		da.Base = append(da.Base, make([]int, len(da.Base)+1)...)
		da.Check = append(da.Check, make([]int, len(da.Check)+1)...)
	}
}

type byLex [][]int

func (l byLex) Len() int      { return len(l) }
func (l byLex) Swap(i, j int) { l[i], l[j] = l[j], l[i] }
func (l byLex) Less(i, j int) bool {
	si := l[i]
	sj := l[j]
	var k int
	for k = 0; k < len(si) && k < len(sj); k++ {
		if si[k] < sj[k] {
			return true
		}
		if si[k] > sj[k] {
			return false
		}
	}
	return k < len(sj)
}

// HasCommonPrefix determines if any sequence in the DoubleArray is a prefix of the given sequence.
func (da *DoubleArray) HasCommonPrefix(seq []string) bool {
	if len(da.Base) == 0 {
		return false
	}

	var i int
	for _, t := range seq {
		code, ok := da.Encoding[t]
		if !ok {
			break
		}
		j := da.Base[i] + code
		if len(da.Check) <= j || da.Check[j] != i+1 {
			break
		}
		i = j
	}
	j := da.Base[i] + len(da.Encoding)
	if len(da.Check) <= j || da.Check[j] != i+1 {
		return false
	}
	return true
}
<|endoftext|>


<commit_before>// Package bamstats provides functions for computing several mapping statistics on a BAM file.
package bamstats

import (
	"sync"
	"time"

	log "github.com/Sirupsen/logrus"
	"github.com/guigolab/bamstats/annotation"
	"github.com/guigolab/bamstats/config"
	"github.com/guigolab/bamstats/sam"
	"github.com/guigolab/bamstats/stats"
	"github.com/guigolab/bamstats/utils"
)

func init() {
	log.SetLevel(log.WarnLevel)
}

// var wg sync.WaitGroup

func worker(id int, in interface{}, out chan stats.Map, index *annotation.RtreeMap, cfg *config.Config, wg *sync.WaitGroup) {
	defer wg.Done()
	logger := log.WithFields(log.Fields{
		"worker": id,
	})
	logger.Debug("Starting")

	sm := makeStatsMap(index, cfg)

	collectStats(in, sm)

	logger.Debug("Done")

	out <- sm
}

func collectStats(in interface{}, sm stats.Map) {
	switch in.(type) {
	case chan *sam.Record:
		c := in.(chan *sam.Record)
		for record := range c {
			for _, s := range sm {
				s.Collect(record)
			}
		}
	case chan *sam.Iterator:
		iterators := in.(chan *sam.Iterator)
		for it := range iterators {
			for it.Next() {
				for _, s := range sm {
					s.Collect(it.Record())
				}
			}
		}
	}
}

func process(bamFile string, index *annotation.RtreeMap, cpu int, maxBuf int, reads int, uniq bool) (stats.Map, error) {

	var wg sync.WaitGroup

	conf := config.NewConfig(cpu, maxBuf, reads, uniq)

	br, err := sam.NewReader(bamFile, conf)
	defer br.Close()
	if err != nil {
		return nil, err
	}
	statChan := make(chan stats.Map, cpu)
	for i := 0; i < br.Workers; i++ {
		id := i + 1
		wg.Add(1)
		go worker(id, br.Channels[i], statChan, index, conf, &wg)
	}

	br.Read()

	go waitProcess(statChan, &wg)
	stat := <-statChan
	stat.Merge(statChan)
	for k, v := range stat {
		v.Finalize()
		if k == "general" {
			s := v.(*stats.GeneralStats)
			s.Reads.Total += br.Unmapped()
		}
	}

	return stat, nil
}

func waitProcess(st chan stats.Map, wg *sync.WaitGroup) {
	wg.Wait()
	close(st)
}

// Process process the input BAM file and collect different mapping stats.
func Process(bamFile string, anno string, cpu int, maxBuf int, reads int, uniq bool) (stats.Map, error) {
	var index *annotation.RtreeMap
	if anno != "" {
		log.Infof("Creating index for %s", anno)
		start := time.Now()
		index = annotation.CreateIndex(anno, bamFile, cpu)
		log.Infof("Index done in %v", time.Since(start))
	}
	start := time.Now()
	log.Infof("Collecting stats for %s", bamFile)
	stats, err := process(bamFile, index, cpu, maxBuf, reads, uniq)
	if err != nil {
		return nil, err
	}
	log.Infof("Stats done in %v", time.Since(start))
	return stats, nil
}

func writeOutput(output string, st stats.Map) {
	out := utils.NewOutput(output)
	utils.OutputJSON(out, st)
}

func makeStatsMap(index *annotation.RtreeMap, cfg *config.Config) stats.Map {
	m := make(stats.Map)
	m.Add("general", stats.NewGeneralStats())
	if index != nil {
		m.Add("coverage", stats.NewCoverageStats(index, false))
		if cfg.Uniq {
			m.Add("coverageUniq", stats.NewCoverageStats(index, true))
		}
	}
	return m
}
<commit_msg>Small naming changes in process.go<commit_after>// Package bamstats provides functions for computing several mapping statistics on a BAM file.
package bamstats

import (
	"sync"
	"time"

	log "github.com/Sirupsen/logrus"
	"github.com/guigolab/bamstats/annotation"
	"github.com/guigolab/bamstats/config"
	"github.com/guigolab/bamstats/sam"
	"github.com/guigolab/bamstats/stats"
	"github.com/guigolab/bamstats/utils"
)

func init() {
	log.SetLevel(log.WarnLevel)
}

// var wg sync.WaitGroup

func worker(id int, in interface{}, out chan stats.Map, index *annotation.RtreeMap, cfg *config.Config, wg *sync.WaitGroup) {
	defer wg.Done()
	logger := log.WithFields(log.Fields{
		"worker": id,
	})
	logger.Debug("Starting")

	sm := makeStatsMap(index, cfg)

	collectStats(in, sm)

	logger.Debug("Done")

	out <- sm
}

func collectStats(in interface{}, sm stats.Map) {
	switch in.(type) {
	case chan *sam.Record:
		c := in.(chan *sam.Record)
		for record := range c {
			for _, s := range sm {
				s.Collect(record)
			}
		}
	case chan *sam.Iterator:
		iterators := in.(chan *sam.Iterator)
		for it := range iterators {
			for it.Next() {
				for _, s := range sm {
					s.Collect(it.Record())
				}
			}
		}
	}
}

func process(bamFile string, index *annotation.RtreeMap, cpu int, maxBuf int, reads int, uniq bool) (stats.Map, error) {

	var wg sync.WaitGroup

	conf := config.NewConfig(cpu, maxBuf, reads, uniq)

	br, err := sam.NewReader(bamFile, conf)
	defer br.Close()
	if err != nil {
		return nil, err
	}
	statChan := make(chan stats.Map, cpu)
	for i := 0; i < br.Workers; i++ {
		id := i + 1
		wg.Add(1)
		go worker(id, br.Channels[i], statChan, index, conf, &wg)
	}

	br.Read()

	go waitProcess(statChan, &wg)
	stat := <-statChan
	stat.Merge(statChan)
	for k, v := range stat {
		v.Finalize()
		if k == "general" {
			s := v.(*stats.GeneralStats)
			s.Reads.Total += br.Unmapped()
		}
	}

	return stat, nil
}

func waitProcess(st chan stats.Map, wg *sync.WaitGroup) {
	wg.Wait()
	close(st)
}

// Process process the input BAM file and collect different mapping stats.
func Process(bamFile string, anno string, cpu int, maxBuf int, reads int, uniq bool) (stats.Map, error) {
	var index *annotation.RtreeMap
	if anno != "" {
		log.Infof("Creating index for %s", anno)
		start := time.Now()
		index = annotation.CreateIndex(anno, bamFile, cpu)
		log.Infof("Index done in %v", time.Since(start))
	}
	start := time.Now()
	log.Infof("Collecting stats for %s", bamFile)
	allStats, err := process(bamFile, index, cpu, maxBuf, reads, uniq)
	if err != nil {
		return nil, err
	}
	log.Infof("Stats done in %v", time.Since(start))
	return allStats, nil
}

func WriteOutput(output string, st stats.Map) {
	out := utils.NewOutput(output)
	utils.OutputJSON(out, st)
}

func makeStatsMap(index *annotation.RtreeMap, cfg *config.Config) stats.Map {
	m := make(stats.Map)
	m.Add("general", stats.NewGeneralStats())
	if index != nil {
		m.Add("coverage", stats.NewCoverageStats(index, false))
		if cfg.Uniq {
			m.Add("coverageUniq", stats.NewCoverageStats(index, true))
		}
	}
	return m
}
<|endoftext|>


<commit_before>package imgo

import (
	"errors"
)

//input a image as src , return a image matrix by sunseteffect process
func SunsetEffect(src string)(imgMatrix [][][]uint8 , err error) {
	imgMatrix,err = Read(src)

	if err != nil {
		return
	}

	height:=len(imgMatrix)
	width:=len(imgMatrix[0])

	for i:=0;i<height;i++{
		for j:=0;j<width;j++{
			imgMatrix[i][j][1] = uint8( float64(imgMatrix[i][j][1]) * 0.7 )
			imgMatrix[i][j][2] = uint8( float64(imgMatrix[i][j][2]) * 0.7 )
		}
	}

	return
}

// input a image as src , return a image matrix by negativefilmeffect process
func NegativeFilmEffect(src string)(imgMatrix [][][]uint8 , err error) {
	imgMatrix,err = Read(src)

	if err != nil {
		return
	}

	height:=len(imgMatrix)
	width:=len(imgMatrix[0])

	for i:=0;i<height;i++{
		for j:=0;j<width;j++{
			imgMatrix[i][j][0] = 255 - imgMatrix[i][j][0]
			imgMatrix[i][j][1] = 255 - imgMatrix[i][j][1]
			imgMatrix[i][j][2] = 255 - imgMatrix[i][j][2]
		}
	}

	return
}

func AdjustBrightness(src string , light float64)(imgMatrix [][][]uint8 , err error) {
	imgMatrix,err = Read(src)

	if err != nil {
		return
	}

	if light <= 0{
		err = errors.New("value of light must be more than 0")
		return
	}

	height:=len(imgMatrix)
	width:=len(imgMatrix[0])

	for i:=0;i<height;i++{
		for j:=0;j<width;j++{
			imgMatrix[i][j][0] = uint8(float64(imgMatrix[i][j][0])*light)
			imgMatrix[i][j][1] = uint8(float64(imgMatrix[i][j][1])*light)
			imgMatrix[i][j][2] = uint8(float64(imgMatrix[i][j][2])*light)
		}
	}

	return
}

// fuse two images and the size of new image is as src1
func ImageFusion(src1 string , src2 string)(imgMatrix [][][]uint8 , err error) {
	imgMatrix1,err1 := Read(src1)

	if err1 != nil {
		err = err1
		return
	}

	height:=len(imgMatrix1)
	width:=len(imgMatrix1[0])

	imgMatrix2,err2 := ResizeForMatrix(src2,width,height)

	if err2 != nil {
		err = err2
		return
	}

	for i:=0;i<height;i++{
		for j:=0;j<width;j++{
			imgMatrix1[i][j][0] = uint8(float64(imgMatrix1[i][j][0])*0.5)+uint8(float64(imgMatrix2[i][j][0])*0.5)
			imgMatrix1[i][j][1] = uint8(float64(imgMatrix1[i][j][1])*0.5)+uint8(float64(imgMatrix2[i][j][1])*0.5)
			imgMatrix1[i][j][2] = uint8(float64(imgMatrix1[i][j][2])*0.5)+uint8(float64(imgMatrix1[i][j][2])*0.5)
		}
	}
	imgMatrix = imgMatrix1
	return
}<commit_msg>Unified function input<commit_after>package imgo

import (
	"errors"
)

//input a image matrix as src , return a image matrix by sunseteffect process
func SunsetEffect(src [][][]uint8)(imgMatrix [][][]uint8 , err error) {
	imgMatrix = src

	height:=len(imgMatrix)
	width:=len(imgMatrix[0])
	if height == 0 || width == 0 {
		err = errors.New("The input of matrix is illegal!")
	}

	for i:=0;i<height;i++{
		for j:=0;j<width;j++{
			imgMatrix[i][j][1] = uint8( float64(imgMatrix[i][j][1]) * 0.7 )
			imgMatrix[i][j][2] = uint8( float64(imgMatrix[i][j][2]) * 0.7 )
		}
	}

	return
}

// input a image as src , return a image matrix by negativefilmeffect process
func NegativeFilmEffect(src [][][]uint8)(imgMatrix [][][]uint8 , err error) {
	imgMatrix = src

	height:=len(imgMatrix)
	width:=len(imgMatrix[0])
	if height == 0 || width == 0 {
		err = errors.New("The input of matrix is illegal!")
	}

	for i:=0;i<height;i++{
		for j:=0;j<width;j++{
			imgMatrix[i][j][0] = 255 - imgMatrix[i][j][0]
			imgMatrix[i][j][1] = 255 - imgMatrix[i][j][1]
			imgMatrix[i][j][2] = 255 - imgMatrix[i][j][2]
		}
	}

	return
}

func AdjustBrightness(src [][][]uint8 , light float64)(imgMatrix [][][]uint8 , err error) {
	imgMatrix = src

	if light <= 0{
		err = errors.New("value of light must be more than 0")
		return
	}

	height:=len(imgMatrix)
	width:=len(imgMatrix[0])
	if height == 0 || width == 0 {
		err = errors.New("The input of matrix is illegal!")
	}

	for i:=0;i<height;i++{
		for j:=0;j<width;j++{
			imgMatrix[i][j][0] = uint8(float64(imgMatrix[i][j][0])*light)
			imgMatrix[i][j][1] = uint8(float64(imgMatrix[i][j][1])*light)
			imgMatrix[i][j][2] = uint8(float64(imgMatrix[i][j][2])*light)
		}
	}

	return
}

// fuse two images(filepath) and the size of new image is as src1
func ImageFusion(src1 string , src2 string)(imgMatrix [][][]uint8 , err error) {
	imgMatrix1,err1 := Read(src1)

	if err1 != nil {
		err = err1
		return
	}

	height:=len(imgMatrix1)
	width:=len(imgMatrix1[0])

	imgMatrix2,err2 := ResizeForMatrix(src2,width,height)

	if err2 != nil {
		err = err2
		return
	}

	for i:=0;i<height;i++{
		for j:=0;j<width;j++{
			imgMatrix1[i][j][0] = uint8(float64(imgMatrix1[i][j][0])*0.5)+uint8(float64(imgMatrix2[i][j][0])*0.5)
			imgMatrix1[i][j][1] = uint8(float64(imgMatrix1[i][j][1])*0.5)+uint8(float64(imgMatrix2[i][j][1])*0.5)
			imgMatrix1[i][j][2] = uint8(float64(imgMatrix1[i][j][2])*0.5)+uint8(float64(imgMatrix1[i][j][2])*0.5)
		}
	}
	imgMatrix = imgMatrix1
	return
}

func VerticalMirror(src [][][]uint8)(imgMatrix [][][]uint8 , err error){
	imgMatrix = src

	height:=len(imgMatrix)
	width:=len(imgMatrix[0])
	if height == 0 || width == 0 {
		err = errors.New("The input of matrix is illegal!")
	}

	mirror_w:=width/2

	for i:=0;i<height;i++{
		for j:=0;j<mirror_w;j++{
			imgMatrix[i][j][0] = imgMatrix[i][width-j-1][0]
			imgMatrix[i][j][1] = imgMatrix[i][width-j-1][1]
			imgMatrix[i][j][2] = imgMatrix[i][width-j-1][2]
		}
	}

	return
}<|endoftext|>


<commit_before>// Package bamstats provides functions for computing several mapping statistics on a BAM file.
package bamstats

import (
	"errors"
	"os"
	"sync"
	"time"

	log "github.com/Sirupsen/logrus"
	"github.com/biogo/hts/bam"
	"github.com/biogo/hts/sam"
)

func init() {
	log.SetLevel(log.WarnLevel)
}

// Stats represents mapping statistics.
type Stats interface {
	Update(other Stats)
	Merge(others chan Stats)
	Collect(record *sam.Record, index *RtreeMap)
	Finalize()
}

// StatsMap is a map of Stats instances with string keys.
type StatsMap map[string]Stats

var wg sync.WaitGroup

// Merge merges instances of StatsMap
func (sm *StatsMap) Merge(stats chan StatsMap) {
	for s := range stats {
		for key, stat := range *sm {
			if otherStat, ok := s[key]; ok {
				stat.Update(otherStat)
			}
		}
	}
}

func getStatsMap(stats []Stats) StatsMap {
	m := make(StatsMap)
	for _, s := range stats {
		s.Finalize()
		switch s.(type) {
		case *GeneralStats:
			m["general"] = s
		case *CoverageStats:
			if s.(*CoverageStats).uniq {
				m["coverageUniq"] = s
			} else {
				m["coverage"] = s
			}
		}
	}
	return m
}

func worker(in chan *sam.Record, out chan StatsMap, index *RtreeMap, uniq bool) {
	defer wg.Done()
	stats := []Stats{NewGeneralStats()}
	if index != nil {
		stats = append(stats, NewCoverageStats())
		if uniq {
			cs := NewCoverageStats()
			cs.uniq = true
			stats = append(stats, cs)
		}
	}
	for record := range in {
		for _, s := range stats {
			s.Collect(record, index)
		}
	}
	log.Debug("Worker DONE!")

	out <- getStatsMap(stats)
}

func readBAM(bamFile string, index *RtreeMap, cpu int, maxBuf int, reads int, uniq bool) (chan StatsMap, error) {
	f, err := os.Open(bamFile)
	defer f.Close()
	if err != nil {
		return nil, err
	}
	br, err := bam.NewReader(f, cpu)
	defer br.Close()
	if err != nil {
		return nil, err
	}
	input := make([]chan *sam.Record, cpu)
	stats := make(chan StatsMap, cpu)
	for i := 0; i < cpu; i++ {
		wg.Add(1)
		input[i] = make(chan *sam.Record, maxBuf)
		go worker(input[i], stats, index, uniq)
	}
	c := 0
	for {
		if reads > -1 && c == reads {
			break
		}
		record, err := br.Read()
		if err != nil {
			break
		}
		input[c%cpu] <- record
		c++
	}
	for i := 0; i < cpu; i++ {
		close(input[i])
	}
	return stats, nil
}

// Process process the input BAM file and collect different mapping stats.
func Process(bamFile string, annotation string, cpu int, maxBuf int, reads int, uniq bool) (StatsMap, error) {
	var index *RtreeMap
	if bamFile == "" {
		return nil, errors.New("Please specify a BAM input file")
	}
	if annotation != "" {
		log.Infof("Creating index for %s", annotation)
		start := time.Now()
		index = CreateIndex(annotation, cpu)
		log.Infof("Index done in %v", time.Since(start))
	}
	start := time.Now()
	log.Infof("Collecting stats for %s", bamFile)
	stats, err := readBAM(bamFile, index, cpu, maxBuf, reads, uniq)
	if err != nil {
		return nil, err
	}
	go func() {
		wg.Wait()
		close(stats)
	}()
	log.Infof("Stats done in %v", time.Since(start))
	st := <-stats
	st.Merge(stats)
	return st, nil
}
<commit_msg>Add workers and caller function for using BAM file index - resolves #4<commit_after>// Package bamstats provides functions for computing several mapping statistics on a BAM file.
package bamstats

import (
	"errors"
	"io"
	"os"
	"sync"
	"time"

	log "github.com/Sirupsen/logrus"
	"github.com/biogo/hts/bam"
	"github.com/biogo/hts/bgzf"
	bgzfidx "github.com/biogo/hts/bgzf/index"
	"github.com/biogo/hts/sam"
)

func init() {
	log.SetLevel(log.WarnLevel)
}

// Stats represents mapping statistics.
type Stats interface {
	Update(other Stats)
	Merge(others chan Stats)
	Collect(record *sam.Record, index *RtreeMap)
	Finalize()
}

type IndexWorkerData struct {
	Ref   *sam.Reference
	Chunk bgzf.Chunk
}

// StatsMap is a map of Stats instances with string keys.
type StatsMap map[string]Stats

var wg sync.WaitGroup

// Merge merges instances of StatsMap
func (sm *StatsMap) Merge(stats chan StatsMap) {
	for s := range stats {
		for key, stat := range *sm {
			if otherStat, ok := s[key]; ok {
				stat.Update(otherStat)
			}
		}
	}
}

func getStatsMap(stats []Stats) StatsMap {
	m := make(StatsMap)
	for _, s := range stats {
		s.Finalize()
		switch s.(type) {
		case *GeneralStats:
			m["general"] = s
		case *CoverageStats:
			if s.(*CoverageStats).uniq {
				m["coverageUniq"] = s
			} else {
				m["coverage"] = s
			}
		}
	}
	return m
}

func workerIndex(id int, fname string, in chan IndexWorkerData, out chan StatsMap, index *RtreeMap, uniq bool) {
	defer wg.Done()
	logger := log.WithFields(log.Fields{
		"Worker": id,
	})
	logger.Debug("Starting")
	f, err := os.Open(fname)
	defer f.Close()
	if err != nil {
		panic(err)
	}
	br, err := bam.NewReader(f, 1)
	defer br.Close()
	if err != nil {
		panic(err)
	}
	stats := []Stats{NewGeneralStats()}
	if index != nil {
		stats = append(stats, NewCoverageStats())
		if uniq {
			cs := NewCoverageStats()
			cs.uniq = true
			stats = append(stats, cs)
		}
	}
	for data := range in {
		logger.WithFields(log.Fields{
			"Reference": data.Ref.Name(),
			"Length":    data.Ref.Len(),
		}).Debugf("Reading reference")
		it, err := bam.NewIterator(br, []bgzf.Chunk{data.Chunk})
		defer it.Close()
		if err != nil {
			if err != io.EOF {
				log.Println(err)
			}
			it.Close()
			panic(err)
		}
		for it.Next() {
			for _, s := range stats {
				s.Collect(it.Record(), index)
			}
		}
	}
	logger.Debug("Done")

	out <- getStatsMap(stats)
}

func worker(id int, in chan *sam.Record, out chan StatsMap, index *RtreeMap, uniq bool) {
	defer wg.Done()
	logger := log.WithFields(log.Fields{
		"worker": id,
	})
	logger.Debug("Starting")
	stats := []Stats{NewGeneralStats()}
	if index != nil {
		stats = append(stats, NewCoverageStats())
		if uniq {
			cs := NewCoverageStats()
			cs.uniq = true
			stats = append(stats, cs)
		}
	}
	for record := range in {
		for _, s := range stats {
			s.Collect(record, index)
		}
	}
	logger.Debug("Done")

	out <- getStatsMap(stats)
}

func readBAMWithIndex(bamFile string, index *RtreeMap, cpu int, maxBuf int, reads int, uniq bool) (chan StatsMap, error) {
	f, err := os.Open(bamFile)
	defer f.Close()
	if err != nil {
		return nil, err
	}
	br, err := bam.NewReader(f, 1)
	defer br.Close()
	if err != nil {
		return nil, err
	}
	log.Infof("Opening BAM index %s", bamFile+".bai")
	i, err := os.Open(bamFile + ".bai")
	defer i.Close()
	if err != nil {
		return nil, err
	}
	bai, err := bam.ReadIndex(i)
	if err != nil {
		return nil, err
	}
	h := br.Header()
	nRefs := len(h.Refs())
	stats := make(chan StatsMap, cpu)
	chunks := make(chan IndexWorkerData, cpu)
	nWorkers := cpu
	if cpu > nRefs {
		log.WithFields(log.Fields{
			"References": nRefs,
		}).Warnf("Limiting the number of workers to the number of BAM references")
		nWorkers = nRefs
	}
	for i := 0; i < nWorkers; i++ {
		wg.Add(1)
		go workerIndex(i+1, bamFile, chunks, stats, index, uniq)
	}
	for _, ref := range h.Refs() {
		refChunks, _ := bai.Chunks(ref, 0, ref.Len())
		if err != nil {
			if err != io.EOF && err != bgzfidx.ErrInvalid {
				log.Error(err)
			}
			return nil, err
		}
		if len(refChunks) > 0 {
			if len(refChunks) > 1 {
				log.Debugf("%v: %v chunks", ref.Name(), len(refChunks))
			}
			for _, chunk := range refChunks {
				chunks <- IndexWorkerData{ref, chunk}
			}
		}
	}
	close(chunks)

	return stats, nil
}

func readBAM(bamFile string, index *RtreeMap, cpu int, maxBuf int, reads int, uniq bool) (chan StatsMap, error) {
	f, err := os.Open(bamFile)
	defer f.Close()
	if err != nil {
		return nil, err
	}
	br, err := bam.NewReader(f, cpu)
	defer br.Close()
	if err != nil {
		return nil, err
	}
	input := make([]chan *sam.Record, cpu)
	stats := make(chan StatsMap, cpu)
	for i := 0; i < cpu; i++ {
		wg.Add(1)
		input[i] = make(chan *sam.Record, maxBuf)
		go worker(i+1, input[i], stats, index, uniq)
	}
	c := 0
	for {
		if reads > -1 && c == reads {
			break
		}
		record, err := br.Read()
		if err != nil {
			break
		}
		input[c%cpu] <- record
		c++
	}
	for i := 0; i < cpu; i++ {
		close(input[i])
	}
	return stats, nil
}

// Process process the input BAM file and collect different mapping stats.
func Process(bamFile string, annotation string, cpu int, maxBuf int, reads int, uniq bool) (StatsMap, error) {
	var index *RtreeMap
	if bamFile == "" {
		return nil, errors.New("Please specify a BAM input file")
	}
	if annotation != "" {
		log.Infof("Creating index for %s", annotation)
		start := time.Now()
		index = CreateIndex(annotation, cpu)
		log.Infof("Index done in %v", time.Since(start))
	}
	start := time.Now()
	log.Infof("Collecting stats for %s", bamFile)
	process := readBAMWithIndex
	if _, err := os.Stat(bamFile + ".bai"); os.IsNotExist(err) || cpu == 1 {
		log.Warning("Not using BAM index")
		process = readBAM
	}
	stats, err := process(bamFile, index, cpu, maxBuf, reads, uniq)
	if err != nil {
		return nil, err
	}
	go func() {
		wg.Wait()
		close(stats)
		log.Infof("Stats done in %v", time.Since(start))
	}()
	st := <-stats
	st.Merge(stats)
	return st, nil
}
<|endoftext|>


<commit_before>package command

import (
	"os"
	"strconv"
	"runtime"
	"path/filepath"
	"github.com/webdevops/go-stubfilegenerator"
	"github.com/remeh/sizedwaitgroup"
)

type Typo3Stubs struct {
	Options MysqlCommonOptions `group:"common"`
	Positional struct {
		Database string `description:"Database" required:"1"`
	} `positional-args:"true"`
	RootPath    string  `long:"path"            description:"TYPO3 root path"`
	Force  bool   `short:"f"  long:"force"      description:"Overwrite existing files"`
}

type storage struct {
	Uid string
	Name string
	Path string
}

type storageFile struct {
	Uid string
	Path string
	RelPath string
	AbsPath string
	ImageWidth string
	ImageHeight string
}

func (conf *Typo3Stubs) Execute(args []string) error {
	Logger.Main("Starting TYPO3 fileadmin stub generator")
	conf.Options.Init()

	rootPath, err := conf.GetTypo3Root()
	if err != nil {
		return err
	}

	sql := `SELECT uid,
                   name,
                   ExtractValue(configuration, '//field[@index=\'basePath\']/value/text()') as storagepath
              FROM sys_file_storage
             WHERE deleted = 0
              AND driver = 'local'`
	result := conf.Options.ExecQuery(conf.Positional.Database, sql)

	for _, row := range result.Row {
		rowList := row.GetList()
		storage := storage{rowList["uid"], rowList["name"], rowList["storagepath"]}
		err := conf.processStorage(storage, rootPath)
		if err != nil {
			return err
		}
	}

	return nil
}

func (conf *Typo3Stubs) processStorage(storage storage, rootPath string) error {
	Logger.Step("processing storage \"%v\" (path: %v)", storage.Name, storage.Path)
	stubgen := stubfilegenerator.NewStubGenerator()
	err := stubgen.SetRootPath(filepath.Join(rootPath, storage.Path))
	if err != nil {
		return err
	}

	stubgen.Image.Text = append(stubgen.Image.Text, "Size: %IMAGE_WIDTH% * %IMAGE_HEIGHT%")

	if conf.Force {
		stubgen.Overwrite = true
	}

	sql := `SELECT f.uid,
                   f.identifier,
                   fm.width as meta_width,
                   fm.height as meta_height
              FROM sys_file f
                   LEFT JOIN sys_file_metadata fm
                     ON fm.file = f.uid
                    AND fm.t3ver_oid = 0
              WHERE f.storage = ` + storage.Uid;
	result := conf.Options.ExecQuery(conf.Positional.Database, sql)

	swg := sizedwaitgroup.New(runtime.GOMAXPROCS(0) * 10)
	for _, row := range result.Row {
		row := row.GetList()

		swg.Add()
		go func(row map[string]string, stubgen stubfilegenerator.StubGenerator) {
			defer swg.Done()

			file := storageFile{}
			file.ImageWidth = "800"
			file.ImageHeight = "400"

			switch len(row) {
			case 4:
				file.ImageWidth = row["meta_width"]
				file.ImageHeight = row["meta_height"]
				fallthrough
			case 2:
				file.Uid = row["uid"]
				file.Path = row["identifier"]
			}

			Logger.Item(file.Path)
			stubgen.TemplateVariables["PATH"] = file.Path
			stubgen.TemplateVariables["IMAGE_WIDTH"] = file.ImageWidth
			stubgen.TemplateVariables["IMAGE_HEIGHT"] = file.ImageHeight
			stubgen.Image.Width, _ = strconv.Atoi(file.ImageWidth)
			stubgen.Image.Height, _ = strconv.Atoi(file.ImageHeight)
			stubgen.Generate(file.Path)
		} (row, stubgen.Clone())
		swg.Wait()
	}

	return nil
}

func (conf *Typo3Stubs) GetTypo3Root() (string, error) {
	if conf.RootPath == "" {
		// use current working dir as root path
		rootPath, err := os.Getwd()
		if err != nil {
			return "", err
		}
		conf.RootPath = rootPath
	}

	return conf.RootPath, nil
}

<commit_msg>Fix text image path in typo3:stubs<commit_after>package command

import (
	"os"
	"strconv"
	"runtime"
	"path/filepath"
	"github.com/webdevops/go-stubfilegenerator"
	"github.com/remeh/sizedwaitgroup"
)

type Typo3Stubs struct {
	Options MysqlCommonOptions `group:"common"`
	Positional struct {
		Database string `description:"Database" required:"1"`
	} `positional-args:"true"`
	RootPath    string  `long:"path"            description:"TYPO3 root path"`
	Force  bool   `short:"f"  long:"force"      description:"Overwrite existing files"`
}

type storage struct {
	Uid string
	Name string
	Path string
}

type storageFile struct {
	Uid string
	Path string
	RelPath string
	AbsPath string
	ImageWidth string
	ImageHeight string
}

func (conf *Typo3Stubs) Execute(args []string) error {
	Logger.Main("Starting TYPO3 fileadmin stub generator")
	conf.Options.Init()

	rootPath, err := conf.GetTypo3Root()
	if err != nil {
		return err
	}

	sql := `SELECT uid,
                   name,
                   ExtractValue(configuration, '//field[@index=\'basePath\']/value/text()') as storagepath
              FROM sys_file_storage
             WHERE deleted = 0
              AND driver = 'local'`
	result := conf.Options.ExecQuery(conf.Positional.Database, sql)

	for _, row := range result.Row {
		rowList := row.GetList()
		storage := storage{rowList["uid"], rowList["name"], rowList["storagepath"]}
		err := conf.processStorage(storage, rootPath)
		if err != nil {
			return err
		}
	}

	return nil
}

func (conf *Typo3Stubs) processStorage(storage storage, rootPath string) error {
	Logger.Step("processing storage \"%v\" (path: %v)", storage.Name, storage.Path)
	stubgen := stubfilegenerator.NewStubGenerator()
	err := stubgen.SetRootPath(filepath.Join(rootPath, storage.Path))
	if err != nil {
		return err
	}

	stubgen.Image.Text = append(stubgen.Image.Text, "Size: %IMAGE_WIDTH% * %IMAGE_HEIGHT%")

	if conf.Force {
		stubgen.Overwrite = true
	}

	sql := `SELECT f.uid,
                   f.identifier,
                   fm.width as meta_width,
                   fm.height as meta_height
              FROM sys_file f
                   LEFT JOIN sys_file_metadata fm
                     ON fm.file = f.uid
                    AND fm.t3ver_oid = 0
              WHERE f.storage = ` + storage.Uid;
	result := conf.Options.ExecQuery(conf.Positional.Database, sql)

	swg := sizedwaitgroup.New(runtime.GOMAXPROCS(0) * 10)
	for _, row := range result.Row {
		row := row.GetList()

		swg.Add()
		go func(row map[string]string, stubgen stubfilegenerator.StubGenerator) {
			defer swg.Done()

			file := storageFile{}
			file.ImageWidth = "800"
			file.ImageHeight = "400"

			switch len(row) {
			case 4:
				file.ImageWidth = row["meta_width"]
				file.ImageHeight = row["meta_height"]
				fallthrough
			case 2:
				file.Uid = row["uid"]
				file.Path = row["identifier"]
			}

			Logger.Item(file.Path)
			stubgen.TemplateVariables["PATH"] = filepath.Join(storage.Path, file.Path)
			stubgen.TemplateVariables["IMAGE_WIDTH"] = file.ImageWidth
			stubgen.TemplateVariables["IMAGE_HEIGHT"] = file.ImageHeight
			stubgen.Image.Width, _ = strconv.Atoi(file.ImageWidth)
			stubgen.Image.Height, _ = strconv.Atoi(file.ImageHeight)
			stubgen.Generate(file.Path)
		} (row, stubgen.Clone())
		swg.Wait()
	}

	return nil
}

func (conf *Typo3Stubs) GetTypo3Root() (string, error) {
	if conf.RootPath == "" {
		// use current working dir as root path
		rootPath, err := os.Getwd()
		if err != nil {
			return "", err
		}
		conf.RootPath = rootPath
	}

	return conf.RootPath, nil
}

<|endoftext|>