Skip to content

Commit 78e13e7

Browse files
raflol33raflol33
authored
Merge pull request #1 from raflol33/codex/analyze-code-for-database-usage
Add round-robin multi-DB support
2 parents 06e4c07 + 2c66dc4 commit 78e13e7

File tree

9 files changed

+198
-125
lines changed

9 files changed

+198
-125
lines changed

‎config/config.go‎

Lines changed: 78 additions & 77 deletions
Original file line numberDiff line numberDiff line change
@@ -5,81 +5,82 @@ const Version = "heplify-server 1.60.1"
55
var Setting HeplifyServer
66

77
type HeplifyServer struct {
8-
HEPAddr string `default:"0.0.0.0:9060"`
9-
HEPTCPAddr string `default:""`
10-
HEPTLSAddr string `default:""`
11-
HEPWSAddr string `default:""`
12-
ESAddr string `default:""`
13-
ESDiscovery bool `default:"true"`
14-
HEPv2Enable bool `default:"true"`
15-
ESUser string `default:""`
16-
ESPass string `default:""`
17-
LokiURL string `default:""`
18-
LokiBulk int `default:"400"`
19-
LokiTimer int `default:"4"`
20-
LokiBuffer int `default:"100000"`
21-
LokiHEPFilter []int `default:"1,5,100"`
22-
LokiIPPortLabels bool `default:"false"`
23-
LokiFromToLabels bool `default:"false"`
24-
LokiCallIDLabels bool `default:"false"`
25-
LokiAllowOutOfOrder bool `default:"false"`
26-
LineprotoURL string `default:""`
27-
LineprotoBulk int `default:"400"`
28-
LineprotoTimer int `default:"4"`
29-
LineprotoBuffer int `default:"100000"`
30-
LineprotoHEPFilter []int `default:"1,5,100"`
31-
LineprotoIPPortLabels bool `default:"false"`
32-
ForceHEPPayload []int `default:""`
33-
PromAddr string `default:":9096"`
34-
PromTargetIP string `default:""`
35-
PromTargetName string `default:""`
36-
DBShema string `default:"homer5"`
37-
DBDriver string `default:"mysql"`
38-
DBAddr string `default:"localhost:3306"`
39-
DBSSLMode string `default:"disable"`
40-
DBUser string `default:"root"`
41-
DBPass string `default:""`
42-
DBDataTable string `default:"homer_data"`
43-
DBConfTable string `default:"homer_configuration"`
44-
DBBulk int `default:"400"`
45-
DBTimer int `default:"4"`
46-
DBBuffer int `default:"400000"`
47-
DBWorker int `default:"8"`
48-
DBRotate bool `default:"true"`
49-
DBPartLog string `default:"2h"`
50-
DBPartIsup string `default:"6h"`
51-
DBPartSip string `default:"2h"`
52-
DBPartQos string `default:"6h"`
53-
DBDropDays int `default:"14"`
54-
DBDropDaysCall int `default:"0"`
55-
DBDropDaysRegister int `default:"0"`
56-
DBDropDaysDefault int `default:"0"`
57-
DBDropOnStart bool `default:"false"`
58-
DBUsageProtection bool `default:"false"`
59-
DBUsageScheme string `default:"percentage"`
60-
DBPercentageUsage string `default:"80%"`
61-
DBMaxSize string `default:"20GB"`
62-
DBProcDropLimit int `default:"2"`
63-
Dedup bool `default:"false"`
64-
DiscardMethod []string `default:""`
65-
CensorMethod []string `default:""`
66-
AlegIDs []string `default:""`
67-
ForceALegID bool `default:"false"`
68-
CustomHeader []string `default:""`
69-
IgnoreCaseCH bool `default:"false"`
70-
SIPHeader []string `default:"ruri_user,ruri_domain,from_user,from_tag,to_user,to_tag,callid,cseq,method,user_agent"`
71-
LogDbg string `default:""`
72-
LogLvl string `default:"info"`
73-
LogStd bool `default:"false"`
74-
LogSys bool `default:"false"`
75-
Config string `default:"./heplify-server.toml"`
76-
ConfigHTTPAddr string `default:""`
77-
ConfigHTTPPW string `default:""`
78-
Version bool `default:"false"`
79-
ScriptEnable bool `default:"false"`
80-
ScriptEngine string `default:"lua"`
81-
ScriptFolder string `default:""`
82-
ScriptHEPFilter []int `default:"1,5,100"`
83-
TLSCertFolder string `default:"."`
84-
TLSMinVersion string `default:"1.2"`
8+
HEPAddr string `default:"0.0.0.0:9060"`
9+
HEPTCPAddr string `default:""`
10+
HEPTLSAddr string `default:""`
11+
HEPWSAddr string `default:""`
12+
ESAddr string `default:""`
13+
ESDiscovery bool `default:"true"`
14+
HEPv2Enable bool `default:"true"`
15+
ESUser string `default:""`
16+
ESPass string `default:""`
17+
LokiURL string `default:""`
18+
LokiBulk int `default:"400"`
19+
LokiTimer int `default:"4"`
20+
LokiBuffer int `default:"100000"`
21+
LokiHEPFilter []int `default:"1,5,100"`
22+
LokiIPPortLabels bool `default:"false"`
23+
LokiFromToLabels bool `default:"false"`
24+
LokiCallIDLabels bool `default:"false"`
25+
LokiAllowOutOfOrder bool `default:"false"`
26+
LineprotoURL string `default:""`
27+
LineprotoBulk int `default:"400"`
28+
LineprotoTimer int `default:"4"`
29+
LineprotoBuffer int `default:"100000"`
30+
LineprotoHEPFilter []int `default:"1,5,100"`
31+
LineprotoIPPortLabels bool `default:"false"`
32+
ForceHEPPayload []int `default:""`
33+
PromAddr string `default:":9096"`
34+
PromTargetIP string `default:""`
35+
PromTargetName string `default:""`
36+
DBShema string `default:"homer5"`
37+
DBDriver string `default:"mysql"`
38+
DBAddr string `default:"localhost:3306"`
39+
DBAddrs []string `default:""`
40+
DBSSLMode string `default:"disable"`
41+
DBUser string `default:"root"`
42+
DBPass string `default:""`
43+
DBDataTable string `default:"homer_data"`
44+
DBConfTable string `default:"homer_configuration"`
45+
DBBulk int `default:"400"`
46+
DBTimer int `default:"4"`
47+
DBBuffer int `default:"400000"`
48+
DBWorker int `default:"8"`
49+
DBRotate bool `default:"true"`
50+
DBPartLog string `default:"2h"`
51+
DBPartIsup string `default:"6h"`
52+
DBPartSip string `default:"2h"`
53+
DBPartQos string `default:"6h"`
54+
DBDropDays int `default:"14"`
55+
DBDropDaysCall int `default:"0"`
56+
DBDropDaysRegister int `default:"0"`
57+
DBDropDaysDefault int `default:"0"`
58+
DBDropOnStart bool `default:"false"`
59+
DBUsageProtection bool `default:"false"`
60+
DBUsageScheme string `default:"percentage"`
61+
DBPercentageUsage string `default:"80%"`
62+
DBMaxSize string `default:"20GB"`
63+
DBProcDropLimit int `default:"2"`
64+
Dedup bool `default:"false"`
65+
DiscardMethod []string `default:""`
66+
CensorMethod []string `default:""`
67+
AlegIDs []string `default:""`
68+
ForceALegID bool `default:"false"`
69+
CustomHeader []string `default:""`
70+
IgnoreCaseCH bool `default:"false"`
71+
SIPHeader []string `default:"ruri_user,ruri_domain,from_user,from_tag,to_user,to_tag,callid,cseq,method,user_agent"`
72+
LogDbg string `default:""`
73+
LogLvl string `default:"info"`
74+
LogStd bool `default:"false"`
75+
LogSys bool `default:"false"`
76+
Config string `default:"./heplify-server.toml"`
77+
ConfigHTTPAddr string `default:""`
78+
ConfigHTTPPW string `default:""`
79+
Version bool `default:"false"`
80+
ScriptEnable bool `default:"false"`
81+
ScriptEngine string `default:"lua"`
82+
ScriptFolder string `default:""`
83+
ScriptHEPFilter []int `default:"1,5,100"`
84+
TLSCertFolder string `default:"."`
85+
TLSMinVersion string `default:"1.2"`
8586
}

‎config/webconfig.go‎

Lines changed: 23 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -6,6 +6,7 @@ import (
66
"os"
77
"reflect"
88
"strconv"
9+
"strings"
910

1011
toml "github.com/pelletier/go-toml"
1112
)
@@ -52,7 +53,7 @@ func WebConfig(r *http.Request) (*HeplifyServer, error) {
5253
lineprotoBufferStr := r.FormValue("LineprotoBuffer")
5354
lineprotoBuffer, err := strconv.Atoi(lineprotoBufferStr)
5455
if err != nil || lineprotoBuffer < 0 || lineprotoBuffer > maxLineprotoBuffer {
55-
return nil, fmt.Errorf("Invalid LineprotoBuffer value: must be between 0 and %d", maxLineprotoBuffer)
56+
return nil, fmt.Errorf("Invalid LineprotoBuffer value: must be between 0 and %d", maxLineprotoBuffer)
5657
}
5758
webSetting.LineprotoBuffer = lineprotoBuffer
5859
DBShema := r.FormValue("DBShema")
@@ -70,6 +71,7 @@ func WebConfig(r *http.Request) (*HeplifyServer, error) {
7071
webSetting.DBConfTable = "homer_config"
7172
}
7273
webSetting.DBAddr = r.FormValue("DBAddr")
74+
webSetting.DBAddrs = parseDBAddrs(r.FormValue("DBAddrs"))
7375
webSetting.DBSSLMode = r.FormValue("DBSSLMode")
7476
webSetting.DBUser = r.FormValue("DBUser")
7577
DBPass := r.FormValue("DBPass")
@@ -133,6 +135,22 @@ func WebConfig(r *http.Request) (*HeplifyServer, error) {
133135
return &webSetting, nil
134136
}
135137

138+
func parseDBAddrs(raw string) []string {
139+
if raw == "" {
140+
return nil
141+
}
142+
143+
parts := strings.Split(raw, ",")
144+
addrs := make([]string, 0, len(parts))
145+
for _, part := range parts {
146+
addr := strings.TrimSpace(part)
147+
if addr != "" {
148+
addrs = append(addrs, addr)
149+
}
150+
}
151+
return addrs
152+
}
153+
136154
var WebForm = `
137155
<!DOCTYPE html>
138156
<html>
@@ -268,6 +286,10 @@ var WebForm = `
268286
<label>DBAddr</label>
269287
<input type="text" name="DBAddr" placeholder="{{.DBAddr}}" value="{{.DBAddr}}">
270288
</div>
289+
<div>
290+
<label>DBAddrs</label>
291+
<input type="text" name="DBAddrs" placeholder="{{.DBAddrs}}" value="{{.DBAddrs}}">
292+
</div>
271293
<div>
272294
<label>DBSSLMode</label>
273295
<input type="text" name="DBSSLMode" placeholder="{{.DBSSLMode}}" value="{{.DBSSLMode}}">

‎database/database.go‎

Lines changed: 22 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -14,10 +14,11 @@ import (
1414
type Database struct {
1515
H DBHandler
1616
Chan chan *decoder.HEP
17+
Addr string
1718
}
1819

1920
type DBHandler interface {
20-
setup() error
21+
setup(dbAddr string) error
2122
insert(chan *decoder.HEP)
2223
}
2324

@@ -29,10 +30,17 @@ func New(name string) *Database {
2930
}
3031

3132
return &Database{
32-
H: register[name],
33+
H: register[name],
34+
Addr: config.Setting.DBAddr,
3335
}
3436
}
3537

38+
func NewWithAddr(name string, dbAddr string) *Database {
39+
db := New(name)
40+
db.Addr = dbAddr
41+
return db
42+
}
43+
3644
func (d *Database) Run() error {
3745
driver := config.Setting.DBDriver
3846
shema := config.Setting.DBShema
@@ -53,7 +61,12 @@ func (d *Database) Run() error {
5361
}
5462
}
5563

56-
err := d.H.setup()
64+
addr := d.Addr
65+
if addr == "" {
66+
addr = config.Setting.DBAddr
67+
}
68+
69+
err := d.H.setup(addr)
5770
if err != nil {
5871
return err
5972
}
@@ -75,12 +88,15 @@ func (d *Database) End() {
7588
logp.Info("close %s channel", config.Setting.DBDriver)
7689
}
7790

78-
func ConnectString(dbName string) (string, error) {
91+
func ConnectString(dbName string, dbAddr string) (string, error) {
7992
var dsn string
8093
driver := config.Setting.DBDriver
81-
addr := strings.Split(config.Setting.DBAddr, ":")
94+
if dbAddr == "" {
95+
dbAddr = config.Setting.DBAddr
96+
}
97+
addr := strings.Split(dbAddr, ":")
8298
if len(addr) != 2 {
83-
return "", fmt.Errorf("wrong database connection format: %v, it should be localhost:3306", config.Setting.DBAddr)
99+
return "", fmt.Errorf("wrong database connection format: %v, it should be localhost:3306", dbAddr)
84100
}
85101
if (addr[1] == "3306" && driver == "postgres") ||
86102
addr[1] == "5432" && driver == "mysql" {

‎database/mock.go‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -15,7 +15,7 @@ type Mock struct {
1515
bulkCnt int
1616
}
1717

18-
func (m *Mock) setup() error {
18+
func (m *Mock) setup(_ string) error {
1919
m.db = new(syncmap.Map)
2020
m.bulkCnt = 200
2121
return nil

‎database/mysql.go‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -90,8 +90,8 @@ type MySQL struct {
9090
rtcBulkVal []byte
9191
}
9292

93-
func (m *MySQL) setup() error {
94-
cs, err := ConnectString(config.Setting.DBDataTable)
93+
func (m *MySQL) setup(dbAddr string) error {
94+
cs, err := ConnectString(config.Setting.DBDataTable, dbAddr)
9595
if err != nil {
9696
return err
9797
}

‎database/postgres.go‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -29,8 +29,8 @@ const (
2929
logCopy = "COPY hep_proto_100_default(sid,create_date,protocol_header,data_header,raw) FROM STDIN"
3030
)
3131

32-
func (p *Postgres) setup() error {
33-
cs, err := ConnectString(config.Setting.DBDataTable)
32+
func (p *Postgres) setup(dbAddr string) error {
33+
cs, err := ConnectString(config.Setting.DBDataTable, dbAddr)
3434
if err != nil {
3535
return err
3636
}

‎rotator/rotator.go‎

Lines changed: 7 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -53,6 +53,10 @@ type Rotator struct {
5353
}
5454

5555
func Setup(quit chan bool) *Rotator {
56+
return SetupWithAddr(quit, config.Setting.DBAddr)
57+
}
58+
59+
func SetupWithAddr(quit chan bool, dbAddr string) *Rotator {
5660
r := &Rotator{
5761
quit: quit,
5862
user: config.Setting.DBUser,
@@ -75,9 +79,9 @@ func Setup(quit chan bool) *Rotator {
7579
dropJob: cron.New(),
7680
}
7781

78-
r.rootDBAddr, _ = database.ConnectString("")
79-
r.confDBAddr, _ = database.ConnectString(config.Setting.DBConfTable)
80-
r.dataDBAddr, _ = database.ConnectString(config.Setting.DBDataTable)
82+
r.rootDBAddr, _ = database.ConnectString("", dbAddr)
83+
r.confDBAddr, _ = database.ConnectString(config.Setting.DBConfTable, dbAddr)
84+
r.dataDBAddr, _ = database.ConnectString(config.Setting.DBDataTable, dbAddr)
8185
if r.dropDaysCall == 0 {
8286
r.dropDaysCall = r.dropDays
8387
}

0 commit comments

Comments
 (0)