<commit_before>package model

import (
	"fmt"
	"github.com/mattn/go-sqlite3"
	"github.com/stretchr/testify/assert"
	"testing"
)

func fixtureEmptySlave() *Slave {
	return &Slave{
		Hostname:             "host1",
		Port:                 1,
		MongodPortRangeBegin: 2,
		MongodPortRangeEnd:   3,
		PersistentStorage:    true,
		Mongods:              []*Mongod{},
		ConfiguredState:      SlaveStateActive,
	}
}

func fixtureEmptyMongod() *Mongod {
	return &Mongod{
		Port:        8080,
		ReplSetName: "repl1",
	}
}

func fixtureEmptyMongodState() MongodState {
	return MongodState{}
}

func fixtureEmptyRiskGroup() *RiskGroup {
	return &RiskGroup{
		Name:   "rg1",
		Slaves: []*Slave{},
	}
}

func fixtureEmptyReplicaSet() *ReplicaSet {
	return &ReplicaSet{
		Name: "repl1",
		PersistentMemberCount:           1,
		VolatileMemberCount:             2,
		ConfigureAsShardingConfigServer: false,
	}
}

func fixtureEmptyProblem() *Problem {
	return &Problem{
		Description: "Test",
	}
}

////////////////////////////////////////////////////////////////////////////////

func TestCanInitializeDB(t *testing.T) {
	_, _, err := InitializeTestDB()
	assert.NoError(t, err)
}

/*
 This elaborate test demonstrates how resolving an association works in gorm.
 Check the assertions to learn about the behavior of gorm.
*/
func TestRelationshipMongodParentSlave(t *testing.T) {

	db, _, _ := InitializeTestDB()
	tx := db.Begin()
	defer tx.Rollback()

	s := fixtureEmptySlave()

	tx.Create(s)

	m := fixtureEmptyMongod()
	m.ParentSlave = s

	tx.Create(m)

	assert.Equal(t, m.ParentSlaveID, s.ID)

	assert.Equal(t, s.Mongods, []*Mongod{})

	var sdb Slave

	// Check what happens when just SELECTing the slave
	err := tx.First(&sdb).Error

	assert.NoError(t, err)
	assert.Nil(t, sdb.Mongods)

	// Now resolve the slave->mongod 1:n association
	err = tx.Model(&sdb).Related(&sdb.Mongods, "Mongods").Error

	assert.NoError(t, err)
	assert.Equal(t, len(sdb.Mongods), 1)
	assert.Equal(t, sdb.Mongods[0].ReplSetName, m.ReplSetName)
	assert.Zero(t, sdb.Mongods[0].ParentSlave)
	assert.Equal(t, sdb.Mongods[0].ParentSlaveID, s.ID)

	// Now resolve the mongod->(parent)slave relation
	parentSlave := &Slave{}
	err = tx.Model(&sdb.Mongods[0]).Related(parentSlave, "ParentSlave").Error
	assert.NoError(t, err)
	assert.NotZero(t, parentSlave)
	assert.Equal(t, s.ID, parentSlave.ID)

}

// Test RiskGroup Slave relationship
func TestRiskGroupSlaveRelationship(t *testing.T) {

	db, _, _ := InitializeTestDB()
	tx := db.Begin()
	defer tx.Rollback()

	s := fixtureEmptySlave()
	r := fixtureEmptyRiskGroup()
	r.Slaves = []*Slave{s}

	err := tx.Create(&r).Error
	assert.NoError(t, err)

	var rdb RiskGroup

	err = tx.First(&rdb).Error

	assert.NoError(t, err)
	assert.Zero(t, rdb.Slaves)

	err = tx.Model(&rdb).Related(&rdb.Slaves, "Slaves").Error
	assert.NoError(t, err)
	assert.NotZero(t, rdb.Slaves)
	assert.Equal(t, len(rdb.Slaves), 1)
	assert.Equal(t, rdb.Slaves[0].ID, s.ID)

}

// Test ReplicaSet - Mongod Relationship
func TestReplicaSetMongodRelationship(t *testing.T) {

	db, _, _ := InitializeTestDB()
	tx := db.Begin()
	defer tx.Rollback()

	r := fixtureEmptyReplicaSet()
	m := fixtureEmptyMongod()
	r.Mongods = []*Mongod{m}

	err := tx.Create(&r).Error
	assert.NoError(t, err)

	var rdb ReplicaSet

	err = tx.First(&rdb).Error

	assert.NoError(t, err)
	assert.Zero(t, rdb.Mongods)

	err = tx.Model(&rdb).Related(&rdb.Mongods, "Mongods").Error
	assert.NoError(t, err)
	assert.NotZero(t, rdb.Mongods)
	assert.Equal(t, len(rdb.Mongods), 1)
	assert.Equal(t, rdb.Mongods[0].ID, m.ID)

}

// Test Mongod - MongodState relationship
func TestMongodMongodStateRelationship(t *testing.T) {

	db, _, _ := InitializeTestDB()
	tx := db.Begin()
	defer tx.Rollback()

	m := fixtureEmptyMongod()

	o := MongodState{
		IsShardingConfigServer: false,
		ExecutionState:         MongodExecutionStateNotRunning,
		ReplicaSetMembers:      []ReplicaSetMember{},
	}

	d := MongodState{
		IsShardingConfigServer: false,
		ExecutionState:         MongodExecutionStateRunning,
		ReplicaSetMembers:      []ReplicaSetMember{},
	}

	assert.NoError(t, tx.Create(m).Error)
	o.ParentMongodID = m.ID
	d.ParentMongodID = m.ID
	assert.NoError(t, tx.Create(&o).Error)
	assert.NoError(t, tx.Create(&d).Error)
	assert.NoError(t, tx.Model(&m).Update("DesiredStateID", d.ID).Error)
	assert.NoError(t, tx.Model(&m).Update("ObservedStateID", o.ID).Error)

	var mdb Mongod

	// Observed
	assert.NoError(t, tx.First(&mdb).Error)
	assert.Zero(t, mdb.ObservedState)

	assert.NoError(t, tx.Model(&mdb).Related(&mdb.ObservedState, "ObservedState").Error)
	assert.NotZero(t, mdb.ObservedState)
	assert.Equal(t, mdb.ObservedState.ExecutionState, MongodExecutionStateNotRunning)

	assert.NoError(t, tx.Model(&mdb).Related(&mdb.DesiredState, "DesiredState").Error)
	assert.NotZero(t, mdb.DesiredState)
	assert.Equal(t, mdb.DesiredState.ExecutionState, MongodExecutionStateRunning)

}

// Test MongodState - ReplicaSetMember relationship
func TestMongodStateReplicaSetMembersRelationship(t *testing.T) {
	db, _, _ := InitializeTestDB()
	tx := db.Begin()
	defer tx.Rollback()

	m := ReplicaSetMember{Hostname: "h1"}

	s := MongodState{ReplicaSetMembers: []ReplicaSetMember{m}}

	assert.NoError(t, tx.Create(&m).Error)
	s.ParentMongodID = m.ID
	assert.NoError(t, tx.Create(&s).Error)
	assert.NoError(t, tx.Model(&m).Update("DesiredStateID", s.ID).Error)

	var sdb MongodState

	assert.NoError(t, tx.First(&sdb).Error)
	assert.Zero(t, sdb.ReplicaSetMembers)

	assert.NoError(t, tx.Model(&sdb).Related(&sdb.ReplicaSetMembers, "ReplicaSetMembers").Error)
	assert.NotZero(t, sdb.ReplicaSetMembers)
	assert.Equal(t, len(sdb.ReplicaSetMembers), 1)
	assert.Equal(t, sdb.ReplicaSetMembers[0].Hostname, m.Hostname)

}

func TestDeleteBehavior(t *testing.T) {

	db, _, _ := InitializeTestDB()
	tx := db.Begin()
	defer tx.Rollback()

	m := fixtureEmptyMongod()
	m.ID = 1000

	// Create it
	tx.Create(&m)

	var mdb Mongod

	// Read it once
	d := tx.First(&mdb)

	assert.NoError(t, d.Error)
	assert.Equal(t, mdb.ID, m.ID)

	// Destroy it once, by ID
	d = tx.Delete(&Mongod{ID: 1000})

	assert.NoError(t, d.Error)
	assert.EqualValues(t, 1, d.RowsAffected)

	// Destroy it a second time.
	// No Error will occur, have to check RowsAffected if we deleted something
	d = tx.Delete(&Mongod{ID: 1000})

	assert.NoError(t, d.Error)
	assert.EqualValues(t, 0, d.RowsAffected)

}

func TestGormFirstBehavior(t *testing.T) {
	db, _, _ := InitializeTestDB()
	tx := db.Begin()
	defer tx.Rollback()

	var m Mongod
	assert.Error(t, tx.First(&m).Error)
}

func TestGormFindBehavior(t *testing.T) {
	db, _, _ := InitializeTestDB()
	tx := db.Begin()
	defer tx.Rollback()

	var ms []Mongod
	d := tx.Find(&ms)

	assert.NoError(t, d.Error)
	assert.EqualValues(t, 0, d.RowsAffected) // RowsAffected does NOT indicate "nothing found"!!!!
	assert.Equal(t, 0, len(ms))              // Use this instead

}

func TestGormTransactions(t *testing.T) {
	db, _, _ := InitializeTestDB()

	//Create a slave
	tx0 := db.Begin()
	m := fixtureEmptySlave()
	m.ID = 1
	m.Hostname = "baz"
	m.Port = 5
	assert.NoError(t, tx0.Create(&m).Error)
	assert.NoError(t, tx0.Commit().Error)

	fmt.Println("Insert slaves from tx0 done and committed")

	//Modify same slave in two transactions
	tx1 := db.Begin()
	fmt.Println("begin tx1")
	assert.NoError(t, tx1.First(&Slave{}, 1).Update("hostname", "foo").Error)
	assert.NoError(t, tx1.First(&Slave{}, 1).Update("port", 15).Error)
	fmt.Println("Update slave 1 from tx1 done")

	tx2 := db.Begin()
	fmt.Println("begin tx2")

	//Should be able to read slave 1 and see old state
	var slaveReadX Slave
	assert.NoError(t, tx2.First(&slaveReadX, 1).Error)
	assert.Equal(t, "baz", slaveReadX.Hostname)
	fmt.Println("Read slave 1 from tx2 done")

	err := tx2.First(&Slave{}, 1).Update("hostname", "bar").Error
	assert.Error(t, err)
	driverErr, ok := err.(sqlite3.Error)
	assert.True(t, ok)
	assert.Equal(t, sqlite3.ErrBusy, driverErr.Code) //https://www.sqlite.org/rescode.html#busy
	fmt.Println("Update slave 1 from tx2 done")

	//Commit Tx2
	assert.NoError(t, tx2.Commit().Error)
	fmt.Println("tx2 done")
	//Commit Tx1
	assert.NoError(t, tx1.Commit().Error)
	fmt.Println("tx1 done")

	var slave Slave
	tx := db.Begin()
	tx.First(&slave, 1)
	tx.Rollback()
	assert.Equal(t, "foo", slave.Hostname)
}

func TestForeignKeyChecking(t *testing.T) {
	db, _, _ := InitializeTestDB()

	tx0 := db.Begin()

	var foreignKeys int
	tx0.Raw("PRAGMA foreign_keys;").Row().Scan(&foreignKeys)
	assert.EqualValues(t, 1, foreignKeys)

	res := tx0.Create(&Slave{
		Hostname:             "host1",
		Port:                 1,
		MongodPortRangeBegin: 2,
		MongodPortRangeEnd:   3,
		PersistentStorage:    true,
		Mongods:              []*Mongod{},
		ConfiguredState:      SlaveStateActive,
		RiskGroupID:          NullIntValue(99),
	})
	assert.NoError(t, res.Error)
	assert.Error(t, tx0.Commit().Error)
}

func TestCascade(t *testing.T) {
	db, _, _ := InitializeTestDB()

	tx0 := db.Begin()
	assert.NoError(t, tx0.Exec("CREATE TABLE foo(id int primary key);").Error)
	assert.NoError(t, tx0.Exec("CREATE TABLE bar3(foreignKey int null references foo(id) on delete cascade deferrable initially deferred);").Error)
	assert.NoError(t, tx0.Commit().Error)

	tx1 := db.Begin()
	assert.NoError(t, tx1.Exec("INSERT INTO foo VALUES(1); INSERT INTO bar3 VALUES(1);").Error)
	assert.NoError(t, tx1.Commit().Error)

	tx3 := db.Begin()
	var count int
	tx3.Raw("SELECT count(*) FROM bar3;").Row().Scan(&count)
	assert.EqualValues(t, 1, count)
	tx3.Commit()

	tx2 := db.Begin()
	tx2.Exec("DELETE FROM foo WHERE id = 1;")
	assert.NoError(t, tx2.Commit().Error)

	tx3 = db.Begin()
	tx3.Raw("SELECT count(*) FROM bar3;").Row().Scan(&count)
	assert.EqualValues(t, 0, count)
	tx3.Commit()
}

func TestCascadeSlaves(t *testing.T) {
	db, _, _ := InitializeTestDB()
	tx := db.Begin()

	r := fixtureEmptyRiskGroup()
	assert.NoError(t, tx.Create(r).Error)

	rs := fixtureEmptyReplicaSet()
	assert.NoError(t, tx.Create(rs).Error)

	s := fixtureEmptySlave()
	s.RiskGroupID = NullIntValue(r.ID)
	assert.NoError(t, tx.Create(s).Error)

	ds := fixtureEmptyMongodState()
	assert.NoError(t, tx.Create(&ds).Error)

	m := fixtureEmptyMongod()
	m.ParentSlaveID = s.ID
	m.ReplicaSetID = rs.ID
	m.DesiredStateID = ds.ID
	assert.NoError(t, tx.Create(m).Error)
	assert.NoError(t, tx.Model(&ds).Update("ParentMongodID", m.ID).Error)

	p := fixtureEmptyProblem()
	p.SlaveID = NullIntValue(s.ID)
	assert.NoError(t, tx.Create(p).Error)

	assert.NoError(t, tx.Commit().Error)

	tx1 := db.Begin()
	tx1.Delete(&Mongod{}, m.ID)
	assert.NoError(t, tx1.Commit().Error)

	tx2 := db.Begin()
	tx2.Delete(&Slave{}, s.ID)
	assert.NoError(t, tx2.Commit().Error)

	tx3 := db.Begin()
	assert.True(t, tx3.First(&Problem{}, p.ID).RecordNotFound())
	tx3.Rollback()
}

// Test case demonstrating how to do overwrites
func TestObservationErrorOverwriteBehavior(t *testing.T) {

	// Assume a situation where Slave already has an ObservationError
	db, _, _ := InitializeTestDB()
	tx := db.Begin()

	o1 := MSPError{
		Identifier: "id1",
	}

	assert.NoError(t, tx.Create(&o1).Error)

	s := Slave{
		Hostname:             "s1",
		Port:                 1,
		MongodPortRangeBegin: 1,
		MongodPortRangeEnd:   2,
		ObservationErrorID:   NullIntValue(o1.ID),
	}

	assert.NoError(t, tx.Create(&s).Error)

	var countBeforeUpdate int64
	assert.NoError(t, tx.Model(&MSPError{}).Count(&countBeforeUpdate).Error)
	assert.EqualValues(t, 1, countBeforeUpdate)

	// Now we attempt to observe Slave again and want to update the value pointed to by slave
	// We could create, update, delete
	// Or we could use the following hack to just UPDATE all values of the existing observation,
	//    saving us an annoying DELETE
	o2 := MSPError{
		Identifier: "id2",
	}
	o2.ID = o1.ID
	tx.Save(&o2)
	// That was it

	var countAfterUpdate int64
	assert.NoError(t, tx.Model(&MSPError{}).Count(&countAfterUpdate).Error)

	assert.EqualValues(t, 1, countAfterUpdate, "updates should remove the row previously referenced in the overwritten column")

}
<commit_msg>DEL: tests/model: remove sqlite3 specific tests<commit_after>package model

import (
	"github.com/stretchr/testify/assert"
	"testing"
)

func fixtureEmptySlave() *Slave {
	return &Slave{
		Hostname:             "host1",
		Port:                 1,
		MongodPortRangeBegin: 2,
		MongodPortRangeEnd:   3,
		PersistentStorage:    true,
		Mongods:              []*Mongod{},
		ConfiguredState:      SlaveStateActive,
	}
}

func fixtureEmptyMongod() *Mongod {
	return &Mongod{
		Port:        8080,
		ReplSetName: "repl1",
	}
}

func fixtureEmptyMongodState() MongodState {
	return MongodState{}
}

func fixtureEmptyRiskGroup() *RiskGroup {
	return &RiskGroup{
		Name:   "rg1",
		Slaves: []*Slave{},
	}
}

func fixtureEmptyReplicaSet() *ReplicaSet {
	return &ReplicaSet{
		Name: "repl1",
		PersistentMemberCount:           1,
		VolatileMemberCount:             2,
		ConfigureAsShardingConfigServer: false,
	}
}

func fixtureEmptyProblem() *Problem {
	return &Problem{
		Description: "Test",
	}
}

////////////////////////////////////////////////////////////////////////////////

func TestCanInitializeDB(t *testing.T) {
	_, _, err := InitializeTestDB()
	assert.NoError(t, err)
}

/*
 This elaborate test demonstrates how resolving an association works in gorm.
 Check the assertions to learn about the behavior of gorm.
*/
func TestRelationshipMongodParentSlave(t *testing.T) {

	db, _, _ := InitializeTestDB()
	tx := db.Begin()
	defer tx.Rollback()

	s := fixtureEmptySlave()

	tx.Create(s)

	m := fixtureEmptyMongod()
	m.ParentSlave = s

	tx.Create(m)

	assert.Equal(t, m.ParentSlaveID, s.ID)

	assert.Equal(t, s.Mongods, []*Mongod{})

	var sdb Slave

	// Check what happens when just SELECTing the slave
	err := tx.First(&sdb).Error

	assert.NoError(t, err)
	assert.Nil(t, sdb.Mongods)

	// Now resolve the slave->mongod 1:n association
	err = tx.Model(&sdb).Related(&sdb.Mongods, "Mongods").Error

	assert.NoError(t, err)
	assert.Equal(t, len(sdb.Mongods), 1)
	assert.Equal(t, sdb.Mongods[0].ReplSetName, m.ReplSetName)
	assert.Zero(t, sdb.Mongods[0].ParentSlave)
	assert.Equal(t, sdb.Mongods[0].ParentSlaveID, s.ID)

	// Now resolve the mongod->(parent)slave relation
	parentSlave := &Slave{}
	err = tx.Model(&sdb.Mongods[0]).Related(parentSlave, "ParentSlave").Error
	assert.NoError(t, err)
	assert.NotZero(t, parentSlave)
	assert.Equal(t, s.ID, parentSlave.ID)

}

// Test RiskGroup Slave relationship
func TestRiskGroupSlaveRelationship(t *testing.T) {

	db, _, _ := InitializeTestDB()
	tx := db.Begin()
	defer tx.Rollback()

	s := fixtureEmptySlave()
	r := fixtureEmptyRiskGroup()
	r.Slaves = []*Slave{s}

	err := tx.Create(&r).Error
	assert.NoError(t, err)

	var rdb RiskGroup

	err = tx.First(&rdb).Error

	assert.NoError(t, err)
	assert.Zero(t, rdb.Slaves)

	err = tx.Model(&rdb).Related(&rdb.Slaves, "Slaves").Error
	assert.NoError(t, err)
	assert.NotZero(t, rdb.Slaves)
	assert.Equal(t, len(rdb.Slaves), 1)
	assert.Equal(t, rdb.Slaves[0].ID, s.ID)

}

// Test ReplicaSet - Mongod Relationship
func TestReplicaSetMongodRelationship(t *testing.T) {

	db, _, _ := InitializeTestDB()
	tx := db.Begin()
	defer tx.Rollback()

	r := fixtureEmptyReplicaSet()
	m := fixtureEmptyMongod()
	r.Mongods = []*Mongod{m}

	err := tx.Create(&r).Error
	assert.NoError(t, err)

	var rdb ReplicaSet

	err = tx.First(&rdb).Error

	assert.NoError(t, err)
	assert.Zero(t, rdb.Mongods)

	err = tx.Model(&rdb).Related(&rdb.Mongods, "Mongods").Error
	assert.NoError(t, err)
	assert.NotZero(t, rdb.Mongods)
	assert.Equal(t, len(rdb.Mongods), 1)
	assert.Equal(t, rdb.Mongods[0].ID, m.ID)

}

// Test Mongod - MongodState relationship
func TestMongodMongodStateRelationship(t *testing.T) {

	db, _, _ := InitializeTestDB()
	tx := db.Begin()
	defer tx.Rollback()

	m := fixtureEmptyMongod()

	o := MongodState{
		IsShardingConfigServer: false,
		ExecutionState:         MongodExecutionStateNotRunning,
		ReplicaSetMembers:      []ReplicaSetMember{},
	}

	d := MongodState{
		IsShardingConfigServer: false,
		ExecutionState:         MongodExecutionStateRunning,
		ReplicaSetMembers:      []ReplicaSetMember{},
	}

	assert.NoError(t, tx.Create(m).Error)
	o.ParentMongodID = m.ID
	d.ParentMongodID = m.ID
	assert.NoError(t, tx.Create(&o).Error)
	assert.NoError(t, tx.Create(&d).Error)
	assert.NoError(t, tx.Model(&m).Update("DesiredStateID", d.ID).Error)
	assert.NoError(t, tx.Model(&m).Update("ObservedStateID", o.ID).Error)

	var mdb Mongod

	// Observed
	assert.NoError(t, tx.First(&mdb).Error)
	assert.Zero(t, mdb.ObservedState)

	assert.NoError(t, tx.Model(&mdb).Related(&mdb.ObservedState, "ObservedState").Error)
	assert.NotZero(t, mdb.ObservedState)
	assert.Equal(t, mdb.ObservedState.ExecutionState, MongodExecutionStateNotRunning)

	assert.NoError(t, tx.Model(&mdb).Related(&mdb.DesiredState, "DesiredState").Error)
	assert.NotZero(t, mdb.DesiredState)
	assert.Equal(t, mdb.DesiredState.ExecutionState, MongodExecutionStateRunning)

}

// Test MongodState - ReplicaSetMember relationship
func TestMongodStateReplicaSetMembersRelationship(t *testing.T) {
	db, _, _ := InitializeTestDB()
	tx := db.Begin()
	defer tx.Rollback()

	m := ReplicaSetMember{Hostname: "h1"}

	s := MongodState{ReplicaSetMembers: []ReplicaSetMember{m}}

	assert.NoError(t, tx.Create(&m).Error)
	s.ParentMongodID = m.ID
	assert.NoError(t, tx.Create(&s).Error)
	assert.NoError(t, tx.Model(&m).Update("DesiredStateID", s.ID).Error)

	var sdb MongodState

	assert.NoError(t, tx.First(&sdb).Error)
	assert.Zero(t, sdb.ReplicaSetMembers)

	assert.NoError(t, tx.Model(&sdb).Related(&sdb.ReplicaSetMembers, "ReplicaSetMembers").Error)
	assert.NotZero(t, sdb.ReplicaSetMembers)
	assert.Equal(t, len(sdb.ReplicaSetMembers), 1)
	assert.Equal(t, sdb.ReplicaSetMembers[0].Hostname, m.Hostname)

}

func TestDeleteBehavior(t *testing.T) {

	db, _, _ := InitializeTestDB()
	tx := db.Begin()
	defer tx.Rollback()

	m := fixtureEmptyMongod()
	m.ID = 1000

	// Create it
	tx.Create(&m)

	var mdb Mongod

	// Read it once
	d := tx.First(&mdb)

	assert.NoError(t, d.Error)
	assert.Equal(t, mdb.ID, m.ID)

	// Destroy it once, by ID
	d = tx.Delete(&Mongod{ID: 1000})

	assert.NoError(t, d.Error)
	assert.EqualValues(t, 1, d.RowsAffected)

	// Destroy it a second time.
	// No Error will occur, have to check RowsAffected if we deleted something
	d = tx.Delete(&Mongod{ID: 1000})

	assert.NoError(t, d.Error)
	assert.EqualValues(t, 0, d.RowsAffected)

}

func TestGormFirstBehavior(t *testing.T) {
	db, _, _ := InitializeTestDB()
	tx := db.Begin()
	defer tx.Rollback()

	var m Mongod
	assert.Error(t, tx.First(&m).Error)
}

func TestGormFindBehavior(t *testing.T) {
	db, _, _ := InitializeTestDB()
	tx := db.Begin()
	defer tx.Rollback()

	var ms []Mongod
	d := tx.Find(&ms)

	assert.NoError(t, d.Error)
	assert.EqualValues(t, 0, d.RowsAffected) // RowsAffected does NOT indicate "nothing found"!!!!
	assert.Equal(t, 0, len(ms))              // Use this instead

}

func TestCascade(t *testing.T) {
	db, _, _ := InitializeTestDB()

	tx0 := db.Begin()
	assert.NoError(t, tx0.Exec("CREATE TABLE foo(id int primary key);").Error)
	assert.NoError(t, tx0.Exec("CREATE TABLE bar3(foreignKey int null references foo(id) on delete cascade deferrable initially deferred);").Error)
	assert.NoError(t, tx0.Commit().Error)

	tx1 := db.Begin()
	assert.NoError(t, tx1.Exec("INSERT INTO foo VALUES(1); INSERT INTO bar3 VALUES(1);").Error)
	assert.NoError(t, tx1.Commit().Error)

	tx3 := db.Begin()
	var count int
	tx3.Raw("SELECT count(*) FROM bar3;").Row().Scan(&count)
	assert.EqualValues(t, 1, count)
	tx3.Commit()

	tx2 := db.Begin()
	tx2.Exec("DELETE FROM foo WHERE id = 1;")
	assert.NoError(t, tx2.Commit().Error)

	tx3 = db.Begin()
	tx3.Raw("SELECT count(*) FROM bar3;").Row().Scan(&count)
	assert.EqualValues(t, 0, count)
	tx3.Commit()
}

func TestCascadeSlaves(t *testing.T) {
	db, _, _ := InitializeTestDB()
	tx := db.Begin()

	r := fixtureEmptyRiskGroup()
	assert.NoError(t, tx.Create(r).Error)

	rs := fixtureEmptyReplicaSet()
	assert.NoError(t, tx.Create(rs).Error)

	s := fixtureEmptySlave()
	s.RiskGroupID = NullIntValue(r.ID)
	assert.NoError(t, tx.Create(s).Error)

	ds := fixtureEmptyMongodState()
	assert.NoError(t, tx.Create(&ds).Error)

	m := fixtureEmptyMongod()
	m.ParentSlaveID = s.ID
	m.ReplicaSetID = rs.ID
	m.DesiredStateID = ds.ID
	assert.NoError(t, tx.Create(m).Error)
	assert.NoError(t, tx.Model(&ds).Update("ParentMongodID", m.ID).Error)

	p := fixtureEmptyProblem()
	p.SlaveID = NullIntValue(s.ID)
	assert.NoError(t, tx.Create(p).Error)

	assert.NoError(t, tx.Commit().Error)

	tx1 := db.Begin()
	tx1.Delete(&Mongod{}, m.ID)
	assert.NoError(t, tx1.Commit().Error)

	tx2 := db.Begin()
	tx2.Delete(&Slave{}, s.ID)
	assert.NoError(t, tx2.Commit().Error)

	tx3 := db.Begin()
	assert.True(t, tx3.First(&Problem{}, p.ID).RecordNotFound())
	tx3.Rollback()
}

// Test case demonstrating how to do overwrites
func TestObservationErrorOverwriteBehavior(t *testing.T) {

	// Assume a situation where Slave already has an ObservationError
	db, _, _ := InitializeTestDB()
	tx := db.Begin()

	o1 := MSPError{
		Identifier: "id1",
	}

	assert.NoError(t, tx.Create(&o1).Error)

	s := Slave{
		Hostname:             "s1",
		Port:                 1,
		MongodPortRangeBegin: 1,
		MongodPortRangeEnd:   2,
		ObservationErrorID:   NullIntValue(o1.ID),
	}

	assert.NoError(t, tx.Create(&s).Error)

	var countBeforeUpdate int64
	assert.NoError(t, tx.Model(&MSPError{}).Count(&countBeforeUpdate).Error)
	assert.EqualValues(t, 1, countBeforeUpdate)

	// Now we attempt to observe Slave again and want to update the value pointed to by slave
	// We could create, update, delete
	// Or we could use the following hack to just UPDATE all values of the existing observation,
	//    saving us an annoying DELETE
	o2 := MSPError{
		Identifier: "id2",
	}
	o2.ID = o1.ID
	tx.Save(&o2)
	// That was it

	var countAfterUpdate int64
	assert.NoError(t, tx.Model(&MSPError{}).Count(&countAfterUpdate).Error)

	assert.EqualValues(t, 1, countAfterUpdate, "updates should remove the row previously referenced in the overwritten column")

}
<|endoftext|>


<commit_before>// Package mastodon provides functions and structs for accessing the mastodon API.
package mastodon

import (
	"bytes"
	"context"
	"encoding/json"
	"errors"
	"fmt"
	"io"
	"mime/multipart"
	"net/http"
	"net/url"
	"os"
	"path"
	"path/filepath"
	"strconv"
	"strings"

	"github.com/tomnomnom/linkheader"
)

// Config is a setting for access mastodon APIs.
type Config struct {
	Server       string
	ClientID     string
	ClientSecret string
	AccessToken  string
}

// Client is a API client for mastodon.
type Client struct {
	http.Client
	config *Config
}

func (c *Client) doAPI(ctx context.Context, method string, uri string, params interface{}, res interface{}, pg *Pagination) error {
	u, err := url.Parse(c.config.Server)
	if err != nil {
		return err
	}
	u.Path = path.Join(u.Path, uri)

	var req *http.Request
	ct := "application/x-www-form-urlencoded"
	if values, ok := params.(url.Values); ok {
		var body io.Reader
		if method == http.MethodGet {
			if pg != nil {
				values = pg.setValues(values)
			}
			u.RawQuery = values.Encode()
		} else {
			body = strings.NewReader(values.Encode())
		}
		req, err = http.NewRequest(method, u.String(), body)
		if err != nil {
			return err
		}
	} else if file, ok := params.(string); ok {
		f, err := os.Open(file)
		if err != nil {
			return err
		}
		defer f.Close()

		var buf bytes.Buffer
		mw := multipart.NewWriter(&buf)
		part, err := mw.CreateFormFile("file", filepath.Base(file))
		if err != nil {
			return err
		}
		_, err = io.Copy(part, f)
		if err != nil {
			return err
		}
		err = mw.Close()
		if err != nil {
			return err
		}
		req, err = http.NewRequest(method, u.String(), &buf)
		if err != nil {
			return err
		}
		ct = mw.FormDataContentType()
	} else if reader, ok := params.(io.Reader); ok {
		var buf bytes.Buffer
		mw := multipart.NewWriter(&buf)
		part, err := mw.CreateFormFile("file", "upload")
		if err != nil {
			return err
		}
		_, err = io.Copy(part, reader)
		if err != nil {
			return err
		}
		err = mw.Close()
		if err != nil {
			return err
		}
		req, err = http.NewRequest(method, u.String(), &buf)
		if err != nil {
			return err
		}
		ct = mw.FormDataContentType()
	} else {
		if method == http.MethodGet && pg != nil {
			u.RawQuery = pg.toValues().Encode()
		}
		req, err = http.NewRequest(method, u.String(), nil)
		if err != nil {
			return err
		}
	}
	req = req.WithContext(ctx)
	req.Header.Set("Authorization", "Bearer "+c.config.AccessToken)
	if params != nil {
		req.Header.Set("Content-Type", ct)
	}

	resp, err := c.Do(req)
	if err != nil {
		return err
	}
	defer resp.Body.Close()

	if resp.StatusCode != http.StatusOK {
		return parseAPIError("bad request", resp)
	} else if res == nil {
		return nil
	} else if pg != nil {
		if lh := resp.Header.Get("Link"); lh != "" {
			pg2, err := newPagination(lh)
			if err != nil {
				return err
			}
			*pg = *pg2
		}
	}
	return json.NewDecoder(resp.Body).Decode(&res)
}

// NewClient return new mastodon API client.
func NewClient(config *Config) *Client {
	return &Client{
		Client: *http.DefaultClient,
		config: config,
	}
}

// Authenticate get access-token to the API.
func (c *Client) Authenticate(ctx context.Context, username, password string) error {
	params := url.Values{
		"client_id":     {c.config.ClientID},
		"client_secret": {c.config.ClientSecret},
		"grant_type":    {"password"},
		"username":      {username},
		"password":      {password},
		"scope":         {"read write follow"},
	}

	return c.authenticate(ctx, params)
}

// AuthenticateToken logs in using a grant token returned by Application.AuthURI.
//
// redirectURI should be the same as Application.RedirectURI.
func (c *Client) AuthenticateToken(ctx context.Context, authCode, redirectURI string) error {
	params := url.Values{
		"client_id":     {c.config.ClientID},
		"client_secret": {c.config.ClientSecret},
		"grant_type":    {"authorization_code"},
		"code":          {authCode},
		"redirect_uri":  {redirectURI},
	}

	return c.authenticate(ctx, params)
}

func (c *Client) authenticate(ctx context.Context, params url.Values) error {
	u, err := url.Parse(c.config.Server)
	if err != nil {
		return err
	}
	u.Path = path.Join(u.Path, "/oauth/token")

	req, err := http.NewRequest(http.MethodPost, u.String(), strings.NewReader(params.Encode()))
	if err != nil {
		return err
	}
	req = req.WithContext(ctx)
	req.Header.Set("Content-Type", "application/x-www-form-urlencoded")
	resp, err := c.Do(req)
	if err != nil {
		return err
	}
	defer resp.Body.Close()

	if resp.StatusCode != http.StatusOK {
		return parseAPIError("bad authorization", resp)
	}

	var res struct {
		AccessToken string `json:"access_token"`
	}
	err = json.NewDecoder(resp.Body).Decode(&res)
	if err != nil {
		return err
	}
	c.config.AccessToken = res.AccessToken
	return nil
}

// Toot is struct to post status.
type Toot struct {
	Status      string `json:"status"`
	InReplyToID ID     `json:"in_reply_to_id"`
	MediaIDs    []ID   `json:"media_ids"`
	Sensitive   bool   `json:"sensitive"`
	SpoilerText string `json:"spoiler_text"`
	Visibility  string `json:"visibility"`
}

// Mention hold information for mention.
type Mention struct {
	URL      string `json:"url"`
	Username string `json:"username"`
	Acct     string `json:"acct"`
	ID       ID     `json:"id"`
}

// Tag hold information for tag.
type Tag struct {
	Name    string    `json:"name"`
	URL     string    `json:"url"`
	History []History `json:"history"`
}

// History hold information for history.
type History struct {
	Day      string `json:"day"`
	Uses     int64  `json:"uses"`
	Accounts int64  `json:"accounts"`
}

// Attachment hold information for attachment.
type Attachment struct {
	ID          ID             `json:"id"`
	Type        string         `json:"type"`
	URL         string         `json:"url"`
	RemoteURL   string         `json:"remote_url"`
	PreviewURL  string         `json:"preview_url"`
	TextURL     string         `json:"text_url"`
	Description string         `json:"description"`
	Meta        AttachmentMeta `json:"meta"`
}

// AttachmentMeta holds information for attachment metadata.
type AttachmentMeta struct {
	Original AttachmentSize `json:"original"`
	Small    AttachmentSize `json:"small"`
}

// AttachmentSize holds information for attatchment size.
type AttachmentSize struct {
	Width  int64   `json:"width"`
	Height int64   `json:"height"`
	Size   string  `json:"size"`
	Aspect float64 `json:"aspect"`
}

// Emoji hold information for CustomEmoji.
type Emoji struct {
	ShortCode       string `json:"shortcode"`
	StaticURL       string `json:"static_url"`
	URL             string `json:"url"`
	VisibleInPicker bool   `json:"visible_in_picker"`
}

// Results hold information for search result.
type Results struct {
	Accounts []*Account `json:"accounts"`
	Statuses []*Status  `json:"statuses"`
	Hashtags []string   `json:"hashtags"`
}

// Pagination is a struct for specifying the get range.
type Pagination struct {
	MaxID   ID
	SinceID ID
	MinID   ID
	Limit   int64
}

func newPagination(rawlink string) (*Pagination, error) {
	if rawlink == "" {
		return nil, errors.New("empty link header")
	}

	p := &Pagination{}
	for _, link := range linkheader.Parse(rawlink) {
		switch link.Rel {
		case "next":
			maxID, err := getPaginationID(link.URL, "max_id")
			if err != nil {
				return nil, err
			}
			p.MaxID = maxID
		case "prev":
			sinceID, err := getPaginationID(link.URL, "since_id")
			if err != nil {
				return nil, err
			}
			p.SinceID = sinceID

			minID, err := getPaginationID(link.URL, "min_id")
			if err != nil {
				return nil, err
			}
			p.MinID = minID
		}
	}

	return p, nil
}

func getPaginationID(rawurl, key string) (ID, error) {
	u, err := url.Parse(rawurl)
	if err != nil {
		return "", err
	}

	val := u.Query().Get(key)
	if val == "" {
		return "", nil
	}

	id, err := strconv.ParseInt(val, 10, 64)
	if err != nil {
		return "", err
	}

	return ID(fmt.Sprint(id)), nil
}

func (p *Pagination) toValues() url.Values {
	return p.setValues(url.Values{})
}

func (p *Pagination) setValues(params url.Values) url.Values {
	if p.MaxID != "" {
		params.Set("max_id", string(p.MaxID))
	}
	if p.SinceID != "" {
		params.Set("since_id", string(p.SinceID))
	}
	if p.MinID != "" {
		params.Set("min_id", string(p.MinID))
	}
	if p.Limit > 0 {
		params.Set("limit", fmt.Sprint(p.Limit))
	}

	return params
}
<commit_msg>Handle HTTP 429 responses with a request backoff approach<commit_after>// Package mastodon provides functions and structs for accessing the mastodon API.
package mastodon

import (
	"bytes"
	"context"
	"encoding/json"
	"errors"
	"fmt"
	"io"
	"mime/multipart"
	"net/http"
	"net/url"
	"os"
	"path"
	"path/filepath"
	"strconv"
	"strings"
	"time"

	"github.com/tomnomnom/linkheader"
)

// Config is a setting for access mastodon APIs.
type Config struct {
	Server       string
	ClientID     string
	ClientSecret string
	AccessToken  string
}

// Client is a API client for mastodon.
type Client struct {
	http.Client
	config *Config
}

func (c *Client) doAPI(ctx context.Context, method string, uri string, params interface{}, res interface{}, pg *Pagination) error {
	u, err := url.Parse(c.config.Server)
	if err != nil {
		return err
	}
	u.Path = path.Join(u.Path, uri)

	var req *http.Request
	ct := "application/x-www-form-urlencoded"
	if values, ok := params.(url.Values); ok {
		var body io.Reader
		if method == http.MethodGet {
			if pg != nil {
				values = pg.setValues(values)
			}
			u.RawQuery = values.Encode()
		} else {
			body = strings.NewReader(values.Encode())
		}
		req, err = http.NewRequest(method, u.String(), body)
		if err != nil {
			return err
		}
	} else if file, ok := params.(string); ok {
		f, err := os.Open(file)
		if err != nil {
			return err
		}
		defer f.Close()

		var buf bytes.Buffer
		mw := multipart.NewWriter(&buf)
		part, err := mw.CreateFormFile("file", filepath.Base(file))
		if err != nil {
			return err
		}
		_, err = io.Copy(part, f)
		if err != nil {
			return err
		}
		err = mw.Close()
		if err != nil {
			return err
		}
		req, err = http.NewRequest(method, u.String(), &buf)
		if err != nil {
			return err
		}
		ct = mw.FormDataContentType()
	} else if reader, ok := params.(io.Reader); ok {
		var buf bytes.Buffer
		mw := multipart.NewWriter(&buf)
		part, err := mw.CreateFormFile("file", "upload")
		if err != nil {
			return err
		}
		_, err = io.Copy(part, reader)
		if err != nil {
			return err
		}
		err = mw.Close()
		if err != nil {
			return err
		}
		req, err = http.NewRequest(method, u.String(), &buf)
		if err != nil {
			return err
		}
		ct = mw.FormDataContentType()
	} else {
		if method == http.MethodGet && pg != nil {
			u.RawQuery = pg.toValues().Encode()
		}
		req, err = http.NewRequest(method, u.String(), nil)
		if err != nil {
			return err
		}
	}
	req = req.WithContext(ctx)
	req.Header.Set("Authorization", "Bearer "+c.config.AccessToken)
	if params != nil {
		req.Header.Set("Content-Type", ct)
	}

	var resp *http.Response
	backoff := 1000 * time.Millisecond
	for {
		resp, err = c.Do(req)
		if err != nil {
			return err
		}
		defer resp.Body.Close()

		// handle status code 429, which indicates the server is throttling
		// our requests. Do an exponential backoff and retry the request.
		if resp.StatusCode == 429 {
			if backoff > time.Hour {
				break
			}
			backoff *= 2

			select {
			case <-time.After(backoff):
			case <-ctx.Done():
				return ctx.Err()
			}
			continue
		}
		break
	}

	if resp.StatusCode != http.StatusOK {
		return parseAPIError("bad request", resp)
	} else if res == nil {
		return nil
	} else if pg != nil {
		if lh := resp.Header.Get("Link"); lh != "" {
			pg2, err := newPagination(lh)
			if err != nil {
				return err
			}
			*pg = *pg2
		}
	}
	return json.NewDecoder(resp.Body).Decode(&res)
}

// NewClient return new mastodon API client.
func NewClient(config *Config) *Client {
	return &Client{
		Client: *http.DefaultClient,
		config: config,
	}
}

// Authenticate get access-token to the API.
func (c *Client) Authenticate(ctx context.Context, username, password string) error {
	params := url.Values{
		"client_id":     {c.config.ClientID},
		"client_secret": {c.config.ClientSecret},
		"grant_type":    {"password"},
		"username":      {username},
		"password":      {password},
		"scope":         {"read write follow"},
	}

	return c.authenticate(ctx, params)
}

// AuthenticateToken logs in using a grant token returned by Application.AuthURI.
//
// redirectURI should be the same as Application.RedirectURI.
func (c *Client) AuthenticateToken(ctx context.Context, authCode, redirectURI string) error {
	params := url.Values{
		"client_id":     {c.config.ClientID},
		"client_secret": {c.config.ClientSecret},
		"grant_type":    {"authorization_code"},
		"code":          {authCode},
		"redirect_uri":  {redirectURI},
	}

	return c.authenticate(ctx, params)
}

func (c *Client) authenticate(ctx context.Context, params url.Values) error {
	u, err := url.Parse(c.config.Server)
	if err != nil {
		return err
	}
	u.Path = path.Join(u.Path, "/oauth/token")

	req, err := http.NewRequest(http.MethodPost, u.String(), strings.NewReader(params.Encode()))
	if err != nil {
		return err
	}
	req = req.WithContext(ctx)
	req.Header.Set("Content-Type", "application/x-www-form-urlencoded")
	resp, err := c.Do(req)
	if err != nil {
		return err
	}
	defer resp.Body.Close()

	if resp.StatusCode != http.StatusOK {
		return parseAPIError("bad authorization", resp)
	}

	var res struct {
		AccessToken string `json:"access_token"`
	}
	err = json.NewDecoder(resp.Body).Decode(&res)
	if err != nil {
		return err
	}
	c.config.AccessToken = res.AccessToken
	return nil
}

// Toot is struct to post status.
type Toot struct {
	Status      string `json:"status"`
	InReplyToID ID     `json:"in_reply_to_id"`
	MediaIDs    []ID   `json:"media_ids"`
	Sensitive   bool   `json:"sensitive"`
	SpoilerText string `json:"spoiler_text"`
	Visibility  string `json:"visibility"`
}

// Mention hold information for mention.
type Mention struct {
	URL      string `json:"url"`
	Username string `json:"username"`
	Acct     string `json:"acct"`
	ID       ID     `json:"id"`
}

// Tag hold information for tag.
type Tag struct {
	Name    string    `json:"name"`
	URL     string    `json:"url"`
	History []History `json:"history"`
}

// History hold information for history.
type History struct {
	Day      string `json:"day"`
	Uses     int64  `json:"uses"`
	Accounts int64  `json:"accounts"`
}

// Attachment hold information for attachment.
type Attachment struct {
	ID          ID             `json:"id"`
	Type        string         `json:"type"`
	URL         string         `json:"url"`
	RemoteURL   string         `json:"remote_url"`
	PreviewURL  string         `json:"preview_url"`
	TextURL     string         `json:"text_url"`
	Description string         `json:"description"`
	Meta        AttachmentMeta `json:"meta"`
}

// AttachmentMeta holds information for attachment metadata.
type AttachmentMeta struct {
	Original AttachmentSize `json:"original"`
	Small    AttachmentSize `json:"small"`
}

// AttachmentSize holds information for attatchment size.
type AttachmentSize struct {
	Width  int64   `json:"width"`
	Height int64   `json:"height"`
	Size   string  `json:"size"`
	Aspect float64 `json:"aspect"`
}

// Emoji hold information for CustomEmoji.
type Emoji struct {
	ShortCode       string `json:"shortcode"`
	StaticURL       string `json:"static_url"`
	URL             string `json:"url"`
	VisibleInPicker bool   `json:"visible_in_picker"`
}

// Results hold information for search result.
type Results struct {
	Accounts []*Account `json:"accounts"`
	Statuses []*Status  `json:"statuses"`
	Hashtags []string   `json:"hashtags"`
}

// Pagination is a struct for specifying the get range.
type Pagination struct {
	MaxID   ID
	SinceID ID
	MinID   ID
	Limit   int64
}

func newPagination(rawlink string) (*Pagination, error) {
	if rawlink == "" {
		return nil, errors.New("empty link header")
	}

	p := &Pagination{}
	for _, link := range linkheader.Parse(rawlink) {
		switch link.Rel {
		case "next":
			maxID, err := getPaginationID(link.URL, "max_id")
			if err != nil {
				return nil, err
			}
			p.MaxID = maxID
		case "prev":
			sinceID, err := getPaginationID(link.URL, "since_id")
			if err != nil {
				return nil, err
			}
			p.SinceID = sinceID

			minID, err := getPaginationID(link.URL, "min_id")
			if err != nil {
				return nil, err
			}
			p.MinID = minID
		}
	}

	return p, nil
}

func getPaginationID(rawurl, key string) (ID, error) {
	u, err := url.Parse(rawurl)
	if err != nil {
		return "", err
	}

	val := u.Query().Get(key)
	if val == "" {
		return "", nil
	}

	id, err := strconv.ParseInt(val, 10, 64)
	if err != nil {
		return "", err
	}

	return ID(fmt.Sprint(id)), nil
}

func (p *Pagination) toValues() url.Values {
	return p.setValues(url.Values{})
}

func (p *Pagination) setValues(params url.Values) url.Values {
	if p.MaxID != "" {
		params.Set("max_id", string(p.MaxID))
	}
	if p.SinceID != "" {
		params.Set("since_id", string(p.SinceID))
	}
	if p.MinID != "" {
		params.Set("min_id", string(p.MinID))
	}
	if p.Limit > 0 {
		params.Set("limit", fmt.Sprint(p.Limit))
	}

	return params
}
<|endoftext|>


<commit_before>package model

import (
	"github.com/slyrz/newscat/html"
	"github.com/slyrz/newscat/util"
)

const (
	numChunkFeatureComp = 36
	numBoostFeatureComp = 10
)

// feature represents a feature vector.
type feature []float32

type chunkFeature [numChunkFeatureComp]float32
type boostFeature [numBoostFeatureComp]float32

// featureWriter writes observations to feature vectors.
type featureWriter struct {
	feature feature
	Pos     int
}

func (fw *featureWriter) Assign(f feature) {
	// If we haven't fully filled the current feature vector yet, but are about to
	// replace it with a new one, start to panic. We don't want to produce partially
	// initalized feature vectors here.
	if fw.feature != nil && fw.Pos != cap(fw.feature) {
		panic("partially filled feature")
	}
	fw.feature = f
	fw.Pos = 0
}

// Write a value of type int, float32 or bool at given offset and skip the
// requested amount of components afterwards.
func (fw *featureWriter) write(val interface{}, off int, skip int) {
	comp := &fw.feature[fw.Pos+off]
	switch val := val.(type) {
	case int:
		*comp = float32(val)
	case float32:
		*comp = val
	case bool:
		if val {
			*comp = 1.0
		} else {
			*comp = 0.0
		}
	}
	if skip > 0 {
		fw.Skip(skip)
	}
}

// Write value at current position and move to the next.
func (fw *featureWriter) Write(val interface{}) {
	fw.write(val, 0, 1)
}

// Write value at offset, but don't move.
func (fw *featureWriter) WriteAt(val interface{}, off int) {
	fw.write(val, off, 0)
}

// Skip components.
func (fw *featureWriter) Skip(n int) {
	fw.Pos += n
}

type chunkFeatureWriter struct {
	featureWriter
}

// Entries with a "plus comment" indicate that the next N elements share
// the same offset intentionally.
var elementTypes = map[string]int{
	"p":   0,
	"a":   1,
	"div": 2,
	"h1":  3, // +5
	"h2":  3,
	"h3":  3,
	"h4":  3,
	"h5":  3,
	"h6":  3,
}

func (fw *chunkFeatureWriter) WriteElementType(chunk *html.Chunk) {
	// One hot encoding of the element type.
	fw.WriteAt(true, elementTypes[chunk.Base.Data])
	fw.Skip(4)
}

var parentTypes = map[string]int{
	"p":    0,
	"span": 1,
	"div":  2,
	"li":   3,
}

func (fw *chunkFeatureWriter) WriteParentType(chunk *html.Chunk) {
	// One hot encoding of the chunk's parent's element type.
	if chunk.Base.Parent != nil {
		fw.WriteAt(true, parentTypes[chunk.Base.Parent.Data])
	}
	fw.Skip(4)
}

func (fw *chunkFeatureWriter) WriteSiblingTypes(chunk *html.Chunk) {
	count := 0
	types := map[string]int{"a": 0, "p": 0, "img": 0}
	for _, siblingType := range chunk.GetSiblingTypes() {
		count += 1
		if val, ok := types[siblingType]; ok {
			types[siblingType] = val + 1
		}
	}
	fw.Write(count)
	fw.Write(types["a"])
	fw.Write(types["p"])
	fw.Write(types["img"])
	if count > 0 {
		fw.Write(float32(types["a"]) / float32(count))
		fw.Write(float32(types["p"]) / float32(count))
		fw.Write(float32(types["img"]) / float32(count))
	} else {
		fw.Skip(3)
	}
}

func (fw *chunkFeatureWriter) WriteAncestors(chunk *html.Chunk) {
	fw.Write((chunk.Ancestors & html.AncestorArticle) != 0)
	fw.Write((chunk.Ancestors & html.AncestorAside) != 0)
	fw.Write((chunk.Ancestors & html.AncestorBlockquote) != 0)
	fw.Write((chunk.Ancestors & html.AncestorList) != 0)
}

func (fw *chunkFeatureWriter) WriteTextStat(chunk *html.Chunk) {
	fw.Write(chunk.Text.Words)
	fw.Write(chunk.Text.Sentences)
	fw.Write(chunk.LinkText)
}

func (fw *chunkFeatureWriter) WriteTextStatSiblings(chunk *html.Chunk) {
	if chunk.Prev != nil {
		fw.Write(chunk.Prev.Block == chunk.Block)
		fw.Write(chunk.Prev.Text.Words)
		fw.Write(chunk.Prev.Text.Sentences)
	} else {
		fw.Skip(3)
	}
	if chunk.Next != nil {
		fw.Write(chunk.Next.Block == chunk.Block)
		fw.Write(chunk.Next.Text.Words)
		fw.Write(chunk.Next.Text.Sentences)
	} else {
		fw.Skip(3)
	}
}

func (fw *chunkFeatureWriter) WriteClassStat(chunk *html.Chunk, classes map[string]*html.TextStat) {
	var best *html.TextStat = nil
	for _, class := range chunk.Classes {
		if stat, ok := classes[class]; ok {
			if best == nil || (stat.Words/stat.Count) > (best.Words/best.Count) {
				best = stat
			}
		}
	}
	if best != nil {
		fw.Write(true)
		fw.Write(float32(best.Words) / float32(best.Count))
		fw.Write(float32(best.Sentences) / float32(best.Count))
	} else {
		fw.Write(false)
		fw.Skip(2)
	}
}

func (fw *chunkFeatureWriter) WriteClusterStat(chunk *html.Chunk, clusters map[*html.Chunk]*html.TextStat) {
	if stat, ok := clusters[chunk]; ok {
		fw.Write(stat.Words)
		fw.Write(stat.Sentences)
		fw.Write(stat.Count)
		fw.Write(float32(stat.Words) / float32(stat.Count))
		fw.Write(float32(stat.Sentences) / float32(stat.Count))
	} else {
		fw.Skip(5)
	}
}

type boostFeatureWriter struct {
	featureWriter
}

var (
	goodQualClass = util.NewRegexFromWords(
		"article",
		"catchline",
		"chapter",
		"content",
		"head",
		"intro",
		"introduction",
		"leadin",
		"main",
		"post",
		"story",
		"summary",
		"title",
	)
	poorQualClass = util.NewRegexFromWords(
		"author",
		"blog",
		"byline",
		"caption",
		"col",
		"comment",
		"description",
		"email",
		"excerpt",
		"image",
		"info",
		"menu",
		"metadata",
		"nav",
		"photo",
		"small",
		"teaser",
		"widget",
	)
)

func (fw *boostFeatureWriter) WriteChunk(chunk *html.Chunk) {
	goodQual := false
	poorQual := false
	for _, class := range chunk.Classes {
		goodQual = goodQual || goodQualClass.In(class)
		poorQual = poorQual || poorQualClass.In(class)
	}
	fw.Write(chunk.LinkText)
	fw.Write(chunk.Text.Words)
	fw.Write(chunk.Text.Sentences)
	fw.Write(goodQual)
	fw.Write(poorQual)
}

func (fw *boostFeatureWriter) WriteCluster(chunk *html.Chunk, cluster *cluster) {
	i := 0
	for ; i < len(cluster.Chunks); i++ {
		if cluster.Chunks[i] == chunk {
			break
		}
	}
	fw.Write(cluster.Score())
	fw.Write(cluster.Scores[i])
	if i > 0 {
		fw.Write(cluster.Scores[i-1])
	} else {
		fw.Write(-10)
	}
	if i < len(cluster.Chunks)-2 {
		fw.Write(cluster.Scores[i+1])
	} else {
		fw.Write(-10)
	}
}

func (fw *boostFeatureWriter) WriteTitleSimilarity(chunk *html.Chunk, title *util.Text) {
	switch chunk.Base.Data {
	case "h1", "h2", "h3":
		fw.Write(chunk.Text.Similarity(title))
	default:
		fw.Skip(1)
	}
}
<commit_msg>renamed variables<commit_after>package model

import (
	"github.com/slyrz/newscat/html"
	"github.com/slyrz/newscat/util"
)

const (
	chunkFeatureCap = 36
	boostFeatureCap = 10
)

// feature represents a feature vector.
type feature []float32

type chunkFeature [chunkFeatureCap]float32
type boostFeature [boostFeatureCap]float32

// featureWriter writes observations to feature vectors.
type featureWriter struct {
	feature feature
	Pos     int
}

func (fw *featureWriter) Assign(f feature) {
	// If we haven't fully filled the current feature vector yet, but are about to
	// replace it with a new one, start to panic. We don't want to produce partially
	// initalized feature vectors here.
	if fw.feature != nil && fw.Pos != cap(fw.feature) {
		panic("partially filled feature")
	}
	fw.feature = f
	fw.Pos = 0
}

// Write a value of type int, float32 or bool at given offset and skip the
// requested amount of components afterwards.
func (fw *featureWriter) write(val interface{}, off int, skip int) {
	comp := &fw.feature[fw.Pos+off]
	switch val := val.(type) {
	case int:
		*comp = float32(val)
	case float32:
		*comp = val
	case bool:
		if val {
			*comp = 1.0
		} else {
			*comp = 0.0
		}
	}
	if skip > 0 {
		fw.Skip(skip)
	}
}

// Write value at current position and move to the next.
func (fw *featureWriter) Write(val interface{}) {
	fw.write(val, 0, 1)
}

// Write value at offset, but don't move.
func (fw *featureWriter) WriteAt(val interface{}, off int) {
	fw.write(val, off, 0)
}

// Skip components.
func (fw *featureWriter) Skip(n int) {
	fw.Pos += n
}

type chunkFeatureWriter struct {
	featureWriter
}

// Entries with a "plus comment" indicate that the next N elements share
// the same offset intentionally.
var elementTypes = map[string]int{
	"p":   0,
	"a":   1,
	"div": 2,
	"h1":  3, // +5
	"h2":  3,
	"h3":  3,
	"h4":  3,
	"h5":  3,
	"h6":  3,
}

func (fw *chunkFeatureWriter) WriteElementType(chunk *html.Chunk) {
	// One hot encoding of the element type.
	fw.WriteAt(true, elementTypes[chunk.Base.Data])
	fw.Skip(4)
}

var parentTypes = map[string]int{
	"p":    0,
	"span": 1,
	"div":  2,
	"li":   3,
}

func (fw *chunkFeatureWriter) WriteParentType(chunk *html.Chunk) {
	// One hot encoding of the chunk's parent's element type.
	if chunk.Base.Parent != nil {
		fw.WriteAt(true, parentTypes[chunk.Base.Parent.Data])
	}
	fw.Skip(4)
}

func (fw *chunkFeatureWriter) WriteSiblingTypes(chunk *html.Chunk) {
	count := 0
	types := map[string]int{"a": 0, "p": 0, "img": 0}
	for _, siblingType := range chunk.GetSiblingTypes() {
		count += 1
		if val, ok := types[siblingType]; ok {
			types[siblingType] = val + 1
		}
	}
	fw.Write(count)
	fw.Write(types["a"])
	fw.Write(types["p"])
	fw.Write(types["img"])
	if count > 0 {
		fw.Write(float32(types["a"]) / float32(count))
		fw.Write(float32(types["p"]) / float32(count))
		fw.Write(float32(types["img"]) / float32(count))
	} else {
		fw.Skip(3)
	}
}

func (fw *chunkFeatureWriter) WriteAncestors(chunk *html.Chunk) {
	fw.Write((chunk.Ancestors & html.AncestorArticle) != 0)
	fw.Write((chunk.Ancestors & html.AncestorAside) != 0)
	fw.Write((chunk.Ancestors & html.AncestorBlockquote) != 0)
	fw.Write((chunk.Ancestors & html.AncestorList) != 0)
}

func (fw *chunkFeatureWriter) WriteTextStat(chunk *html.Chunk) {
	fw.Write(chunk.Text.Words)
	fw.Write(chunk.Text.Sentences)
	fw.Write(chunk.LinkText)
}

func (fw *chunkFeatureWriter) WriteTextStatSiblings(chunk *html.Chunk) {
	if chunk.Prev != nil {
		fw.Write(chunk.Prev.Block == chunk.Block)
		fw.Write(chunk.Prev.Text.Words)
		fw.Write(chunk.Prev.Text.Sentences)
	} else {
		fw.Skip(3)
	}
	if chunk.Next != nil {
		fw.Write(chunk.Next.Block == chunk.Block)
		fw.Write(chunk.Next.Text.Words)
		fw.Write(chunk.Next.Text.Sentences)
	} else {
		fw.Skip(3)
	}
}

func (fw *chunkFeatureWriter) WriteClassStat(chunk *html.Chunk, classes map[string]*html.TextStat) {
	var best *html.TextStat = nil
	for _, class := range chunk.Classes {
		if stat, ok := classes[class]; ok {
			if best == nil || (stat.Words/stat.Count) > (best.Words/best.Count) {
				best = stat
			}
		}
	}
	if best != nil {
		fw.Write(true)
		fw.Write(float32(best.Words) / float32(best.Count))
		fw.Write(float32(best.Sentences) / float32(best.Count))
	} else {
		fw.Write(false)
		fw.Skip(2)
	}
}

func (fw *chunkFeatureWriter) WriteClusterStat(chunk *html.Chunk, clusters map[*html.Chunk]*html.TextStat) {
	if stat, ok := clusters[chunk]; ok {
		fw.Write(stat.Words)
		fw.Write(stat.Sentences)
		fw.Write(stat.Count)
		fw.Write(float32(stat.Words) / float32(stat.Count))
		fw.Write(float32(stat.Sentences) / float32(stat.Count))
	} else {
		fw.Skip(5)
	}
}

type boostFeatureWriter struct {
	featureWriter
}

var (
	goodQualClass = util.NewRegexFromWords(
		"article",
		"catchline",
		"chapter",
		"content",
		"head",
		"intro",
		"introduction",
		"leadin",
		"main",
		"post",
		"story",
		"summary",
		"title",
	)
	poorQualClass = util.NewRegexFromWords(
		"author",
		"blog",
		"byline",
		"caption",
		"col",
		"comment",
		"description",
		"email",
		"excerpt",
		"image",
		"info",
		"menu",
		"metadata",
		"nav",
		"photo",
		"small",
		"teaser",
		"widget",
	)
)

func (fw *boostFeatureWriter) WriteChunk(chunk *html.Chunk) {
	goodQual := false
	poorQual := false
	for _, class := range chunk.Classes {
		goodQual = goodQual || goodQualClass.In(class)
		poorQual = poorQual || poorQualClass.In(class)
	}
	fw.Write(chunk.LinkText)
	fw.Write(chunk.Text.Words)
	fw.Write(chunk.Text.Sentences)
	fw.Write(goodQual)
	fw.Write(poorQual)
}

func (fw *boostFeatureWriter) WriteCluster(chunk *html.Chunk, cluster *cluster) {
	i := 0
	for ; i < len(cluster.Chunks); i++ {
		if cluster.Chunks[i] == chunk {
			break
		}
	}
	fw.Write(cluster.Score())
	fw.Write(cluster.Scores[i])
	if i > 0 {
		fw.Write(cluster.Scores[i-1])
	} else {
		fw.Write(-10)
	}
	if i < len(cluster.Chunks)-2 {
		fw.Write(cluster.Scores[i+1])
	} else {
		fw.Write(-10)
	}
}

func (fw *boostFeatureWriter) WriteTitleSimilarity(chunk *html.Chunk, title *util.Text) {
	switch chunk.Base.Data {
	case "h1", "h2", "h3":
		fw.Write(chunk.Text.Similarity(title))
	default:
		fw.Skip(1)
	}
}
<|endoftext|>


<commit_before>/*

Hidden Markov model.

 α β γ ζ Φ

*/
package hmm

import (
	"code.google.com/p/biogo.matrix"
	"fmt"
	"github.com/akualab/gjoa/model"
	"github.com/golang/glog"
	"math"
)

const (
	SMALL_NUMBER = 0.000001
)

type HMM struct {

	// q(t) is the state at time t
	// states are labeled {0,2,...,N-1}

	// Number fo hidden states.
	// N
	nstates int

	// State-transition probability distribution matrixi (log scale).
	// [nstates x nstates]
	// a(i,j) = log(P[q(t+1) = j | q(t) = i]); 0 <= i,j <= N-1
	// when transitions between states are not allowed: a(i,j) = 0
	// => saved in log domain.
	logTransProbs *matrix.Dense

	// Observation probability distribution functions. [nstates x 1]
	// b(j,k) = P[o(t) | q(t) = j]
	obsModels []model.Modeler

	// Initial state distribution. [nstates x 1]
	// π(i) = P[q(0) = i]; 0<=i<N
	// => saved in log domain.
	logInitProbs *matrix.Dense

	// Num elements in obs vector.
	numElements int

	// Complete HMM model.
	// Φ = (A, B, π)
}

// Define functions for elementwise transformations.
var log = func(r, c int, v float64) float64 { return math.Log(v) }

// Create a new HMM.
func NewHMM(transProbs, initialStateProbs *matrix.Dense, obsModels []model.Modeler) (hmm *HMM, e error) {

	r, c := transProbs.Dims()
	if r != c {
		e = fmt.Errorf("Matrix transProbs must be square. rows=[%d], cols=[%d]", r, c)
		return
	}

	glog.Infof("New HMM. Num states = [%d].", r)

	// init logTransProbs and logInitProbs
	logTransProbs := matrix.MustDense(matrix.ZeroDense(r, r))
	logInitProbs := matrix.MustDense(matrix.ZeroDense(r, 1))

	// apply log to transProbs and initialStateProbs
	logTransProbs = transProbs.ApplyDense(log, logTransProbs)
	logInitProbs = initialStateProbs.ApplyDense(log, logInitProbs)

	hmm = &HMM{
		nstates:       r,
		logTransProbs: logTransProbs,
		obsModels:     obsModels,
		logInitProbs:  logInitProbs,
		numElements:   obsModels[0].NumElements(),
	}

	return
}

// Compute alphas. Indices are: α(state, time)
//
// α = | α(0,0),   α(0,1)   ... α(0,T-1)   |
//     | α(1,0),   α(1,1)   ... α(1,T-1)   |
//     ...
//     | α(N-1,0), α(N-1,1) ... α(N-1,T-1) |
//
//
// 1. Initialization: α(i,0) =  π(i) b(i,o(0)); 0<=i<N
// 2. Induction:      α(j,t+1) =  sum_{i=0}^{N-1}[α(i,t)a(i,j)] b(j,o(t+1)); 0<=t<T-1; 0<=j<N
// 3. Termination:    P(O/Φ) = sum_{i=0}^{N-1} α(i,T-1)
// For scaling details see Rabiner/Juang
func (hmm *HMM) alpha(observations *matrix.Dense) (α *matrix.Dense, logProb float64, e error) {

	// Num states.
	N := hmm.nstates

	// expected num rows: numElements
	// expected num cols: T
	ne, T := observations.Dims()

	if ne != hmm.numElements {
		e = fmt.Errorf("Mismatch in num elements in observations [%d] expected [%d].", ne, hmm.numElements)
		return
	}

	// Allocate log-alpha matrix.
	// TODO: use a reusable data structure to minimize garbage.
	α = matrix.MustDense(matrix.ZeroDense(N, T))

	// 1. Initialization. Add in the log doman.
	for i := 0; i < N; i++ {
		α.Set(i, 0, hmm.logInitProbs.At(i, 0)+hmm.obsModels[i].LogProb(ColumnAt(observations, 0)))
	}

	// 2. Induction.
	var sumAlphas, sum float64
	for t := 0; t < T-1; t++ {
		sumAlphas = 0
		for j := 0; j < N; j++ {

			sum = 0
			for i := 0; i < N; i++ {
				sum += math.Exp(α.At(i, t) + hmm.logTransProbs.At(i, j))
			}
			v := math.Log(sum) + hmm.obsModels[j].LogProb(ColumnAt(observations, t+1))
			α.Set(j, t+1, v)

			sumAlphas += math.Exp(v)
			if glog.V(4) {
				glog.Infof("t: %4d | j: %2d | logAlpha: %5e | sumAlphas: %5e", t, j, v, sumAlphas)
			}
		}
		// Applied scale for t independent of j.
		if sumAlphas < SMALL_NUMBER {
			continue
		}
		for j := 0; j < N; j++ {
			v := α.At(j, t+1) / sumAlphas
			α.Set(j, t+1, v)
		}
		logProb -= math.Log(sumAlphas)
	}

	return
}

// Compute betas. Indices are: β(state, time)
//
// β = | β(0,0),   β(0,1)   ... β(0,T-1)   |
//     | β(1,0),   β(1,1)   ... β(1,T-1)   |
//     ...
//     | β(N-1,0), β(N-1,1) ... β(N-1,T-1) |
//
//
// 1. Initialization: β(i,T-1) = 1;  0<=i<N
// 2. Induction:      β(i,t) =  sum_{j=0}^{N-1} a(i,j) b(j,o(t+1)) β(j,t+1); t=T-2,T-3,...,0; 0<=i<N
func (hmm *HMM) beta(observations *matrix.Dense) (β *matrix.Dense, e error) {

	// Num states.
	N := hmm.nstates

	// expected num rows: numElements
	// expected num cols: T
	ne, T := observations.Dims()

	if ne != hmm.numElements {
		e = fmt.Errorf("Mismatch in num elements in observations [%d] expected [%d].", ne, hmm.numElements)
		return
	}

	// Allocate log-beta matrix.
	// TODO: use a reusable data structure to minimize garbage.
	β = matrix.MustDense(matrix.ZeroDense(N, T))

	// 1. Initialization. Add in the log doman.
	for i := 0; i < N; i++ {
		β.Set(i, T-1, 1.0)
	}

	// 2. Induction.
	for t := T - 2; t >= 0; t-- {
		for i := 0; i < N; i++ {

			var sum float64
			for j := 0; j < N; j++ {

				sum += math.Exp(hmm.logTransProbs.At(i, j) + // a(i,j)
					hmm.obsModels[j].LogProb(ColumnAt(observations, t+1)) + // b(j,o(t+1))
					β.At(j, t+1)) // β(j,t+1)
			}
			β.Set(i, t, math.Log(sum))
		}
	}

	return
}

// Compute gammas. Indices are: γ(state, time)
//
// γ(i,t) =  α(i,t)β(i,t) / sum_{j=0}^{N-1} α(j,t)β(j,t);  0<=j<N
func (hmm *HMM) gamma(α, β *matrix.Dense) (γ *matrix.Dense, e error) {

	αr, αc := α.Dims()
	βr, βc := β.Dims()

	if αr != βr || αc != βc {
		e = fmt.Errorf("Shape mismatch: alpha[%d,%d] beta[%d,%d]", αr, αc, βr, βc)
		return
	}

	T := αc
	N := hmm.nstates
	if αr != N {
		e = fmt.Errorf("Num rows [%d] doesn't match num states [%d].", αr, N)
	}

	// Allocate log-gamma matrix.
	// TODO: use a reusable data structure to minimize garbage.
	γ = matrix.MustDense(matrix.ZeroDense(N, T))

	for t := 0; t < T; t++ {
		var sum float64 = 0.0
		for i := 0; i < N; i++ {
			x := α.At(i, t) + β.At(i, t)
			γ.Set(i, t, x)
			sum += math.Exp(x)
		}

		// Normalize.
		for i := 0; i < N; i++ {
			x := γ.At(i, t) - math.Log(sum)
			γ.Set(i, t, x)
		}
	}
	return
}

/*
   Compute xi. Indices are: ζ(from, to, time)

                         α(i,t) a(i,j) b(j,o(t+1)) β(j,t+1)
 ζ(i,j,t) = ------------------------------------------------------------------
            sum_{i=0}^{N-1} sum_{j=0}^{N-1} α(i,t) a(i,j) b(j,o(t+1)) β(j,t+1)

*/
func (hmm *HMM) xi(observations, α, β *matrix.Dense) (ζ [][][]float64, e error) {

	a := hmm.logTransProbs
	αr, αc := α.Dims()
	βr, βc := β.Dims()
	or, oc := observations.Dims()

	if αr != βr || αc != βc {
		e = fmt.Errorf("Shape mismatch: alpha[%d,%d] beta[%d,%d]", αr, αc, βr, βc)
		return
	}

	T := αc
	N := hmm.nstates
	if or != hmm.numElements {
		e = fmt.Errorf("Mismatch in num elements in observations [%d] expected [%d].",
			or, hmm.numElements)
		return
	}
	if oc != T {
		e = fmt.Errorf("Mismatch in T observations has [%d], expected [%d].", oc, T)
		return
	}
	if αr != N {
		e = fmt.Errorf("Num rows [%d] doesn't match num states [%d].", αr, N)
	}

	// Allocate log-xi matrix.
	// TODO: use a reusable data structure to minimize garbage.
	ζ = make([][][]float64, N)
	for i := 0; i < N; i++ {
		ζ[i] = make([][]float64, N)
		for j := 0; j < N; j++ {
			for t := 0; t < N; t++ {
				ζ[i][j] = make([]float64, T)
			}
		}
	}

	for t := 0; t < T-1; t++ {
		var sum float64 = 0.0
		for j := 0; j < N; j++ {
			b := hmm.obsModels[j].LogProb(ColumnAt(observations, t+1))
			for i := 0; i < N; i++ {
				x := α.At(i, t) + a.At(i, j) + b + β.At(j, t+1)
				ζ[i][j][t] = x
				sum += math.Exp(x)
			}
		}
		// Normalize.
		for i := 0; i < N; i++ {
			for j := 0; j < N; j++ {
				ζ[i][j][t] -= math.Log(sum)
			}
		}
	}
	return
}

// Update model statistics.
// sequence is a matrix
func (hmm *HMM) Update(observations *matrix.Dense) (e error) {

	var α, β, γ *matrix.Dense
	var ζ [][][]float64
	var logProb float64

	// Compute  α, β, γ, ζ
	// TODO: compute α, β, concurrently using go routines.
	α, logProb, e = hmm.alpha(observations)
	if e != nil {
		return
	}
	β, e = hmm.beta(observations)
	if e != nil {
		return
	}
	// TODO: compute γ, ζ concurrently using go routines.
	// Can we compute ζ more efficiently using γ?
	γ, e = hmm.gamma(α, β)
	if e != nil {
		return
	}
	ζ, e = hmm.xi(observations, α, β)
	if e != nil {
		return
	}

	fmt.Println(α, β, γ, ζ, logProb)
	return
}

// ColumnAt returns a *matrix.Dense column that is a copy of the values at column c of the matrix.
// Column will panic with ErrIndexOutOfRange is c is not a valid column index.
func ColumnAt(d *matrix.Dense, c int) *matrix.Dense {

	sl := d.Column(c)
	nrows := len(sl)

	col := matrix.MustDense(matrix.ZeroDense(nrows, 1))

	for i := 0; i < nrows; i++ {
		col.Set(i, 0, sl[i])
	}

	return col
}

func (hmm *HMM) NumElements() int {
	return hmm.numElements
}
<commit_msg>Fix alpha scaling. Implement beta scaling.<commit_after>/*

Hidden Markov model.

 α β γ ζ Φ

*/
package hmm

import (
	"code.google.com/p/biogo.matrix"
	"fmt"
	"github.com/akualab/gjoa/model"
	"github.com/golang/glog"
	"math"
)

const (
	SMALL_NUMBER = 0.000001
)

type HMM struct {

	// q(t) is the state at time t
	// states are labeled {0,2,...,N-1}

	// Number fo hidden states.
	// N
	nstates int

	// State-transition probability distribution matrixi (log scale).
	// [nstates x nstates]
	// a(i,j) = log(P[q(t+1) = j | q(t) = i]); 0 <= i,j <= N-1
	// when transitions between states are not allowed: a(i,j) = 0
	// => saved in log domain.
	logTransProbs *matrix.Dense

	// Observation probability distribution functions. [nstates x 1]
	// b(j,k) = P[o(t) | q(t) = j]
	obsModels []model.Modeler

	// Initial state distribution. [nstates x 1]
	// π(i) = P[q(0) = i]; 0<=i<N
	// => saved in log domain.
	logInitProbs *matrix.Dense

	// Num elements in obs vector.
	numElements int

	// Complete HMM model.
	// Φ = (A, B, π)
}

// Define functions for elementwise transformations.
var log = func(r, c int, v float64) float64 { return math.Log(v) }

// Create a new HMM.
func NewHMM(transProbs, initialStateProbs *matrix.Dense, obsModels []model.Modeler) (hmm *HMM, e error) {

	r, c := transProbs.Dims()
	if r != c {
		e = fmt.Errorf("Matrix transProbs must be square. rows=[%d], cols=[%d]", r, c)
		return
	}

	glog.Infof("New HMM. Num states = [%d].", r)

	// init logTransProbs and logInitProbs
	logTransProbs := matrix.MustDense(matrix.ZeroDense(r, r))
	logInitProbs := matrix.MustDense(matrix.ZeroDense(r, 1))

	// apply log to transProbs and initialStateProbs
	logTransProbs = transProbs.ApplyDense(log, logTransProbs)
	logInitProbs = initialStateProbs.ApplyDense(log, logInitProbs)

	hmm = &HMM{
		nstates:       r,
		logTransProbs: logTransProbs,
		obsModels:     obsModels,
		logInitProbs:  logInitProbs,
		numElements:   obsModels[0].NumElements(),
	}

	return
}

// Compute alphas. Indices are: α(state, time)
//
// α = | α(0,0),   α(0,1)   ... α(0,T-1)   |
//     | α(1,0),   α(1,1)   ... α(1,T-1)   |
//     ...
//     | α(N-1,0), α(N-1,1) ... α(N-1,T-1) |
//
//
// 1. Initialization: α(i,0) =  π(i) b(i,o(0)); 0<=i<N
// 2. Induction:      α(j,t+1) =  sum_{i=0}^{N-1}[α(i,t)a(i,j)] b(j,o(t+1)); 0<=t<T-1; 0<=j<N
// 3. Termination:    P(O/Φ) = sum_{i=0}^{N-1} α(i,T-1)
// For scaling details see Rabiner/Juang
func (hmm *HMM) alpha(observations *matrix.Dense) (α *matrix.Dense, logProb float64, e error) {

	// Num states.
	N := hmm.nstates

	// expected num rows: numElements
	// expected num cols: T
	ne, T := observations.Dims()

	if ne != hmm.numElements {
		e = fmt.Errorf("Mismatch in num elements in observations [%d] expected [%d].", ne, hmm.numElements)
		return
	}

	// Allocate log-alpha matrix.
	// TODO: use a reusable data structure to minimize garbage.
	α = matrix.MustDense(matrix.ZeroDense(N, T))

	// 1. Initialization. Add in the log doman.
	for i := 0; i < N; i++ {
		α.Set(i, 0, hmm.logInitProbs.At(i, 0)+hmm.obsModels[i].LogProb(ColumnAt(observations, 0)))
	}

	// 2. Induction.
	var sumAlphas, sum float64
	for t := 0; t < T-1; t++ {
		sumAlphas = 0
		for j := 0; j < N; j++ {

			sum = 0
			for i := 0; i < N; i++ {
				sum += math.Exp(α.At(i, t) + hmm.logTransProbs.At(i, j))
			}
			v := math.Log(sum) + hmm.obsModels[j].LogProb(ColumnAt(observations, t+1))
			α.Set(j, t+1, v)

			sumAlphas += math.Exp(v)
			if glog.V(4) {
				glog.Infof("t: %4d | j: %2d | logAlpha: %5e | sumAlphas: %5e", t, j, v, sumAlphas)
			}
		}
		// Applied scale for t independent of j.
		logSumAlphas := math.Log(sumAlphas)
		for j := 0; j < N; j++ {
			v := α.At(j, t+1) - logSumAlphas
			α.Set(j, t+1, v)
		}
		logProb += logSumAlphas
	}

	return
}

// Compute betas. Indices are: β(state, time)
//
// β = | β(0,0),   β(0,1)   ... β(0,T-1)   |
//     | β(1,0),   β(1,1)   ... β(1,T-1)   |
//     ...
//     | β(N-1,0), β(N-1,1) ... β(N-1,T-1) |
//
//
// 1. Initialization: β(i,T-1) = 1;  0<=i<N
// 2. Induction:      β(i,t) =  sum_{j=0}^{N-1} a(i,j) b(j,o(t+1)) β(j,t+1); t=T-2,T-3,...,0; 0<=i<N
func (hmm *HMM) beta(observations *matrix.Dense) (β *matrix.Dense, e error) {

	// Num states.
	N := hmm.nstates

	// expected num rows: numElements
	// expected num cols: T
	ne, T := observations.Dims()

	if ne != hmm.numElements {
		e = fmt.Errorf("Mismatch in num elements in observations [%d] expected [%d].", ne, hmm.numElements)
		return
	}

	// Allocate log-beta matrix.
	// TODO: use a reusable data structure to minimize garbage.
	β = matrix.MustDense(matrix.ZeroDense(N, T))

	// 1. Initialization. Add in the log doman.
	for i := 0; i < N; i++ {
		β.Set(i, T-1, 1.0)
	}

	// 2. Induction.
	var sumBetas float64
	for t := T - 2; t >= 0; t-- {
		sumBetas = 0
		for i := 0; i < N; i++ {

			var sum float64
			for j := 0; j < N; j++ {

				sum += math.Exp(hmm.logTransProbs.At(i, j) + // a(i,j)
					hmm.obsModels[j].LogProb(ColumnAt(observations, t+1)) + // b(j,o(t+1))
					β.At(j, t+1)) // β(j,t+1)
			}
			β.Set(i, t, math.Log(sum))
			sumBetas += sum
		}
		// Applied scale for t independent of i.
		logSumBetas := math.Log(sumBetas)
		for i := 0; i < N; i++ {
			v := β.At(i, t) - logSumBetas
			β.Set(i, t, v)
		}
	}

	return
}

// Compute gammas. Indices are: γ(state, time)
//
// γ(i,t) =  α(i,t)β(i,t) / sum_{j=0}^{N-1} α(j,t)β(j,t);  0<=j<N
func (hmm *HMM) gamma(α, β *matrix.Dense) (γ *matrix.Dense, e error) {

	αr, αc := α.Dims()
	βr, βc := β.Dims()

	if αr != βr || αc != βc {
		e = fmt.Errorf("Shape mismatch: alpha[%d,%d] beta[%d,%d]", αr, αc, βr, βc)
		return
	}

	T := αc
	N := hmm.nstates
	if αr != N {
		e = fmt.Errorf("Num rows [%d] doesn't match num states [%d].", αr, N)
	}

	// Allocate log-gamma matrix.
	// TODO: use a reusable data structure to minimize garbage.
	γ = matrix.MustDense(matrix.ZeroDense(N, T))

	for t := 0; t < T; t++ {
		var sum float64 = 0.0
		for i := 0; i < N; i++ {
			x := α.At(i, t) + β.At(i, t)
			γ.Set(i, t, x)
			sum += math.Exp(x)
		}

		// Normalize.
		for i := 0; i < N; i++ {
			x := γ.At(i, t) - math.Log(sum)
			γ.Set(i, t, x)
		}
	}
	return
}

/*
   Compute xi. Indices are: ζ(from, to, time)

                         α(i,t) a(i,j) b(j,o(t+1)) β(j,t+1)
 ζ(i,j,t) = ------------------------------------------------------------------
            sum_{i=0}^{N-1} sum_{j=0}^{N-1} α(i,t) a(i,j) b(j,o(t+1)) β(j,t+1)

*/
func (hmm *HMM) xi(observations, α, β *matrix.Dense) (ζ [][][]float64, e error) {

	a := hmm.logTransProbs
	αr, αc := α.Dims()
	βr, βc := β.Dims()
	or, oc := observations.Dims()

	if αr != βr || αc != βc {
		e = fmt.Errorf("Shape mismatch: alpha[%d,%d] beta[%d,%d]", αr, αc, βr, βc)
		return
	}

	T := αc
	N := hmm.nstates
	if or != hmm.numElements {
		e = fmt.Errorf("Mismatch in num elements in observations [%d] expected [%d].",
			or, hmm.numElements)
		return
	}
	if oc != T {
		e = fmt.Errorf("Mismatch in T observations has [%d], expected [%d].", oc, T)
		return
	}
	if αr != N {
		e = fmt.Errorf("Num rows [%d] doesn't match num states [%d].", αr, N)
	}

	// Allocate log-xi matrix.
	// TODO: use a reusable data structure to minimize garbage.
	ζ = make([][][]float64, N)
	for i := 0; i < N; i++ {
		ζ[i] = make([][]float64, N)
		for j := 0; j < N; j++ {
			for t := 0; t < N; t++ {
				ζ[i][j] = make([]float64, T)
			}
		}
	}

	for t := 0; t < T-1; t++ {
		var sum float64 = 0.0
		for j := 0; j < N; j++ {
			b := hmm.obsModels[j].LogProb(ColumnAt(observations, t+1))
			for i := 0; i < N; i++ {
				x := α.At(i, t) + a.At(i, j) + b + β.At(j, t+1)
				ζ[i][j][t] = x
				sum += math.Exp(x)
			}
		}
		// Normalize.
		for i := 0; i < N; i++ {
			for j := 0; j < N; j++ {
				ζ[i][j][t] -= math.Log(sum)
			}
		}
	}
	return
}

// Update model statistics.
// sequence is a matrix
func (hmm *HMM) Update(observations *matrix.Dense) (e error) {

	var α, β, γ *matrix.Dense
	var ζ [][][]float64
	var logProb float64

	// Compute  α, β, γ, ζ
	// TODO: compute α, β, concurrently using go routines.
	α, logProb, e = hmm.alpha(observations)
	if e != nil {
		return
	}
	β, e = hmm.beta(observations)
	if e != nil {
		return
	}
	// TODO: compute γ, ζ concurrently using go routines.
	// Can we compute ζ more efficiently using γ?
	γ, e = hmm.gamma(α, β)
	if e != nil {
		return
	}
	ζ, e = hmm.xi(observations, α, β)
	if e != nil {
		return
	}

	fmt.Println(α, β, γ, ζ, logProb)
	return
}

// ColumnAt returns a *matrix.Dense column that is a copy of the values at column c of the matrix.
// Column will panic with ErrIndexOutOfRange is c is not a valid column index.
func ColumnAt(d *matrix.Dense, c int) *matrix.Dense {

	sl := d.Column(c)
	nrows := len(sl)

	col := matrix.MustDense(matrix.ZeroDense(nrows, 1))

	for i := 0; i < nrows; i++ {
		col.Set(i, 0, sl[i])
	}

	return col
}

func (hmm *HMM) NumElements() int {
	return hmm.numElements
}
<|endoftext|>


<commit_before><commit_msg>Tweak<commit_after><|endoftext|>


<commit_before><commit_msg>Simplified the struct members given that transcoding is done elsewhere now<commit_after><|endoftext|>


<commit_before><commit_msg>devices: despise Fahrenheit a bit more<commit_after><|endoftext|>


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

package cmd

import (
	"bytes"
	"fmt"
	"hash/fnv"
	"net/http"
	"strings"
	"time"

	humanize "github.com/dustin/go-humanize"
	"github.com/fatih/color"
	"github.com/minio/cli"
	json "github.com/minio/mc/pkg/colorjson"
	"github.com/minio/mc/pkg/console"
	"github.com/minio/mc/pkg/probe"
	"github.com/minio/minio/pkg/madmin"
)

var adminTraceFlags = []cli.Flag{
	cli.BoolFlag{
		Name:  "verbose, v",
		Usage: "print verbose trace",
	},
	cli.BoolFlag{
		Name:  "all, a",
		Usage: "trace all traffic (including internode traffic between MinIO servers)",
	},
	cli.BoolFlag{
		Name:  "errors, e",
		Usage: "trace failed requests only",
	},
}

var adminTraceCmd = cli.Command{
	Name:            "trace",
	Usage:           "show http trace for MinIO server",
	Action:          mainAdminTrace,
	Before:          setGlobalsFromContext,
	Flags:           append(adminTraceFlags, globalFlags...),
	HideHelpCommand: true,
	CustomHelpTemplate: `NAME:
  {{.HelpName}} - {{.Usage}}

USAGE:
  {{.HelpName}} [FLAGS] TARGET

FLAGS:
  {{range .VisibleFlags}}{{.}}
  {{end}}
EXAMPLES:
  1. Show console trace for a MinIO server with alias 'play'
     {{.Prompt}} {{.HelpName}} -v -a play

  2. Show trace only for failed requests for a MinIO server with alias 'myminio'
    {{.Prompt}} {{.HelpName}} -v -e myminio
`,
}

const timeFormat = "15:04:05.000"

var (
	colors = []color.Attribute{color.FgCyan, color.FgWhite, color.FgYellow, color.FgGreen}
)

func checkAdminTraceSyntax(ctx *cli.Context) {
	if len(ctx.Args()) != 1 {
		cli.ShowCommandHelpAndExit(ctx, "trace", 1) // last argument is exit code
	}
}

// mainAdminTrace - the entry function of trace command
func mainAdminTrace(ctx *cli.Context) error {
	// Check for command syntax
	checkAdminTraceSyntax(ctx)
	verbose := ctx.Bool("verbose")
	all := ctx.Bool("all")
	errfltr := ctx.Bool("errors")
	aliasedURL := ctx.Args().Get(0)
	console.SetColor("Stat", color.New(color.FgYellow))

	console.SetColor("Request", color.New(color.FgCyan))
	console.SetColor("Method", color.New(color.Bold, color.FgWhite))
	console.SetColor("Host", color.New(color.Bold, color.FgGreen))
	console.SetColor("FuncName", color.New(color.Bold, color.FgGreen))

	console.SetColor("ReqHeaderKey", color.New(color.Bold, color.FgWhite))
	console.SetColor("RespHeaderKey", color.New(color.Bold, color.FgCyan))
	console.SetColor("HeaderValue", color.New(color.FgWhite))
	console.SetColor("RespStatus", color.New(color.Bold, color.FgYellow))
	console.SetColor("ErrStatus", color.New(color.Bold, color.FgRed))

	console.SetColor("Response", color.New(color.FgGreen))
	console.SetColor("Body", color.New(color.FgYellow))
	for _, c := range colors {
		console.SetColor(fmt.Sprintf("Node%d", c), color.New(c))
	}
	// Create a new MinIO Admin Client
	client, err := newAdminClient(aliasedURL)
	if err != nil {
		fatalIf(err.Trace(aliasedURL), "Cannot initialize admin client.")
		return nil
	}
	doneCh := make(chan struct{})
	defer close(doneCh)

	// Start listening on all trace activity.
	traceCh := client.ServiceTrace(all, errfltr, doneCh)
	for traceInfo := range traceCh {
		if traceInfo.Err != nil {
			fatalIf(probe.NewError(traceInfo.Err), "Cannot listen to http trace")
		}
		if verbose {
			printMsg(traceMessage{traceInfo})
			continue
		}
		printMsg(shortTrace(traceInfo))
	}
	return nil
}

// Short trace record
type shortTraceMsg struct {
	Host       string    `json:"host"`
	Time       time.Time `json:"time"`
	Client     string    `json:"client"`
	CallStats  callStats `json:"callStats"`
	FuncName   string    `json:"api"`
	Path       string    `json:"path"`
	Query      string    `json:"query"`
	StatusCode int       `json:"statusCode"`
	StatusMsg  string    `json:"statusMsg"`
}

type traceMessage struct {
	madmin.ServiceTraceInfo
}

type requestInfo struct {
	Time     time.Time         `json:"time"`
	Method   string            `json:"method"`
	Path     string            `json:"path,omitempty"`
	RawQuery string            `json:"rawQuery,omitempty"`
	Headers  map[string]string `json:"headers,omitempty"`
	Body     string            `json:"body,omitempty"`
}

type responseInfo struct {
	Time       time.Time         `json:"time"`
	Headers    map[string]string `json:"headers,omitempty"`
	Body       string            `json:"body,omitempty"`
	StatusCode int               `json:"statusCode,omitempty"`
}
type callStats struct {
	Rx       int           `json:"rx"`
	Tx       int           `json:"tx"`
	Duration time.Duration `json:"duration"`
	Ttfb     time.Duration `json:"timeToFirstByte"`
}

type trace struct {
	NodeName     string       `json:"host"`
	FuncName     string       `json:"api"`
	RequestInfo  requestInfo  `json:"request"`
	ResponseInfo responseInfo `json:"response"`
	CallStats    callStats    `json:"callStats"`
}

// return a struct with minimal trace info.
func shortTrace(ti madmin.ServiceTraceInfo) shortTraceMsg {
	s := shortTraceMsg{}
	t := ti.Trace
	s.Time = t.ReqInfo.Time
	if host, ok := t.ReqInfo.Headers["Host"]; ok {
		s.Host = strings.Join(host, "")
	}
	s.Path = t.ReqInfo.Path
	s.Query = t.ReqInfo.RawQuery
	s.FuncName = t.FuncName
	s.StatusCode = t.RespInfo.StatusCode
	s.StatusMsg = http.StatusText(t.RespInfo.StatusCode)
	cSlice := strings.Split(t.ReqInfo.Client, ":")
	s.Client = cSlice[0]
	s.CallStats.Duration = t.CallStats.Latency
	s.CallStats.Rx = t.CallStats.InputBytes
	s.CallStats.Tx = t.CallStats.OutputBytes
	return s
}

func (s shortTraceMsg) JSON() string {
	buf := &bytes.Buffer{}
	enc := json.NewEncoder(buf)
	enc.SetIndent("", " ")
	// Disable escaping special chars to display XML tags correctly
	enc.SetEscapeHTML(false)

	fatalIf(probe.NewError(enc.Encode(s)), "Unable to marshal into JSON.")
	return buf.String()
}

func (s shortTraceMsg) String() string {
	var hostStr string
	var b = &strings.Builder{}

	if s.Host != "" {
		hostStr = colorizedNodeName(s.Host)
	}
	fmt.Fprintf(b, "%s ", s.Time.Format(timeFormat))
	statusStr := console.Colorize("RespStatus", fmt.Sprintf("%d %s", s.StatusCode, s.StatusMsg))
	if s.StatusCode >= http.StatusBadRequest {
		statusStr = console.Colorize("ErrStatus", fmt.Sprintf("%d %s", s.StatusCode, s.StatusMsg))
	}
	fmt.Fprintf(b, "[%s] %s ", statusStr, console.Colorize("FuncName", s.FuncName))
	fmt.Fprintf(b, "%s%s", hostStr, s.Path)

	if s.Query != "" {
		fmt.Fprintf(b, "?%s ", s.Query)
	}
	fmt.Fprintf(b, " %s ", s.Client)
	spaces := 15 - len(s.Client)
	fmt.Fprintf(b, "%*s", spaces, " ")
	fmt.Fprint(b, console.Colorize("HeaderValue", fmt.Sprintf("  %2s", s.CallStats.Duration.Round(time.Microsecond))))
	spaces = 12 - len(fmt.Sprintf("%2s", s.CallStats.Duration.Round(time.Microsecond)))
	fmt.Fprintf(b, "%*s", spaces, " ")
	fmt.Fprint(b, console.Colorize("Stat", fmt.Sprintf(" 🠉 ")))
	fmt.Fprint(b, console.Colorize("HeaderValue", humanize.IBytes(uint64(s.CallStats.Rx))))
	fmt.Fprint(b, console.Colorize("Stat", fmt.Sprintf("  🠋 ")))
	fmt.Fprint(b, console.Colorize("HeaderValue", humanize.IBytes(uint64(s.CallStats.Tx))))

	return b.String()
}

// colorize node name
func colorizedNodeName(nodeName string) string {
	nodeHash := fnv.New32a()
	nodeHash.Write([]byte(nodeName))
	nHashSum := nodeHash.Sum32()
	idx := uint32(nHashSum) % uint32(len(colors))
	return console.Colorize(fmt.Sprintf("Node%d", colors[idx]), nodeName)
}

func (t traceMessage) JSON() string {
	rqHdrs := make(map[string]string)
	rspHdrs := make(map[string]string)
	rq := t.Trace.ReqInfo
	rs := t.Trace.RespInfo
	for k, v := range rq.Headers {
		rqHdrs[k] = strings.Join(v, " ")
	}
	for k, v := range rs.Headers {
		rspHdrs[k] = strings.Join(v, " ")
	}
	trc := trace{
		NodeName: t.Trace.NodeName,
		FuncName: t.Trace.FuncName,
		RequestInfo: requestInfo{
			Time:     rq.Time,
			Method:   rq.Method,
			Path:     rq.Path,
			RawQuery: rq.RawQuery,
			Body:     string(rq.Body),
			Headers:  rqHdrs,
		},
		ResponseInfo: responseInfo{
			Time:       rs.Time,
			Body:       string(rs.Body),
			Headers:    rspHdrs,
			StatusCode: rs.StatusCode,
		},
		CallStats: callStats{
			Duration: t.Trace.CallStats.Latency,
			Rx:       t.Trace.CallStats.InputBytes,
			Tx:       t.Trace.CallStats.OutputBytes,
			Ttfb:     t.Trace.CallStats.TimeToFirstByte,
		},
	}
	buf := &bytes.Buffer{}
	enc := json.NewEncoder(buf)
	enc.SetIndent("", " ")
	// Disable escaping special chars to display XML tags correctly
	enc.SetEscapeHTML(false)
	fatalIf(probe.NewError(enc.Encode(trc)), "Unable to marshal into JSON.")

	// strip off extra newline added by json encoder
	return strings.TrimSuffix(buf.String(), "\n")
}

func (t traceMessage) String() string {
	var nodeNameStr string
	var b = &strings.Builder{}

	trc := t.Trace
	if trc.NodeName != "" {
		nodeNameStr = fmt.Sprintf("%s ", colorizedNodeName(trc.NodeName))
	}

	ri := trc.ReqInfo
	rs := trc.RespInfo
	fmt.Fprintf(b, "%s%s", nodeNameStr, console.Colorize("Request", fmt.Sprintf("[REQUEST %s] ", trc.FuncName)))
	fmt.Fprintf(b, "%s\n", ri.Time.Format(timeFormat))
	fmt.Fprintf(b, "%s%s", nodeNameStr, console.Colorize("Method", fmt.Sprintf("%s %s", ri.Method, ri.Path)))
	if ri.RawQuery != "" {
		fmt.Fprintf(b, "?%s", ri.RawQuery)
	}
	fmt.Fprint(b, "\n")
	host, ok := ri.Headers["Host"]
	if ok {
		delete(ri.Headers, "Host")
	}
	hostStr := strings.Join(host, "")
	fmt.Fprintf(b, "%s%s", nodeNameStr, console.Colorize("Host", fmt.Sprintf("Host: %s\n", hostStr)))
	for k, v := range ri.Headers {
		fmt.Fprintf(b, "%s%s", nodeNameStr, console.Colorize("ReqHeaderKey",
			fmt.Sprintf("%s: ", k))+console.Colorize("HeaderValue", fmt.Sprintf("%s\n", strings.Join(v, ""))))
	}

	fmt.Fprintf(b, "%s%s", nodeNameStr, console.Colorize("Body", fmt.Sprintf("%s\n", string(ri.Body))))
	fmt.Fprintf(b, "%s%s", nodeNameStr, console.Colorize("Response", fmt.Sprintf("[RESPONSE] ")))
	fmt.Fprintf(b, "[%s] ", rs.Time.Format(timeFormat))
	fmt.Fprint(b, console.Colorize("Stat", fmt.Sprintf("[ Duration %2s  🠉 %s  🠋 %s ]\n", trc.CallStats.Latency.Round(time.Microsecond), humanize.IBytes(uint64(trc.CallStats.InputBytes)), humanize.IBytes(uint64(trc.CallStats.OutputBytes)))))

	statusStr := console.Colorize("RespStatus", fmt.Sprintf("%d %s", rs.StatusCode, http.StatusText(rs.StatusCode)))
	if rs.StatusCode != http.StatusOK {
		statusStr = console.Colorize("ErrStatus", fmt.Sprintf("%d %s", rs.StatusCode, http.StatusText(rs.StatusCode)))
	}
	fmt.Fprintf(b, "%s%s\n", nodeNameStr, statusStr)

	for k, v := range rs.Headers {
		fmt.Fprintf(b, "%s%s", nodeNameStr, console.Colorize("RespHeaderKey",
			fmt.Sprintf("%s: ", k))+console.Colorize("HeaderValue", fmt.Sprintf("%s\n", strings.Join(v, ""))))
	}
	fmt.Fprintf(b, "%s%s\n", nodeNameStr, console.Colorize("Body", string(rs.Body)))
	fmt.Fprint(b, nodeNameStr)
	return b.String()
}
<commit_msg>Fix arrow signs to use more common U+2191 & U+2193 (#2942)<commit_after>/*
 * MinIO Client (C) 2019 MinIO, Inc.
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

package cmd

import (
	"bytes"
	"fmt"
	"hash/fnv"
	"net/http"
	"strings"
	"time"

	humanize "github.com/dustin/go-humanize"
	"github.com/fatih/color"
	"github.com/minio/cli"
	json "github.com/minio/mc/pkg/colorjson"
	"github.com/minio/mc/pkg/console"
	"github.com/minio/mc/pkg/probe"
	"github.com/minio/minio/pkg/madmin"
)

var adminTraceFlags = []cli.Flag{
	cli.BoolFlag{
		Name:  "verbose, v",
		Usage: "print verbose trace",
	},
	cli.BoolFlag{
		Name:  "all, a",
		Usage: "trace all traffic (including internode traffic between MinIO servers)",
	},
	cli.BoolFlag{
		Name:  "errors, e",
		Usage: "trace failed requests only",
	},
}

var adminTraceCmd = cli.Command{
	Name:            "trace",
	Usage:           "show http trace for MinIO server",
	Action:          mainAdminTrace,
	Before:          setGlobalsFromContext,
	Flags:           append(adminTraceFlags, globalFlags...),
	HideHelpCommand: true,
	CustomHelpTemplate: `NAME:
  {{.HelpName}} - {{.Usage}}

USAGE:
  {{.HelpName}} [FLAGS] TARGET

FLAGS:
  {{range .VisibleFlags}}{{.}}
  {{end}}
EXAMPLES:
  1. Show console trace for a MinIO server with alias 'play'
     {{.Prompt}} {{.HelpName}} -v -a play

  2. Show trace only for failed requests for a MinIO server with alias 'myminio'
    {{.Prompt}} {{.HelpName}} -v -e myminio
`,
}

const timeFormat = "15:04:05.000"

var (
	colors = []color.Attribute{color.FgCyan, color.FgWhite, color.FgYellow, color.FgGreen}
)

func checkAdminTraceSyntax(ctx *cli.Context) {
	if len(ctx.Args()) != 1 {
		cli.ShowCommandHelpAndExit(ctx, "trace", 1) // last argument is exit code
	}
}

// mainAdminTrace - the entry function of trace command
func mainAdminTrace(ctx *cli.Context) error {
	// Check for command syntax
	checkAdminTraceSyntax(ctx)
	verbose := ctx.Bool("verbose")
	all := ctx.Bool("all")
	errfltr := ctx.Bool("errors")
	aliasedURL := ctx.Args().Get(0)
	console.SetColor("Stat", color.New(color.FgYellow))

	console.SetColor("Request", color.New(color.FgCyan))
	console.SetColor("Method", color.New(color.Bold, color.FgWhite))
	console.SetColor("Host", color.New(color.Bold, color.FgGreen))
	console.SetColor("FuncName", color.New(color.Bold, color.FgGreen))

	console.SetColor("ReqHeaderKey", color.New(color.Bold, color.FgWhite))
	console.SetColor("RespHeaderKey", color.New(color.Bold, color.FgCyan))
	console.SetColor("HeaderValue", color.New(color.FgWhite))
	console.SetColor("RespStatus", color.New(color.Bold, color.FgYellow))
	console.SetColor("ErrStatus", color.New(color.Bold, color.FgRed))

	console.SetColor("Response", color.New(color.FgGreen))
	console.SetColor("Body", color.New(color.FgYellow))
	for _, c := range colors {
		console.SetColor(fmt.Sprintf("Node%d", c), color.New(c))
	}
	// Create a new MinIO Admin Client
	client, err := newAdminClient(aliasedURL)
	if err != nil {
		fatalIf(err.Trace(aliasedURL), "Cannot initialize admin client.")
		return nil
	}
	doneCh := make(chan struct{})
	defer close(doneCh)

	// Start listening on all trace activity.
	traceCh := client.ServiceTrace(all, errfltr, doneCh)
	for traceInfo := range traceCh {
		if traceInfo.Err != nil {
			fatalIf(probe.NewError(traceInfo.Err), "Cannot listen to http trace")
		}
		if verbose {
			printMsg(traceMessage{traceInfo})
			continue
		}
		printMsg(shortTrace(traceInfo))
	}
	return nil
}

// Short trace record
type shortTraceMsg struct {
	Host       string    `json:"host"`
	Time       time.Time `json:"time"`
	Client     string    `json:"client"`
	CallStats  callStats `json:"callStats"`
	FuncName   string    `json:"api"`
	Path       string    `json:"path"`
	Query      string    `json:"query"`
	StatusCode int       `json:"statusCode"`
	StatusMsg  string    `json:"statusMsg"`
}

type traceMessage struct {
	madmin.ServiceTraceInfo
}

type requestInfo struct {
	Time     time.Time         `json:"time"`
	Method   string            `json:"method"`
	Path     string            `json:"path,omitempty"`
	RawQuery string            `json:"rawQuery,omitempty"`
	Headers  map[string]string `json:"headers,omitempty"`
	Body     string            `json:"body,omitempty"`
}

type responseInfo struct {
	Time       time.Time         `json:"time"`
	Headers    map[string]string `json:"headers,omitempty"`
	Body       string            `json:"body,omitempty"`
	StatusCode int               `json:"statusCode,omitempty"`
}
type callStats struct {
	Rx       int           `json:"rx"`
	Tx       int           `json:"tx"`
	Duration time.Duration `json:"duration"`
	Ttfb     time.Duration `json:"timeToFirstByte"`
}

type trace struct {
	NodeName     string       `json:"host"`
	FuncName     string       `json:"api"`
	RequestInfo  requestInfo  `json:"request"`
	ResponseInfo responseInfo `json:"response"`
	CallStats    callStats    `json:"callStats"`
}

// return a struct with minimal trace info.
func shortTrace(ti madmin.ServiceTraceInfo) shortTraceMsg {
	s := shortTraceMsg{}
	t := ti.Trace
	s.Time = t.ReqInfo.Time
	if host, ok := t.ReqInfo.Headers["Host"]; ok {
		s.Host = strings.Join(host, "")
	}
	s.Path = t.ReqInfo.Path
	s.Query = t.ReqInfo.RawQuery
	s.FuncName = t.FuncName
	s.StatusCode = t.RespInfo.StatusCode
	s.StatusMsg = http.StatusText(t.RespInfo.StatusCode)
	cSlice := strings.Split(t.ReqInfo.Client, ":")
	s.Client = cSlice[0]
	s.CallStats.Duration = t.CallStats.Latency
	s.CallStats.Rx = t.CallStats.InputBytes
	s.CallStats.Tx = t.CallStats.OutputBytes
	return s
}

func (s shortTraceMsg) JSON() string {
	buf := &bytes.Buffer{}
	enc := json.NewEncoder(buf)
	enc.SetIndent("", " ")
	// Disable escaping special chars to display XML tags correctly
	enc.SetEscapeHTML(false)

	fatalIf(probe.NewError(enc.Encode(s)), "Unable to marshal into JSON.")
	return buf.String()
}

func (s shortTraceMsg) String() string {
	var hostStr string
	var b = &strings.Builder{}

	if s.Host != "" {
		hostStr = colorizedNodeName(s.Host)
	}
	fmt.Fprintf(b, "%s ", s.Time.Format(timeFormat))
	statusStr := console.Colorize("RespStatus", fmt.Sprintf("%d %s", s.StatusCode, s.StatusMsg))
	if s.StatusCode >= http.StatusBadRequest {
		statusStr = console.Colorize("ErrStatus", fmt.Sprintf("%d %s", s.StatusCode, s.StatusMsg))
	}
	fmt.Fprintf(b, "[%s] %s ", statusStr, console.Colorize("FuncName", s.FuncName))
	fmt.Fprintf(b, "%s%s", hostStr, s.Path)

	if s.Query != "" {
		fmt.Fprintf(b, "?%s ", s.Query)
	}
	fmt.Fprintf(b, " %s ", s.Client)
	spaces := 15 - len(s.Client)
	fmt.Fprintf(b, "%*s", spaces, " ")
	fmt.Fprint(b, console.Colorize("HeaderValue", fmt.Sprintf("  %2s", s.CallStats.Duration.Round(time.Microsecond))))
	spaces = 12 - len(fmt.Sprintf("%2s", s.CallStats.Duration.Round(time.Microsecond)))
	fmt.Fprintf(b, "%*s", spaces, " ")
	fmt.Fprint(b, console.Colorize("Stat", fmt.Sprintf(" ↑ ")))
	fmt.Fprint(b, console.Colorize("HeaderValue", humanize.IBytes(uint64(s.CallStats.Rx))))
	fmt.Fprint(b, console.Colorize("Stat", fmt.Sprintf(" ↓ ")))
	fmt.Fprint(b, console.Colorize("HeaderValue", humanize.IBytes(uint64(s.CallStats.Tx))))

	return b.String()
}

// colorize node name
func colorizedNodeName(nodeName string) string {
	nodeHash := fnv.New32a()
	nodeHash.Write([]byte(nodeName))
	nHashSum := nodeHash.Sum32()
	idx := uint32(nHashSum) % uint32(len(colors))
	return console.Colorize(fmt.Sprintf("Node%d", colors[idx]), nodeName)
}

func (t traceMessage) JSON() string {
	rqHdrs := make(map[string]string)
	rspHdrs := make(map[string]string)
	rq := t.Trace.ReqInfo
	rs := t.Trace.RespInfo
	for k, v := range rq.Headers {
		rqHdrs[k] = strings.Join(v, " ")
	}
	for k, v := range rs.Headers {
		rspHdrs[k] = strings.Join(v, " ")
	}
	trc := trace{
		NodeName: t.Trace.NodeName,
		FuncName: t.Trace.FuncName,
		RequestInfo: requestInfo{
			Time:     rq.Time,
			Method:   rq.Method,
			Path:     rq.Path,
			RawQuery: rq.RawQuery,
			Body:     string(rq.Body),
			Headers:  rqHdrs,
		},
		ResponseInfo: responseInfo{
			Time:       rs.Time,
			Body:       string(rs.Body),
			Headers:    rspHdrs,
			StatusCode: rs.StatusCode,
		},
		CallStats: callStats{
			Duration: t.Trace.CallStats.Latency,
			Rx:       t.Trace.CallStats.InputBytes,
			Tx:       t.Trace.CallStats.OutputBytes,
			Ttfb:     t.Trace.CallStats.TimeToFirstByte,
		},
	}
	buf := &bytes.Buffer{}
	enc := json.NewEncoder(buf)
	enc.SetIndent("", " ")
	// Disable escaping special chars to display XML tags correctly
	enc.SetEscapeHTML(false)
	fatalIf(probe.NewError(enc.Encode(trc)), "Unable to marshal into JSON.")

	// strip off extra newline added by json encoder
	return strings.TrimSuffix(buf.String(), "\n")
}

func (t traceMessage) String() string {
	var nodeNameStr string
	var b = &strings.Builder{}

	trc := t.Trace
	if trc.NodeName != "" {
		nodeNameStr = fmt.Sprintf("%s ", colorizedNodeName(trc.NodeName))
	}

	ri := trc.ReqInfo
	rs := trc.RespInfo
	fmt.Fprintf(b, "%s%s", nodeNameStr, console.Colorize("Request", fmt.Sprintf("[REQUEST %s] ", trc.FuncName)))
	fmt.Fprintf(b, "%s\n", ri.Time.Format(timeFormat))
	fmt.Fprintf(b, "%s%s", nodeNameStr, console.Colorize("Method", fmt.Sprintf("%s %s", ri.Method, ri.Path)))
	if ri.RawQuery != "" {
		fmt.Fprintf(b, "?%s", ri.RawQuery)
	}
	fmt.Fprint(b, "\n")
	host, ok := ri.Headers["Host"]
	if ok {
		delete(ri.Headers, "Host")
	}
	hostStr := strings.Join(host, "")
	fmt.Fprintf(b, "%s%s", nodeNameStr, console.Colorize("Host", fmt.Sprintf("Host: %s\n", hostStr)))
	for k, v := range ri.Headers {
		fmt.Fprintf(b, "%s%s", nodeNameStr, console.Colorize("ReqHeaderKey",
			fmt.Sprintf("%s: ", k))+console.Colorize("HeaderValue", fmt.Sprintf("%s\n", strings.Join(v, ""))))
	}

	fmt.Fprintf(b, "%s%s", nodeNameStr, console.Colorize("Body", fmt.Sprintf("%s\n", string(ri.Body))))
	fmt.Fprintf(b, "%s%s", nodeNameStr, console.Colorize("Response", fmt.Sprintf("[RESPONSE] ")))
	fmt.Fprintf(b, "[%s] ", rs.Time.Format(timeFormat))
	fmt.Fprint(b, console.Colorize("Stat", fmt.Sprintf("[ Duration %2s  ↑ %s  ↓ %s ]\n", trc.CallStats.Latency.Round(time.Microsecond), humanize.IBytes(uint64(trc.CallStats.InputBytes)), humanize.IBytes(uint64(trc.CallStats.OutputBytes)))))

	statusStr := console.Colorize("RespStatus", fmt.Sprintf("%d %s", rs.StatusCode, http.StatusText(rs.StatusCode)))
	if rs.StatusCode != http.StatusOK {
		statusStr = console.Colorize("ErrStatus", fmt.Sprintf("%d %s", rs.StatusCode, http.StatusText(rs.StatusCode)))
	}
	fmt.Fprintf(b, "%s%s\n", nodeNameStr, statusStr)

	for k, v := range rs.Headers {
		fmt.Fprintf(b, "%s%s", nodeNameStr, console.Colorize("RespHeaderKey",
			fmt.Sprintf("%s: ", k))+console.Colorize("HeaderValue", fmt.Sprintf("%s\n", strings.Join(v, ""))))
	}
	fmt.Fprintf(b, "%s%s\n", nodeNameStr, console.Colorize("Body", string(rs.Body)))
	fmt.Fprint(b, nodeNameStr)
	return b.String()
}
<|endoftext|>