@ -15,44 +15,26 @@
package main
import (
"context"
"errors"
"flag"
"fmt"
"net"
"net/url"
"os"
"strings"
"time"
"github.com/ IBM/sarama "
"github.com/ OpenIMSDK/tools/component "
"github.com/OpenIMSDK/tools/errs"
"github.com/go-zookeeper/zk"
"go.mongodb.org/mongo-driver/mongo"
"go.mongodb.org/mongo-driver/mongo/options"
"github.com/openimsdk/open-im-server/v3/pkg/common/config"
"github.com/minio/minio-go/v7"
"github.com/minio/minio-go/v7/pkg/credentials"
"github.com/redis/go-redis/v9"
"gopkg.in/yaml.v3"
)
const (
// defaultCfgPath is the default path of the configuration file.
defaultCfgPath = "../../../../../config/config.yaml"
minioHealthCheckDuration = 1
maxRetry = 300
componentStartErrCode = 6000
configErrCode = 6001
mongoConnTimeout = 30 * time . Second
)
const (
colorRed = 31
colorGreen = 32
colorYellow = 33
defaultCfgPath = "../../../../../config/config.yaml"
maxRetry = 300
componentStartErrCode = 6000
configErrCode = 6001
)
var (
@ -103,16 +85,16 @@ func main() {
for _ , check := range checks {
str , err := check . function ( )
if err != nil {
e ( fmt . Sprintf ( "Starting %s failed, %v" , check . name , err ) )
component . E ( fmt . Sprintf ( "Starting %s failed, %v" , check . name , err ) )
allSuccess = false
break
} else {
s ( fmt . Sprintf ( "%s connected successfully, %s" , check . name , str ) )
component . S ( fmt . Sprintf ( "%s connected successfully, %s" , check . name , str ) )
}
}
if allSuccess {
s ( "All components started successfully!" )
component . S ( "All components started successfully!" )
return
}
@ -120,19 +102,6 @@ func main() {
os . Exit ( 1 )
}
func exactIP ( urll string ) string {
u , _ := url . Parse ( urll )
host , _ , err := net . SplitHostPort ( u . Host )
if err != nil {
host = u . Host
}
if strings . HasSuffix ( host , ":" ) {
host = host [ 0 : len ( host ) - 1 ]
}
return host
}
// Helper function to get environment variable or default value
func getEnv ( key , fallback string ) string {
if value , exists := os . LookupEnv ( key ) ; exists {
@ -143,45 +112,23 @@ func getEnv(key, fallback string) string {
// checkMongo checks the MongoDB connection without retries
func checkMongo ( ) ( string , error ) {
uri := getEnv ( "MONGO_URI" , buildMongoURI ( ) )
ctx , cancel := context . WithTimeout ( context . Background ( ) , mongoConnTimeout )
defer cancel ( )
str := "ths addr is:" + strings . Join ( config . Config . Mongo . Address , "," )
client , err := mongo . Connect ( ctx , options . Client ( ) . ApplyURI ( uri ) )
if err != nil {
return "" , errs . Wrap ( errStr ( err , str ) )
mongo := & component . Mongo {
Address : config . Config . Mongo . Address ,
Database : config . Config . Mongo . Database ,
Username : config . Config . Mongo . Username ,
Password : config . Config . Mongo . Password ,
MaxPoolSize : config . Config . Mongo . MaxPoolSize ,
}
defer client . Disconnect ( context . Background ( ) )
ctx , cancel = context . WithTimeout ( context . Background ( ) , mongoConnTimeout )
defer cancel ( )
if err = client . Ping ( ctx , nil ) ; err != nil {
return "" , errs . Wrap ( errStr ( err , str ) )
uri , uriExist := os . LookupEnv ( "MONGO_URI" )
if uriExist {
mongo . URL = uri
}
return str , nil
}
// buildMongoURI constructs the MongoDB URI using configuration settings
func buildMongoURI ( ) string {
// Fallback to config if environment variables are not set
username := config . Config . Mongo . Username
password := config . Config . Mongo . Password
database := config . Config . Mongo . Database
maxPoolSize := config . Config . Mongo . MaxPoolSize
mongodbHosts := strings . Join ( config . Config . Mongo . Address , "," )
if username != "" && password != "" {
return fmt . Sprintf ( "mongodb://%s:%s@%s/%s?maxPoolSize=%d" ,
username , password , mongodbHosts , database , maxPoolSize )
str , err := component . CheckMongo ( mongo )
if err != nil {
return "" , err
}
return fmt . Sprintf ( "mongodb://%s/%s?maxPoolSize=%d" ,
mongodbHosts , database , maxPoolSize )
return str , nil
}
// checkMinio checks the MinIO connection
@ -191,52 +138,24 @@ func checkMinio() (string, error) {
return "" , nil
}
// Prioritize environment variables
endpoint := getEnv ( "MINIO_ENDPOINT" , config . Config . Object . Minio . Endpoint )
accessKeyID := getEnv ( "MINIO_ACCESS_KEY_ID" , config . Config . Object . Minio . AccessKeyID )
secretAccessKey := getEnv ( "MINIO_SECRET_ACCESS_KEY" , config . Config . Object . Minio . SecretAccessKey )
useSSL := getEnv ( "MINIO_USE_SSL" , "false" ) // Assuming SSL is not used by default
if endpoint == "" || accessKeyID == "" || secretAccessKey == "" {
return "" , ErrConfig . Wrap ( "MinIO configuration missing" )
}
// Parse endpoint URL to determine if SSL is enabled
u , err := url . Parse ( endpoint )
endpoint , err := getMinioAddr ( "MINIO_ENDPOINT" , "MINIO_ADDRESS" , "MINIO_PORT" , config . Config . Object . Minio . Endpoint )
if err != nil {
str := "the endpoint is:" + endpoint
return "" , errs . Wrap ( errStr ( err , str ) )
}
secure := u . Scheme == "https" || useSSL == "true"
// Initialize MinIO client
minioClient , err := minio . New ( u . Host , & minio . Options {
Creds : credentials . NewStaticV4 ( accessKeyID , secretAccessKey , "" ) ,
Secure : secure ,
} )
str := "ths addr is:" + u . Host
if err != nil {
strs := fmt . Sprintf ( "%v;host:%s,accessKeyID:%s,secretAccessKey:%s,Secure:%v" , err , u . Host , accessKeyID , secretAccessKey , secure )
return "" , errs . Wrap ( err , strs )
return "" , err
}
// Perform health check
cancel , err := minioClient . HealthCheck ( time . Duration ( minioHealthCheckDuration ) * time . Second )
if err != nil {
return "" , errs . Wrap ( errStr ( err , str ) )
minio := & component . Minio {
ApiURL : config . Config . Object . ApiURL ,
Endpoint : endpoint ,
AccessKeyID : getEnv ( "MINIO_ACCESS_KEY_ID" , config . Config . Object . Minio . AccessKeyID ) ,
SecretAccessKey : getEnv ( "MINIO_SECRET_ACCESS_KEY" , config . Config . Object . Minio . SecretAccessKey ) ,
SignEndpoint : config . Config . Object . Minio . SignEndpoint ,
UseSSL : getEnv ( "MINIO_USE_SSL" , "false" ) ,
}
defer cancel ( )
if minioClient . IsOffline ( ) {
str := fmt . Sprintf ( "Minio server is offline;%s" , str )
return "" , ErrComponentStart . Wrap ( str )
}
// Check for localhost in API URL and Minio SignEndpoint
if exactIP ( config . Config . Object . ApiURL ) == "127.0.0.1" || exactIP ( config . Config . Object . Minio . SignEndpoint ) == "127.0.0.1" {
return "" , ErrConfig . Wrap ( "apiURL or Minio SignEndpoint endpoint contain 127.0.0.1" )
str , err := component . CheckMinio ( minio )
if err != nil {
return "" , err
}
return str , nil
}
@ -247,76 +166,48 @@ func checkRedis() (string, error) {
username := getEnv ( "REDIS_USERNAME" , config . Config . Redis . Username )
password := getEnv ( "REDIS_PASSWORD" , config . Config . Redis . Password )
// Split address to handle multiple addresses for cluster setup
redisAddresses := strings . Split ( address , "," )
var redisClient redis . UniversalClient
if len ( redisAddresses ) > 1 {
// Use cluster client for multiple addresses
redisClient = redis . NewClusterClient ( & redis . ClusterOptions {
Addrs : redisAddresses ,
Username : username ,
Password : password ,
} )
} else {
// Use regular client for single address
redisClient = redis . NewClient ( & redis . Options {
Addr : redisAddresses [ 0 ] ,
Username : username ,
Password : password ,
} )
redis := & component . Redis {
Address : strings . Split ( address , "," ) ,
Username : username ,
Password : password ,
}
defer redisClient . Close ( )
// Ping Redis to check connectivity
_ , err := redisClient . Ping ( context . Background ( ) ) . Result ( )
str := "the addr is:" + strings . Join ( redisAddresses , "," )
addresses , err := getAddress ( "REDIS_ADDRESS" , "REDIS_PORT" , config . Config . Redis . Address )
if err != nil {
return "" , err s. Wrap ( errStr ( err , str ) )
return "" , err
}
redis . Address = addresses
str , err := component . CheckRedis ( redis )
if err != nil {
return "" , err
}
return str , nil
}
// checkZookeeper checks the Zookeeper connection
func checkZookeeper ( ) ( string , error ) {
// Prioritize environment variables
schema := getEnv ( "ZOOKEEPER_SCHEMA" , "digest" )
address := getEnv ( "ZOOKEEPER_ADDRESS" , strings . Join ( config . Config . Zookeeper . ZkAddr , "," ) )
username := getEnv ( "ZOOKEEPER_USERNAME" , config . Config . Zookeeper . Username )
password := getEnv ( "ZOOKEEPER_PASSWORD" , config . Config . Zookeeper . Password )
// Split addresses to handle multiple Zookeeper nodes
zookeeperAddresses := strings . Split ( address , "," )
zk := & component . Zookeeper {
Schema : getEnv ( "ZOOKEEPER_SCHEMA" , "digest" ) ,
ZkAddr : strings . Split ( address , "," ) ,
Username : getEnv ( "ZOOKEEPER_USERNAME" , config . Config . Zookeeper . Username ) ,
Password : getEnv ( "ZOOKEEPER_PASSWORD" , config . Config . Zookeeper . Password ) ,
}
// Connect to Zookeeper
str := "the addr is:" + address
c , eventChan , err := zk . Connect ( zookeeperAddresses , time . Second ) // Adjust the timeout as necessary
addresses , err := getAddress ( "ZOOKEEPER_ADDRESS" , "ZOOKEEPER_PORT" , config . Config . Zookeeper . ZkAddr )
if err != nil {
return "" , errs . Wrap ( errStr ( err , str ) )
}
timeout := time . After ( 5 * time . Second )
for {
select {
case event := <- eventChan :
if event . State == zk . StateConnected {
fmt . Println ( "Connected to Zookeeper" )
goto Connected
}
case <- timeout :
return "" , errs . Wrap ( errors . New ( "timeout waiting for Zookeeper connection" ) , "Zookeeper Addr: " + strings . Join ( config . Config . Zookeeper . ZkAddr , " " ) )
}
return "" , nil
}
Connected :
defer c . Close ( )
zk . ZkAddr = addresses
// Set authentication if username and password are provided
if username != "" && password != "" {
if err := c . AddAuth ( schema , [ ] byte ( username + ":" + password ) ) ; err != nil {
return "" , errs . Wrap ( errStr ( err , str ) )
}
str , err := component . CheckZookeeper ( zk )
if err != nil {
return "" , err
}
return str , nil
}
@ -327,24 +218,21 @@ func checkKafka() (string, error) {
password := getEnv ( "KAFKA_PASSWORD" , config . Config . Kafka . Password )
address := getEnv ( "KAFKA_ADDRESS" , strings . Join ( config . Config . Kafka . Addr , "," ) )
// Split addresses to handle multiple Kafka brokers
kafkaAddresses := strings . Split ( address , "," )
kafka := & component . Kafka {
Username : username ,
Password : password ,
Addr : strings . Split ( address , "," ) ,
}
// Configure Kafka client
cfg := sarama . NewConfig ( )
if username != "" && password != "" {
cfg . Net . SASL . Enable = true
cfg . Net . SASL . User = username
cfg . Net . SASL . Password = password
addresses , err := getAddress ( "KAFKA_ADDRESS" , "KAFKA_PORT" , config . Config . Kafka . Addr )
if err != nil {
return "" , nil
}
// Additional Kafka setup (e.g., TLS configuration) can be added here
// kafka.SetupTLSConfig(cfg)
kafka . Addr = addresses
// Create Kafka client
str := "the addr is:" + address
kafkaClient , err := sarama . NewClient ( kafkaAddresses , cfg )
str , kafkaClient , err := component . CheckKafka ( kafka )
if err != nil {
return "" , err s. Wrap ( errStr ( err , str ) )
return "" , err
}
defer kafkaClient . Close ( )
@ -379,22 +267,42 @@ func isTopicPresent(topic string, topics []string) bool {
return false
}
func colorPrint ( colorCode int , format string , a ... interface { } ) {
fmt . Printf ( "\x1b[%dm%s\x1b[0m\n" , colorCode , fmt . Sprintf ( format , a ... ) )
}
func errorPrint ( s string ) {
colorPrint ( colorRed , "%v" , s )
}
func successPrint ( s string ) {
colorPrint ( colorGreen , "%v" , s )
}
func getAddress ( key1 , key2 string , fallback [ ] string ) ( [ ] string , error ) {
address , addrExist := os . LookupEnv ( key1 )
port , portExist := os . LookupEnv ( key2 )
func warningPrint ( s string ) {
colorPrint ( colorYellow , "Warning: But %v" , s )
if addrExist && portExist {
addresses := strings . Split ( address , "," )
for i , addr := range addresses {
addresses [ i ] = addr + ":" + port
}
return addresses , nil
} else if ! addrExist && portExist {
result := make ( [ ] string , len ( config . Config . Redis . Address ) )
for i , addr := range config . Config . Redis . Address {
add := strings . Split ( addr , ":" )
result [ i ] = add [ 0 ] + ":" + port
}
return result , nil
} else if addrExist && ! portExist {
return nil , errs . Wrap ( errors . New ( "the ZOOKEEPER_PORT of minio is empty" ) )
}
return fallback , nil
}
func errStr ( err error , str string ) error {
return fmt . Errorf ( "%v;%s" , err , str )
}
func getMinioAddr ( key1 , key2 , key3 , fallback string ) ( string , error ) {
// Prioritize environment variables
endpoint := getEnv ( key1 , fallback )
address , addressExist := os . LookupEnv ( key2 )
port , portExist := os . LookupEnv ( key3 )
if portExist && addressExist {
endpoint = "http://" + address + ":" + port
} else if ! portExist && addressExist {
return "" , errs . Wrap ( errors . New ( "the MINIO_PORT of minio is empty" ) )
} else if portExist && ! addressExist {
arr := strings . Split ( config . Config . Object . Minio . Endpoint , ":" )
arr [ 2 ] = port
endpoint = strings . Join ( arr , ":" )
}
return endpoint , nil
}
Brabem
5 months ago
committed by
GitHub