Compare commits
6 Commits
main
...
127f3db78f
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
127f3db78f | ||
|
|
bb5d8a154b | ||
|
|
7033e9f29c | ||
|
|
bb2c87dbd3 | ||
|
|
f78d0d1f59 | ||
|
|
72a976bdd0 |
185
src/adapt.go
Executable file → Normal file
185
src/adapt.go
Executable file → Normal file
@@ -8,7 +8,6 @@ import (
|
|||||||
"io"
|
"io"
|
||||||
"log"
|
"log"
|
||||||
"net"
|
"net"
|
||||||
"regexp"
|
|
||||||
"strconv"
|
"strconv"
|
||||||
)
|
)
|
||||||
|
|
||||||
@@ -16,133 +15,129 @@ func adapt(ctx context.Context, config Config, conn net.Conn) error {
|
|||||||
reader := bufio.NewReader(conn)
|
reader := bufio.NewReader(conn)
|
||||||
|
|
||||||
for ctx.Err() == nil {
|
for ctx.Err() == nil {
|
||||||
if err := adaptOne(ctx, config, conn, reader); err != nil {
|
message, err := readMessage(reader)
|
||||||
|
if err != nil {
|
||||||
|
if err == io.EOF {
|
||||||
|
return nil
|
||||||
|
}
|
||||||
return err
|
return err
|
||||||
}
|
}
|
||||||
|
|
||||||
|
if len(message) > 0 {
|
||||||
|
var reply string
|
||||||
|
switch message[0] {
|
||||||
|
case "hello":
|
||||||
|
reply = "*0\r\n"
|
||||||
|
case "client":
|
||||||
|
reply = "+OK\r\n"
|
||||||
|
case "ping":
|
||||||
|
reply = "+PONG\r\n"
|
||||||
|
case "set":
|
||||||
|
k := message[1].(string)
|
||||||
|
v := message[2]
|
||||||
|
reply = "+OK\r\n"
|
||||||
|
switch v2 := v.(type) {
|
||||||
|
case []byte:
|
||||||
|
if err := config.db.Put(ctx, 1, k, string(v2)); err != nil {
|
||||||
|
reply = "-" + err.Error() + "\r\n"
|
||||||
|
}
|
||||||
|
case string:
|
||||||
|
if err := config.db.Put(ctx, 1, k, v2); err != nil {
|
||||||
|
reply = "-" + err.Error() + "\r\n"
|
||||||
|
}
|
||||||
|
default:
|
||||||
|
log.Fatalf("not impl: set %T", v)
|
||||||
|
}
|
||||||
|
case "del":
|
||||||
|
k := message[1].(string)
|
||||||
|
reply = "+OK\r\n"
|
||||||
|
if err := config.db.Del(ctx, 1, k); err != nil {
|
||||||
|
reply = "-" + err.Error() + "\r\n"
|
||||||
|
}
|
||||||
|
case "get":
|
||||||
|
v, err := Select[any](ctx, 1, fmt.Sprint(message[1]))
|
||||||
|
if err != nil {
|
||||||
|
reply = "-" + err.Error()
|
||||||
|
} else if v == nil || *v == nil {
|
||||||
|
reply = "_\r\n" // RESP3 null
|
||||||
|
} else {
|
||||||
|
switch v2 := (*v).(type) {
|
||||||
|
case []byte:
|
||||||
|
reply = fmt.Sprintf("+%s\r\n", v2)
|
||||||
|
case string:
|
||||||
|
reply = fmt.Sprintf("+%s\r\n", v2)
|
||||||
|
default:
|
||||||
|
log.Fatalf("not impl: get type %T", v2)
|
||||||
|
return nil
|
||||||
|
}
|
||||||
|
}
|
||||||
|
default:
|
||||||
|
log.Fatalf("unhandled: %+v", message)
|
||||||
|
}
|
||||||
|
log.Printf("replying: %q", reply)
|
||||||
|
if _, err := fmt.Fprintf(conn, reply); err != nil {
|
||||||
|
return err
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
log.Println()
|
||||||
}
|
}
|
||||||
|
|
||||||
return io.EOF
|
return io.EOF
|
||||||
}
|
}
|
||||||
|
|
||||||
func adaptOne(ctx context.Context, config Config, conn io.Writer, reader *bufio.Reader) error {
|
func readMessage(reader *bufio.Reader) ([]any, error) {
|
||||||
raw, err := readMessage(reader)
|
|
||||||
if err != nil {
|
|
||||||
return err
|
|
||||||
} else if len(raw) == 0 {
|
|
||||||
return nil
|
|
||||||
}
|
|
||||||
log.Printf("routing: %q", raw)
|
|
||||||
|
|
||||||
return config.WithConn(parseHashKey(raw), func(forwardConn net.Conn) error {
|
|
||||||
if _, err := io.Copy(forwardConn, bytes.NewReader(raw)); err != nil {
|
|
||||||
return err
|
|
||||||
}
|
|
||||||
|
|
||||||
return readMessageTo(conn, bufio.NewReader(forwardConn))
|
|
||||||
})
|
|
||||||
}
|
|
||||||
|
|
||||||
func readMessage(reader *bufio.Reader) ([]byte, error) {
|
|
||||||
w := bytes.NewBuffer(nil)
|
|
||||||
err := readMessageTo(w, reader)
|
|
||||||
return w.Bytes(), err
|
|
||||||
}
|
|
||||||
|
|
||||||
func readMessageTo(w io.Writer, reader *bufio.Reader) error {
|
|
||||||
w2 := bufio.NewWriter(w)
|
|
||||||
defer w2.Flush()
|
|
||||||
err := _readMessageTo(w2, reader)
|
|
||||||
if err != nil {
|
|
||||||
return err
|
|
||||||
}
|
|
||||||
return w2.Flush()
|
|
||||||
}
|
|
||||||
|
|
||||||
func _readMessageTo(w io.Writer, reader *bufio.Reader) error {
|
|
||||||
firstLine, _, err := reader.ReadLine()
|
firstLine, _, err := reader.ReadLine()
|
||||||
|
log.Printf("%q", firstLine)
|
||||||
if err != nil {
|
if err != nil {
|
||||||
return err
|
return nil, err
|
||||||
}
|
}
|
||||||
firstLine = bytes.TrimSuffix(firstLine, []byte("\r\n"))
|
firstLine = bytes.TrimSuffix(firstLine, []byte("\r\n"))
|
||||||
if len(firstLine) == 0 {
|
if len(firstLine) == 0 {
|
||||||
return nil
|
return nil, nil
|
||||||
}
|
}
|
||||||
fmt.Fprintf(w, "%s\r\n", firstLine)
|
|
||||||
|
|
||||||
switch firstLine[0] {
|
switch firstLine[0] {
|
||||||
case '+': // simple string, like +OK
|
case '+': // simple string, like +OK
|
||||||
return nil
|
return []any{string(firstLine[1:])}, nil
|
||||||
case '-': // simple error, like -message
|
case '-': // simple error, like -message
|
||||||
return nil
|
return nil, fmt.Errorf("error: %s", firstLine[1:])
|
||||||
case ':': // number, like /[+-][0-9]+/
|
case ':': // number, like /[+-][0-9]+/
|
||||||
return err
|
firstLine = bytes.TrimPrefix(firstLine[1:], []byte("+"))
|
||||||
|
n, err := strconv.Atoi(string(firstLine[1:]))
|
||||||
|
return []any{n}, err
|
||||||
case '$': // long string, like $-1 for nil, like $LEN\r\nSTRING\r\n
|
case '$': // long string, like $-1 for nil, like $LEN\r\nSTRING\r\n
|
||||||
if firstLine[1] == '-' {
|
if firstLine[1] == '-' {
|
||||||
return nil
|
return []any{nil}, nil
|
||||||
}
|
}
|
||||||
firstLine = bytes.TrimPrefix(firstLine[1:], []byte("+"))
|
nextLine, _, err := reader.ReadLine()
|
||||||
n, err := strconv.Atoi(string(firstLine))
|
log.Printf("%q", nextLine)
|
||||||
if err != nil {
|
nextLine = bytes.TrimSuffix(nextLine, []byte("\r\n"))
|
||||||
return fmt.Errorf("num not a num: %q: %w", firstLine, err)
|
return []any{string(nextLine)}, err
|
||||||
}
|
case '*': // array, like *-1 for nil, like *4 for [1,2,3,4]
|
||||||
nextLine := make([]byte, n+2)
|
|
||||||
nAt := 0
|
|
||||||
for nAt < n+2 {
|
|
||||||
nMore, err := reader.Read(nextLine[nAt:])
|
|
||||||
if err != nil {
|
|
||||||
return fmt.Errorf("couldnt read %v more/%v bytes for long string: %w", n+2-nAt, n+2, err)
|
|
||||||
}
|
|
||||||
nAt += nMore
|
|
||||||
}
|
|
||||||
fmt.Fprintf(w, "%s", nextLine)
|
|
||||||
return err
|
|
||||||
case '*', '~', '%': // *=array ~=set %=map, like *-1 for nil, like *4 for [1,2,3,4]
|
|
||||||
n, err := strconv.Atoi(string(firstLine[1:]))
|
n, err := strconv.Atoi(string(firstLine[1:]))
|
||||||
if err != nil {
|
if err != nil {
|
||||||
return fmt.Errorf("arr not a num: %q: %w", firstLine, err)
|
return nil, err
|
||||||
} else if n == -1 {
|
} else if n == -1 {
|
||||||
return nil
|
return nil, nil
|
||||||
}
|
|
||||||
if isMap := firstLine[0] == '%'; isMap {
|
|
||||||
n *= 2
|
|
||||||
}
|
}
|
||||||
|
var result []any
|
||||||
for i := 0; i < n; i++ {
|
for i := 0; i < n; i++ {
|
||||||
err := _readMessageTo(w, reader)
|
more, err := readMessage(reader)
|
||||||
if err != nil {
|
if err != nil {
|
||||||
return err
|
return nil, err
|
||||||
}
|
}
|
||||||
|
result = append(result, more...)
|
||||||
}
|
}
|
||||||
return nil
|
return result, nil
|
||||||
case '_': // nil
|
case '_': // nil
|
||||||
return nil
|
return nil, nil
|
||||||
case '#': // boolean, like #t or #f
|
case '#': // boolean, like #t or #f
|
||||||
return nil
|
return []any{firstLine[1] == 't'}, nil
|
||||||
case ',': // double
|
case ',': // double
|
||||||
return nil
|
log.Fatal("not impl")
|
||||||
}
|
}
|
||||||
|
|
||||||
log.Fatalf("not impl: %q", firstLine)
|
log.Fatalf("not impl: %q", firstLine)
|
||||||
return nil
|
return nil, nil
|
||||||
}
|
|
||||||
|
|
||||||
var alpha = regexp.MustCompile(`[a-zA-Z]`)
|
|
||||||
var parts = regexp.MustCompile(`(.*?)\r\n*`)
|
|
||||||
|
|
||||||
func parseHashKey(raw []byte) string {
|
|
||||||
matches := parts.FindAllSubmatch(raw, 5)
|
|
||||||
if len(matches) == 0 {
|
|
||||||
return ""
|
|
||||||
}
|
|
||||||
for i := range matches {
|
|
||||||
if isTheCommand := alpha.Match(matches[i][0]); isTheCommand {
|
|
||||||
for j := i + 1; j < len(matches); j++ {
|
|
||||||
if matches[j][0][0] == byte('$') {
|
|
||||||
continue
|
|
||||||
}
|
|
||||||
return string(matches[j][1])
|
|
||||||
}
|
|
||||||
return string(matches[i][1])
|
|
||||||
}
|
|
||||||
}
|
|
||||||
return ""
|
|
||||||
}
|
}
|
||||||
|
|||||||
@@ -1,64 +0,0 @@
|
|||||||
package src
|
|
||||||
|
|
||||||
import (
|
|
||||||
"bufio"
|
|
||||||
"bytes"
|
|
||||||
"fmt"
|
|
||||||
"strings"
|
|
||||||
"testing"
|
|
||||||
)
|
|
||||||
|
|
||||||
func TestParseHashKey(t *testing.T) {
|
|
||||||
cases := map[string]string{
|
|
||||||
"+OK\r\n": "+OK",
|
|
||||||
"+string\r\n": "+string",
|
|
||||||
"*1\r\n$4\r\nping\r\n": "ping",
|
|
||||||
"*2\r\n$3\r\nget\r\n$1\r\nx\r\n": "x",
|
|
||||||
"*2\r\n$3\r\nset\r\n$1\r\nx\r\n$1\r\ny\r\n": "x",
|
|
||||||
}
|
|
||||||
|
|
||||||
for given, expect := range cases {
|
|
||||||
given, expect := given, expect
|
|
||||||
t.Run(fmt.Sprintf("%q", given), func(t *testing.T) {
|
|
||||||
got := parseHashKey([]byte(given))
|
|
||||||
if string(got) != expect {
|
|
||||||
t.Errorf("expected %q=>%q but got =>%q", given, expect, got)
|
|
||||||
}
|
|
||||||
})
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
func TestReadMessage(t *testing.T) {
|
|
||||||
cases := map[string]string{
|
|
||||||
"ok": "+OK\r\n",
|
|
||||||
"simple string": "+my string\r\n",
|
|
||||||
"error": "-my err\r\n",
|
|
||||||
"number": ":1\r\n",
|
|
||||||
"negative number": ":-1\r\n",
|
|
||||||
"positive number": ":+1\r\n",
|
|
||||||
"nil string": "$-1\r\n",
|
|
||||||
"long string": "$100\r\n" + strings.Repeat("a", 100) + "\r\n",
|
|
||||||
"longer string": "$10000\r\n" + strings.Repeat("a", 10000) + "\r\n",
|
|
||||||
"zero array": "*0\r\n",
|
|
||||||
"array": "*3\r\n+hello\r\n:world\r\n$1\r\n!\r\n",
|
|
||||||
"set": "~3\r\n+hello\r\n:world\r\n$1\r\n!\r\n",
|
|
||||||
"map": "%1\r\n+foo\r\n+bar\r\n",
|
|
||||||
"nil": "_\r\n",
|
|
||||||
"boolean": "#t\r\n",
|
|
||||||
"double": ",1.5\r\n",
|
|
||||||
// "empty": "", (eof)
|
|
||||||
}
|
|
||||||
|
|
||||||
for name, payload := range cases {
|
|
||||||
payload := payload
|
|
||||||
t.Run(name, func(t *testing.T) {
|
|
||||||
got, err := readMessage(bufio.NewReader(strings.NewReader(payload)))
|
|
||||||
if err != nil {
|
|
||||||
t.Fatal(err)
|
|
||||||
}
|
|
||||||
if !bytes.Equal(got, []byte(payload)) {
|
|
||||||
t.Errorf("expected %q but got %q", payload, got)
|
|
||||||
}
|
|
||||||
})
|
|
||||||
}
|
|
||||||
}
|
|
||||||
147
src/config.go
Executable file → Normal file
147
src/config.go
Executable file → Normal file
@@ -1,54 +1,43 @@
|
|||||||
package src
|
package src
|
||||||
|
|
||||||
import (
|
import (
|
||||||
"bufio"
|
|
||||||
"bytes"
|
|
||||||
"context"
|
"context"
|
||||||
"encoding/base64"
|
|
||||||
"encoding/json"
|
"encoding/json"
|
||||||
"fmt"
|
|
||||||
"hash/crc32"
|
|
||||||
"io"
|
|
||||||
"log"
|
|
||||||
"net"
|
|
||||||
"os"
|
"os"
|
||||||
"slices"
|
|
||||||
"strings"
|
|
||||||
"sync"
|
|
||||||
)
|
)
|
||||||
|
|
||||||
type Config struct {
|
type Config struct {
|
||||||
ctx context.Context
|
Listen string `json:"LISTEN"`
|
||||||
Listen string `json:"LISTEN"`
|
ConnectionDriver string `json:"CONNECTION_DRIVER"`
|
||||||
Forwards string `json:"FORWARDS"`
|
ConnectionString string `json:"CONNECTION_STRING"`
|
||||||
Hello string `json:"HELLO_B64"`
|
db DB
|
||||||
forwards []*sync.Pool
|
|
||||||
}
|
}
|
||||||
|
|
||||||
func NewConfig(ctx context.Context) (Config, error) {
|
func NewConfig(ctx context.Context) (Config, error) {
|
||||||
config := Config{
|
config, err := newConfig()
|
||||||
ctx: ctx,
|
if err != nil {
|
||||||
Listen: ":10000",
|
|
||||||
Forwards: "",
|
|
||||||
Hello: base64.StdEncoding.EncodeToString([]byte("*1\r\n$4\r\nping\r\n")),
|
|
||||||
}
|
|
||||||
|
|
||||||
if err := config.loadEnv(); err != nil {
|
|
||||||
return config, err
|
return config, err
|
||||||
}
|
}
|
||||||
|
|
||||||
if err := config.setForwards(); err != nil {
|
db, err := NewDB(ctx, config.ConnectionDriver, config.ConnectionString)
|
||||||
|
if err != nil {
|
||||||
return config, err
|
return config, err
|
||||||
}
|
}
|
||||||
|
config.db = db
|
||||||
|
|
||||||
return config, nil
|
return config, nil
|
||||||
}
|
}
|
||||||
|
|
||||||
func (c *Config) loadEnv() error {
|
func newConfig() (Config, error) {
|
||||||
b, _ := json.Marshal(*c)
|
config := Config{
|
||||||
|
Listen: ":10000",
|
||||||
|
ConnectionDriver: "sqlite",
|
||||||
|
ConnectionString: "/tmp/red-apter.db",
|
||||||
|
}
|
||||||
|
b, _ := json.Marshal(config)
|
||||||
var m map[string]any
|
var m map[string]any
|
||||||
if err := json.Unmarshal(b, &m); err != nil {
|
if err := json.Unmarshal(b, &m); err != nil {
|
||||||
return err
|
return config, err
|
||||||
}
|
}
|
||||||
for k := range m {
|
for k := range m {
|
||||||
if v := os.Getenv(k); v != "" {
|
if v := os.Getenv(k); v != "" {
|
||||||
@@ -56,102 +45,8 @@ func (c *Config) loadEnv() error {
|
|||||||
}
|
}
|
||||||
}
|
}
|
||||||
b2, _ := json.Marshal(m)
|
b2, _ := json.Marshal(m)
|
||||||
return json.Unmarshal(b2, c)
|
if err := json.Unmarshal(b2, &config); err != nil {
|
||||||
}
|
return config, err
|
||||||
|
}
|
||||||
func (c *Config) setForwards() error {
|
return config, nil
|
||||||
hello, err := base64.StdEncoding.DecodeString(c.Hello)
|
|
||||||
if err != nil {
|
|
||||||
return err
|
|
||||||
}
|
|
||||||
|
|
||||||
addresses := strings.Split(c.Forwards, ",")
|
|
||||||
addresses = slices.DeleteFunc(addresses, func(s string) bool { return s == "" })
|
|
||||||
if len(addresses) == 0 {
|
|
||||||
return fmt.Errorf("at least one $FORWARD required")
|
|
||||||
}
|
|
||||||
|
|
||||||
c.forwards = make([]*sync.Pool, len(addresses))
|
|
||||||
for i := range addresses {
|
|
||||||
address := addresses[i]
|
|
||||||
c.forwards[i] = &sync.Pool{
|
|
||||||
New: func() any {
|
|
||||||
conn, err := c.dial(address, hello)
|
|
||||||
if err != nil {
|
|
||||||
log.Printf("! failed to dial %q: %v", address, err)
|
|
||||||
}
|
|
||||||
return conn
|
|
||||||
},
|
|
||||||
}
|
|
||||||
}
|
|
||||||
return nil
|
|
||||||
}
|
|
||||||
|
|
||||||
func (c *Config) dial(address string, hello []byte) (net.Conn, error) {
|
|
||||||
if err := c.ctx.Err(); err != nil {
|
|
||||||
return nil, err
|
|
||||||
}
|
|
||||||
|
|
||||||
log.Printf("dialing %q with %q", address, hello)
|
|
||||||
v, err := (&net.Dialer{}).DialContext(c.ctx, "tcp", address)
|
|
||||||
if err != nil {
|
|
||||||
return nil, err
|
|
||||||
}
|
|
||||||
|
|
||||||
if _, err := v.Write([]byte(hello)); err != nil {
|
|
||||||
v.Close()
|
|
||||||
return nil, err
|
|
||||||
}
|
|
||||||
|
|
||||||
if raw, err := readMessage(bufio.NewReader(v)); err != nil {
|
|
||||||
v.Close()
|
|
||||||
return nil, err
|
|
||||||
} else if bytes.HasPrefix(raw, []byte("-")) {
|
|
||||||
v.Close()
|
|
||||||
return nil, fmt.Errorf("failed hello: %s", raw)
|
|
||||||
}
|
|
||||||
|
|
||||||
return v, nil
|
|
||||||
}
|
|
||||||
|
|
||||||
func (c Config) Close() {
|
|
||||||
for i := range c.forwards {
|
|
||||||
if c.forwards[i] != nil {
|
|
||||||
i := i
|
|
||||||
func() {
|
|
||||||
c.forwards[i].New = nil
|
|
||||||
for {
|
|
||||||
got := c.forwards[i].Get()
|
|
||||||
if got == nil {
|
|
||||||
break
|
|
||||||
}
|
|
||||||
got.(net.Conn).Close()
|
|
||||||
}
|
|
||||||
}()
|
|
||||||
}
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
func (c Config) WithConn(hashKey string, foo func(net.Conn) error) error {
|
|
||||||
hash := int(crc32.ChecksumIEEE([]byte(hashKey)))
|
|
||||||
hashIdx := hash % len(c.forwards)
|
|
||||||
forward := c.forwards[hashIdx]
|
|
||||||
|
|
||||||
log.Printf("acquire conn %v", hashIdx)
|
|
||||||
forwardCon := forward.Get()
|
|
||||||
if forwardCon == nil {
|
|
||||||
log.Printf("got a nil conn to %v", hashIdx)
|
|
||||||
return io.EOF
|
|
||||||
}
|
|
||||||
forwardConn := forwardCon.(net.Conn)
|
|
||||||
|
|
||||||
if err := foo(forwardConn); err != nil {
|
|
||||||
log.Printf("errored with conn %v: %v", hashIdx, err)
|
|
||||||
forwardConn.Close()
|
|
||||||
return err
|
|
||||||
}
|
|
||||||
|
|
||||||
log.Printf("release conn %v", hashIdx)
|
|
||||||
forward.Put(forwardConn)
|
|
||||||
return nil
|
|
||||||
}
|
}
|
||||||
|
|||||||
@@ -1,72 +0,0 @@
|
|||||||
package src
|
|
||||||
|
|
||||||
import (
|
|
||||||
"fmt"
|
|
||||||
"io"
|
|
||||||
"log"
|
|
||||||
"net"
|
|
||||||
"os"
|
|
||||||
"strconv"
|
|
||||||
"sync"
|
|
||||||
"testing"
|
|
||||||
)
|
|
||||||
|
|
||||||
func TestConfigLoadEnv(t *testing.T) {
|
|
||||||
c := Config{
|
|
||||||
Listen: "default",
|
|
||||||
Forwards: "default",
|
|
||||||
Hello: "default",
|
|
||||||
}
|
|
||||||
|
|
||||||
os.Setenv("LISTEN", "$LISTEN")
|
|
||||||
os.Setenv("FORWARDS", "")
|
|
||||||
os.Setenv("HELLO_B64", "$HELLO_B64")
|
|
||||||
|
|
||||||
if err := c.loadEnv(); err != nil {
|
|
||||||
t.Fatal(err)
|
|
||||||
}
|
|
||||||
|
|
||||||
if c.Listen != "$LISTEN" {
|
|
||||||
t.Error("listen", c.Listen)
|
|
||||||
}
|
|
||||||
if c.Forwards != "default" {
|
|
||||||
t.Error("forwards", c.Forwards)
|
|
||||||
}
|
|
||||||
if c.Hello != "$HELLO_B64" {
|
|
||||||
t.Error("hello", c.Hello)
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
func TestConfigWithConn(t *testing.T) {
|
|
||||||
log.SetOutput(io.Discard)
|
|
||||||
|
|
||||||
c := Config{}
|
|
||||||
|
|
||||||
seen0, seen1 := false, false
|
|
||||||
conn0, conn1 := &net.IPConn{}, &net.IPConn{}
|
|
||||||
pool0, pool1 := &sync.Pool{}, &sync.Pool{}
|
|
||||||
pool0.Put(conn0)
|
|
||||||
pool1.Put(conn1)
|
|
||||||
c.forwards = []*sync.Pool{pool0, pool1}
|
|
||||||
|
|
||||||
for i := 0; i < 10; i++ {
|
|
||||||
i := strconv.Itoa(i)
|
|
||||||
t.Run(i, func(t *testing.T) {
|
|
||||||
if err := c.WithConn(i, func(c net.Conn) error {
|
|
||||||
c, _ = c.(*net.IPConn)
|
|
||||||
seen0 = seen0 || c == conn0
|
|
||||||
seen1 = seen1 || c == conn1
|
|
||||||
if c != conn0 && c != conn1 {
|
|
||||||
return fmt.Errorf("got unexepcted conn %v", c)
|
|
||||||
}
|
|
||||||
return nil
|
|
||||||
}); err != nil {
|
|
||||||
t.Fatal(err)
|
|
||||||
}
|
|
||||||
})
|
|
||||||
}
|
|
||||||
|
|
||||||
if !seen0 || !seen1 {
|
|
||||||
t.Fatalf("didnt see both conns")
|
|
||||||
}
|
|
||||||
}
|
|
||||||
114
src/db.go
Normal file
114
src/db.go
Normal file
@@ -0,0 +1,114 @@
|
|||||||
|
package src
|
||||||
|
|
||||||
|
import (
|
||||||
|
"context"
|
||||||
|
"database/sql"
|
||||||
|
"encoding/json"
|
||||||
|
|
||||||
|
_ "github.com/lib/pq"
|
||||||
|
_ "modernc.org/sqlite"
|
||||||
|
)
|
||||||
|
|
||||||
|
type DB struct {
|
||||||
|
*sql.DB
|
||||||
|
}
|
||||||
|
|
||||||
|
var global DB
|
||||||
|
|
||||||
|
func NewDB(ctx context.Context, driver, conn string) (DB, error) {
|
||||||
|
sql, err := sql.Open(driver, conn)
|
||||||
|
if err != nil {
|
||||||
|
return DB{}, err
|
||||||
|
}
|
||||||
|
db := DB{sql}
|
||||||
|
|
||||||
|
if err := db.setup(ctx); err != nil {
|
||||||
|
db.Close()
|
||||||
|
return db, err
|
||||||
|
}
|
||||||
|
|
||||||
|
global = db
|
||||||
|
return db, nil
|
||||||
|
}
|
||||||
|
|
||||||
|
func (db DB) setup(ctx context.Context) error {
|
||||||
|
if err := db.PingContext(ctx); err != nil {
|
||||||
|
return err
|
||||||
|
}
|
||||||
|
if _, err := db.ExecContext(ctx, `
|
||||||
|
CREATE TABLE IF NOT EXISTS data (
|
||||||
|
database INTEGER
|
||||||
|
, key TEXT
|
||||||
|
, value TEXT
|
||||||
|
, PRIMARY KEY (database, key)
|
||||||
|
)
|
||||||
|
`); err != nil {
|
||||||
|
return err
|
||||||
|
}
|
||||||
|
return nil
|
||||||
|
}
|
||||||
|
|
||||||
|
func Select[V any](ctx context.Context, database int, k string) (*V, error) {
|
||||||
|
var v V
|
||||||
|
err := global.Get(ctx, database, k, &v)
|
||||||
|
if err != nil {
|
||||||
|
if err != sql.ErrNoRows {
|
||||||
|
return nil, nil
|
||||||
|
}
|
||||||
|
return nil, err
|
||||||
|
}
|
||||||
|
return &v, nil
|
||||||
|
}
|
||||||
|
|
||||||
|
func (db DB) Get(ctx context.Context, database int, k string, vPtr any) error {
|
||||||
|
if err := db.get(ctx, database, k, vPtr); err != nil {
|
||||||
|
if err != sql.ErrNoRows {
|
||||||
|
return err
|
||||||
|
}
|
||||||
|
}
|
||||||
|
return nil
|
||||||
|
}
|
||||||
|
|
||||||
|
func (db DB) get(ctx context.Context, database int, k string, vPtr any) error {
|
||||||
|
row := db.QueryRowContext(ctx, `
|
||||||
|
SELECT value
|
||||||
|
FROM data
|
||||||
|
WHERE database=$1 AND key=$2
|
||||||
|
`, database, k)
|
||||||
|
if err := row.Err(); err != nil {
|
||||||
|
return err
|
||||||
|
}
|
||||||
|
|
||||||
|
var v string
|
||||||
|
if err := row.Scan(&v); err != nil {
|
||||||
|
return err
|
||||||
|
}
|
||||||
|
|
||||||
|
if err := json.Unmarshal([]byte(v), vPtr); err != nil {
|
||||||
|
return err
|
||||||
|
}
|
||||||
|
|
||||||
|
return nil
|
||||||
|
}
|
||||||
|
|
||||||
|
func (db DB) Put(ctx context.Context, database int, k string, v any) error {
|
||||||
|
b, _ := json.Marshal(v)
|
||||||
|
_, err := db.ExecContext(ctx, `
|
||||||
|
INSERT INTO data (database, key, value)
|
||||||
|
VALUES ($1, $2, $3)
|
||||||
|
|
||||||
|
ON CONFLICT DO UPDATE
|
||||||
|
|
||||||
|
SET value = $3
|
||||||
|
WHERE database = $1 AND key = $2
|
||||||
|
`, database, k, string(b))
|
||||||
|
return err
|
||||||
|
}
|
||||||
|
|
||||||
|
func (db DB) Del(ctx context.Context, database int, k string) error {
|
||||||
|
_, err := db.ExecContext(ctx, `
|
||||||
|
DELETE FROM data
|
||||||
|
WHERE database=$1 AND key=$2
|
||||||
|
`, database, k)
|
||||||
|
return err
|
||||||
|
}
|
||||||
58
src/db_test.go
Normal file
58
src/db_test.go
Normal file
@@ -0,0 +1,58 @@
|
|||||||
|
package src_test
|
||||||
|
|
||||||
|
import (
|
||||||
|
"context"
|
||||||
|
"path"
|
||||||
|
"red-apter/src"
|
||||||
|
"testing"
|
||||||
|
"time"
|
||||||
|
)
|
||||||
|
|
||||||
|
func TestDBSqlite(t *testing.T) {
|
||||||
|
ctx, can := context.WithTimeout(context.Background(), time.Minute)
|
||||||
|
defer can()
|
||||||
|
|
||||||
|
db, err := src.NewDB(ctx, "sqlite", path.Join(t.TempDir(), "db"))
|
||||||
|
if err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
t.Cleanup(func() { db.Close() })
|
||||||
|
|
||||||
|
v := -1
|
||||||
|
if err := db.Get(ctx, 5, "k", &v); err != nil {
|
||||||
|
t.Error("failed get 404:", err)
|
||||||
|
} else if v != -1 {
|
||||||
|
t.Error("got 404:", v)
|
||||||
|
}
|
||||||
|
|
||||||
|
if err := db.Put(ctx, 5, "k", 6); err != nil {
|
||||||
|
t.Error("failed insert:", err)
|
||||||
|
} else if err := db.Put(ctx, 5, "k", 7); err != nil {
|
||||||
|
t.Error("failed update:", err)
|
||||||
|
}
|
||||||
|
|
||||||
|
if err := db.Get(ctx, 5, "k", &v); err != nil {
|
||||||
|
t.Error("failed get:", err)
|
||||||
|
} else if v != 7 {
|
||||||
|
t.Errorf("wrong get: (%T) %v but wanted 7", v, v)
|
||||||
|
}
|
||||||
|
|
||||||
|
if v, err := src.Select[int](ctx, 5, "k"); err != nil {
|
||||||
|
t.Error("failed select[int]:", err)
|
||||||
|
} else if v == nil || *v != 7 {
|
||||||
|
t.Errorf("wrong get: (%T) %v but wanted 7", v, v)
|
||||||
|
}
|
||||||
|
|
||||||
|
if err := db.Del(ctx, 5, "k"); err != nil {
|
||||||
|
t.Error("failed del:", err)
|
||||||
|
} else if err := db.Del(ctx, 5, "k"); err != nil {
|
||||||
|
t.Error("failed del 404:", err)
|
||||||
|
}
|
||||||
|
|
||||||
|
v = -1
|
||||||
|
if err := db.Get(ctx, 5, "k", &v); err != nil {
|
||||||
|
t.Error("failed get 410:", err)
|
||||||
|
} else if v != -1 {
|
||||||
|
t.Error("wrong get 410:", v)
|
||||||
|
}
|
||||||
|
}
|
||||||
11
src/listen.go
Executable file → Normal file
11
src/listen.go
Executable file → Normal file
@@ -8,8 +8,6 @@ import (
|
|||||||
)
|
)
|
||||||
|
|
||||||
func listen(ctx context.Context, config Config) error {
|
func listen(ctx context.Context, config Config) error {
|
||||||
defer log.Println("/listen()")
|
|
||||||
|
|
||||||
wg := &sync.WaitGroup{}
|
wg := &sync.WaitGroup{}
|
||||||
defer wg.Wait()
|
defer wg.Wait()
|
||||||
|
|
||||||
@@ -17,6 +15,8 @@ func listen(ctx context.Context, config Config) error {
|
|||||||
if err != nil {
|
if err != nil {
|
||||||
return err
|
return err
|
||||||
}
|
}
|
||||||
|
defer listener.Close()
|
||||||
|
|
||||||
wg.Add(1)
|
wg.Add(1)
|
||||||
go func() {
|
go func() {
|
||||||
defer wg.Done()
|
defer wg.Done()
|
||||||
@@ -49,12 +49,5 @@ func handle(ctx context.Context, config Config, conn net.Conn) {
|
|||||||
func _handle(ctx context.Context, config Config, conn net.Conn) error {
|
func _handle(ctx context.Context, config Config, conn net.Conn) error {
|
||||||
defer conn.Close()
|
defer conn.Close()
|
||||||
|
|
||||||
ctx, can := context.WithCancel(ctx)
|
|
||||||
defer can()
|
|
||||||
go func() {
|
|
||||||
<-ctx.Done()
|
|
||||||
conn.Close()
|
|
||||||
}()
|
|
||||||
|
|
||||||
return adapt(ctx, config, conn)
|
return adapt(ctx, config, conn)
|
||||||
}
|
}
|
||||||
|
|||||||
5
src/main.go
Executable file → Normal file
5
src/main.go
Executable file → Normal file
@@ -2,17 +2,14 @@ package src
|
|||||||
|
|
||||||
import (
|
import (
|
||||||
"context"
|
"context"
|
||||||
"log"
|
|
||||||
)
|
)
|
||||||
|
|
||||||
func Main(ctx context.Context) error {
|
func Main(ctx context.Context) error {
|
||||||
defer log.Println("/Main()")
|
|
||||||
|
|
||||||
config, err := NewConfig(ctx)
|
config, err := NewConfig(ctx)
|
||||||
if err != nil {
|
if err != nil {
|
||||||
return err
|
return err
|
||||||
}
|
}
|
||||||
defer config.Close()
|
defer config.db.Close()
|
||||||
|
|
||||||
if err := listen(ctx, config); err != nil && ctx.Err() == nil {
|
if err := listen(ctx, config); err != nil && ctx.Err() == nil {
|
||||||
return err
|
return err
|
||||||
|
|||||||
22
src/testdata/redises.sh
vendored
22
src/testdata/redises.sh
vendored
@@ -1,22 +0,0 @@
|
|||||||
#! /usr/bin/bash
|
|
||||||
|
|
||||||
bin=valkey-server
|
|
||||||
if ! which valkey-server; then
|
|
||||||
bin=redis-server
|
|
||||||
fi
|
|
||||||
|
|
||||||
d=$(mktemp -d)
|
|
||||||
mkdir -p $d/1
|
|
||||||
mkdir -p $d/2
|
|
||||||
$bin --dir $d/1 --port 60113 --logfile $d/1.log &
|
|
||||||
bg=${!}
|
|
||||||
cleanup() {
|
|
||||||
kill $bg
|
|
||||||
}
|
|
||||||
trap cleanup EXIT
|
|
||||||
(
|
|
||||||
sleep 5
|
|
||||||
echo FORWARDS=127.0.0.1:60113,127.0.0.1:60114
|
|
||||||
) &
|
|
||||||
$bin --dir $d/2 --port 60114 --logfile $d/2.log
|
|
||||||
|
|
||||||
Reference in New Issue
Block a user