/
partition_stress.go
100 lines (91 loc) · 3.21 KB
/
partition_stress.go
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
package partition_stress
import (
"fmt"
"labix.org/v2/mgo"
"labix.org/v2/mgo/bson"
"time"
)
type indvPartitionInfo struct {
Id uint64 `bson:"_id"`
max bson.M `bson:"max"`
CreateTime time.Time `bson:"createTime"`
}
type partitionInfo struct {
NumPartitions uint64 `bson:"numPartitions"`
Partitions []indvPartitionInfo `bson:"partitions"`
Ok int `bson:"ok"`
}
// a Work used to add partitions to a partitioned
// collection. To work, -partition=true must be used when creating
// the benchmark, otherwise, this work will spit errors.
type AddPartitionWork struct {
DB *mgo.Database
Collname string
Interval time.Duration // interval between partition adds
}
// checks the createTime of the last partition, and if it happened
// longer than Interval defined in AddPartitionWork, then it adds a partition.
// For example, if a.Interval is set to one hour, and the last partition was created
// 61 minutes ago, this function will add a partition
func (a AddPartitionWork) Do(c chan<- interface{}) {
coll := a.DB.C(a.Collname)
var result partitionInfo
err := a.DB.Run(bson.M{"getPartitionInfo": coll.Name}, &result)
if err == nil {
var numPartitions uint64
numPartitions = result.NumPartitions
var lastCreateTime time.Time
lastCreateTime = result.Partitions[numPartitions-1].CreateTime
currentTime := time.Now()
difference := currentTime.Sub(lastCreateTime)
if difference > a.Interval {
var addPartitionResult bson.M
err = a.DB.Run(bson.M{"addPartition": coll.Name}, &addPartitionResult)
if err != nil {
fmt.Println("error while adding Partition, ", err)
}
}
} else {
fmt.Println("error while getting Partition info, ", err)
}
}
// closes the session used to add partitions
func (a AddPartitionWork) Close() {
}
// a Work used to drop partitions of a partitioned
// collection. To work, -partition=true must be used when creating
// the benchmark, otherwise, this work will spit errors.
type DropPartitionWork struct {
DB *mgo.Database
Collname string
Interval time.Duration // interval between partition adds
}
// checks the createTime of the first partition, and if it happened
// longer than Interval defined in DropPartitionWork, then it drops the first partition.
// For example, if a.Interval is set to six hours, and the first partition was created
// seven hours ago, this function will drop the first partition
func (a DropPartitionWork) Do(c chan<- interface{}) {
coll := a.DB.C(a.Collname)
var result partitionInfo
err := a.DB.Run(bson.M{"getPartitionInfo": coll.Name}, &result)
if err == nil {
var numPartitions uint64
numPartitions = result.NumPartitions
var firstCreateTime time.Time
firstCreateTime = result.Partitions[0].CreateTime
currentTime := time.Now()
difference := currentTime.Sub(firstCreateTime)
if numPartitions > 0 && difference > a.Interval {
firstID := result.Partitions[0].Id
var dropPartitionResult bson.M
err = a.DB.Run(bson.D{{"dropPartition", coll.Name}, {"id", firstID}}, &dropPartitionResult)
if err != nil {
fmt.Println("error while dropping Partition, ", err)
}
}
} else {
fmt.Println("error while getting Partition info, ", err)
}
}
func (a DropPartitionWork) Close() {
}