HyperNet/Nexus
Archived
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

✨ Database allocator

Browse Source
This commit is contained in:
LittleSheep
2024-10-20 19:04:41 +08:00
parent e5a32bc05a
commit 4adbfe9c19
23 changed files with 536 additions and 29 deletions
Download Patch File Download Diff File Expand all files Collapse all files

39
pkg/internal/database/allocator.go Normal file
View File

@ -0,0 +1,39 @@
package database
import (
"fmt"
"github.com/samber/lo"
"github.com/spf13/viper"
"strings"
)
func AllocDatabase(name string) (string, error) {
// Disabled
if Kdb == nil {
return "host=localhost", fmt.Errorf("database feature is disabled")
}
var connString []string
connString = strings.Split(viper.GetString("database.dsn"), " ")
connString = lo.Filter(connString, func(item string, _ int) bool {
return !strings.HasPrefix(item, "dbname=")
})
name = viper.GetString("database.prefix") + name
var exists bool
if err := Kdb.QueryRow("SELECT EXISTS(SELECT datname FROM pg_catalog.pg_database WHERE datname = $1)", name).Scan(&exists); err != nil {
return strings.Join(connString, " "), nil
}
if !exists {
_, err := Kdb.Exec("CREATE DATABASE " + name)
if err != nil {
return strings.Join(connString, " "), err
}
}
connString = append(connString, "dbname="+name)
return strings.Join(connString, " "), nil
}

16
pkg/internal/database/connect.go Normal file
View File

@ -0,0 +1,16 @@
package database
import (
"database/sql"
_ "github.com/lib/pq"
)
var Kdb *sql.DB
func Connect(str string) (*sql.DB, error) {
db, err := sql.Open("postgres", str)
if err == nil {
Kdb = db
}
return db, err
}

15
pkg/internal/gap/net.go
View File

@ -1,15 +0,0 @@
package gap
import "net"
func GetOutboundIP() (net.IP, error) {
conn, err := net.Dial("udp", "1.1.1.1:80")
if err != nil {
return nil, err
} else {
defer conn.Close()
}
localAddr := conn.LocalAddr().(*net.UDPAddr)
return localAddr.IP, nil
}

20
pkg/internal/grpc/database.go Normal file
View File

@ -0,0 +1,20 @@
package grpc
import (
"context"
"git.solsynth.dev/hypernet/nexus/pkg/internal/database"
"git.solsynth.dev/hypernet/nexus/pkg/proto"
"google.golang.org/grpc/codes"
"google.golang.org/grpc/status"
)
func (v *Server) AllocDatabase(ctx context.Context, request *proto.AllocDatabaseRequest) (*proto.AllocDatabaseResponse, error) {
dsn, err := database.AllocDatabase(request.GetName())
if err != nil {
return nil, status.Errorf(codes.Internal, err.Error())
}
return &proto.AllocDatabaseResponse{
IsSuccess: true,
Dsn: dsn,
}, nil
}

4
pkg/internal/grpc/health.go
View File

@ -7,13 +7,13 @@ import (
health "google.golang.org/grpc/health/grpc_health_v1"
)
func (v *GrpcServer) Check(ctx context.Context, request *health.HealthCheckRequest) (*health.HealthCheckResponse, error) {
func (v *Server) Check(ctx context.Context, request *health.HealthCheckRequest) (*health.HealthCheckResponse, error) {
return &health.HealthCheckResponse{
Status: health.HealthCheckResponse_SERVING,
}, nil
}
func (v *GrpcServer) Watch(request *health.HealthCheckRequest, server health.Health_WatchServer) error {
func (v *Server) Watch(request *health.HealthCheckRequest, server health.Health_WatchServer) error {
for {
if server.Send(&health.HealthCheckResponse{
Status: health.HealthCheckResponse_SERVING,

11
pkg/internal/grpc/server.go
View File

@ -14,19 +14,22 @@ import (
health "google.golang.org/grpc/health/grpc_health_v1"
)
type GrpcServer struct {
type Server struct {
proto.UnimplementedDatabaseControllerServer
proto.UnimplementedStreamControllerServer
health.UnimplementedHealthServer
srv *grpc.Server
}
func NewServer() *GrpcServer {
server := &GrpcServer{
func NewServer() *Server {
server := &Server{
srv: grpc.NewServer(),
}
proto.RegisterServiceDirectoryServer(server.srv, &directory.ServiceRpcServer{})
proto.RegisterCommandControllerServer(server.srv, &directory.CommandRpcServer{})
proto.RegisterDatabaseControllerServer(server.srv, server)
proto.RegisterStreamControllerServer(server.srv, server)
health.RegisterHealthServer(server.srv, server)
@ -35,7 +38,7 @@ func NewServer() *GrpcServer {
return server
}
func (v *GrpcServer) Listen() error {
func (v *Server) Listen() error {
listener, err := net.Listen("tcp", viper.GetString("grpc_bind"))
if err != nil {
return err

18
pkg/internal/grpc/stream.go
View File

@ -3,27 +3,27 @@ package grpc
import (
"context"
"fmt"
"git.solsynth.dev/hypernet/nexus/pkg/http/ws"
"git.solsynth.dev/hypernet/nexus/pkg/internal/services"
"git.solsynth.dev/hypernet/nexus/pkg/proto"
"github.com/samber/lo"
)
func (v *GrpcServer) CountStreamConnection(ctx context.Context, request *proto.CountConnectionRequest) (*proto.CountConnectionResponse, error) {
out := services.ClientCount(uint(request.GetUserId()))
func (v *Server) CountStreamConnection(ctx context.Context, request *proto.CountConnectionRequest) (*proto.CountConnectionResponse, error) {
out := ws.ClientCount(uint(request.GetUserId()))
return &proto.CountConnectionResponse{
Count: int64(out),
}, nil
}
func (v *GrpcServer) PushStream(ctx context.Context, request *proto.PushStreamRequest) (*proto.PushStreamResponse, error) {
func (v *Server) PushStream(ctx context.Context, request *proto.PushStreamRequest) (*proto.PushStreamResponse, error) {
var cnt int
var success int
var errs []error
if request.UserId != nil {
cnt, success, errs = services.WebsocketPush(uint(request.GetUserId()), request.GetBody())
cnt, success, errs = ws.WebsocketPush(uint(request.GetUserId()), request.GetBody())
} else if request.ClientId != nil {
cnt, success, errs = services.WebsocketPushDirect(request.GetClientId(), request.GetBody())
cnt, success, errs = ws.WebsocketPushDirect(request.GetClientId(), request.GetBody())
} else {
return nil, fmt.Errorf("you must give one of the user id or client id")
}
@ -47,12 +47,12 @@ func (v *GrpcServer) PushStream(ctx context.Context, request *proto.PushStreamRe
}, nil
}
func (v *GrpcServer) PushStreamBatch(ctx context.Context, request *proto.PushStreamBatchRequest) (*proto.PushStreamResponse, error) {
func (v *Server) PushStreamBatch(ctx context.Context, request *proto.PushStreamBatchRequest) (*proto.PushStreamResponse, error) {
var cnt int
var success int
var errs []error
if len(request.UserId) != 0 {
cnt, success, errs = services.WebsocketPushBatch(
cnt, success, errs = ws.WebsocketPushBatch(
lo.Map(request.GetUserId(), func(item uint64, idx int) uint {
return uint(item)
},
@ -60,7 +60,7 @@ func (v *GrpcServer) PushStreamBatch(ctx context.Context, request *proto.PushStr
)
}
if len(request.ClientId) != 0 {
cCnt, cSuccess, cErrs := services.WebsocketPushBatchDirect(request.GetClientId(), request.GetBody())
cCnt, cSuccess, cErrs := ws.WebsocketPushBatchDirect(request.GetClientId(), request.GetBody())
cnt += cCnt
success += cSuccess
errs = append(errs, cErrs...)

95
pkg/internal/services/connections.go
View File

@ -1,95 +0,0 @@
package services
import (
"math/rand"
"sync"
"git.solsynth.dev/hypernet/nexus/pkg/internal/models"
"github.com/gofiber/contrib/websocket"
)
var (
wsMutex sync.Mutex
wsConn = make(map[uint]map[uint64]*websocket.Conn)
)
func ClientRegister(user models.Account, conn *websocket.Conn) uint64 {
wsMutex.Lock()
if wsConn[user.ID] == nil {
wsConn[user.ID] = make(map[uint64]*websocket.Conn)
}
clientId := rand.Uint64()
wsConn[user.ID][clientId] = conn
wsMutex.Unlock()
return clientId
}
func ClientUnregister(user models.Account, id uint64) {
wsMutex.Lock()
if wsConn[user.ID] == nil {
wsConn[user.ID] = make(map[uint64]*websocket.Conn)
}
delete(wsConn[user.ID], id)
wsMutex.Unlock()
}
func ClientCount(uid uint) int {
return len(wsConn[uid])
}
func WebsocketPush(uid uint, body []byte) (count int, success int, errs []error) {
for _, conn := range wsConn[uid] {
if err := conn.WriteMessage(1, body); err != nil {
errs = append(errs, err)
} else {
success++
}
count++
}
return
}
func WebsocketPushDirect(clientId uint64, body []byte) (count int, success int, errs []error) {
for _, m := range wsConn {
if conn, ok := m[clientId]; ok {
if err := conn.WriteMessage(1, body); err != nil {
errs = append(errs, err)
} else {
success++
}
count++
}
}
return
}
func WebsocketPushBatch(uidList []uint, body []byte) (count int, success int, errs []error) {
for _, uid := range uidList {
for _, conn := range wsConn[uid] {
if err := conn.WriteMessage(1, body); err != nil {
errs = append(errs, err)
} else {
success++
}
count++
}
}
return
}
func WebsocketPushBatchDirect(clientIdList []uint64, body []byte) (count int, success int, errs []error) {
for _, clientId := range clientIdList {
for _, m := range wsConn {
if conn, ok := m[clientId]; ok {
if err := conn.WriteMessage(1, body); err != nil {
errs = append(errs, err)
} else {
success++
}
count++
}
}
}
return
}