fix server lock up in high load scenarios

.github/workflows/integration-tests.yml

@@ -0,0 +1,27 @@
+name: Run integration tests
+
+on: [ push ]
+
+env:
+  GO_VERSION: 1.21.6
+
+jobs:
+  integration-tests:
+
+    runs-on: ubuntu-latest
+
+    steps:
+      - uses: actions/checkout@v4
+      - name: Setup Go
+        uses: actions/setup-go@v4
+        with:
+          go-version: ${{ env.GO_VERSION }}
+      - name: Build
+        run: go build -v ./...
+      - name: Run integration tests
+        run: go test -json ./test/integration > IntegrationTestResults-${{ env.GO_VERSION }}.json
+      - name: Upload Go integration tests results
+        uses: actions/upload-artifact@v4
+        with:
+          name: Go-results-${{ env.GO_VERSION }}
+          path: IntegrationTestResults-${{ env.GO_VERSION }}.json

cmd/cli/main.go

@@ -2,17 +2,17 @@ package main
 
 import (
 	"context"
-	"fmt"
 	"github.com/nitwhiz/omnilock/pkg/server"
+	"log"
 )
 
 func main() {
 	s, err := server.New(context.Background())
 
 	if err != nil {
-		fmt.Println(err)
+		log.Println(err)
 		return
 	}
 
-	s.Accept()
+	s.Run()
 }

pkg/client/client.go

@@ -5,6 +5,8 @@ import (
 	"context"
 	"github.com/nitwhiz/omnilock/pkg/id"
 	"net"
+	"sync"
+	"time"
 )
 
 type Client struct {
@@ -13,6 +15,7 @@ type Client struct {
 	conn    net.Conn
 	reader  *bufio.Reader
 	cmdChan chan<- *Command
+	mu      *sync.Mutex
 }
 
 func New(ctx context.Context, conn net.Conn, cmdChan chan<- *Command) *Client {
@@ -22,11 +25,24 @@ func New(ctx context.Context, conn net.Conn, cmdChan chan<- *Command) *Client {
 		conn:    conn,
 		reader:  bufio.NewReader(conn),
 		cmdChan: cmdChan,
+		mu:      &sync.Mutex{},
 	}
 
 	return &c
 }
 
+func (c *Client) Done() <-chan struct{} {
+	return c.ctx.Done()
+}
+
+func (c *Client) Lock() {
+	c.mu.Lock()
+}
+
+func (c *Client) Unlock() {
+	c.mu.Unlock()
+}
+
 func (c *Client) GetID() uint64 {
 	return c.id
 }
@@ -35,7 +51,13 @@ func (c *Client) GetContext() context.Context {
 	return c.ctx
 }
 
-func (c *Client) Write(b []byte) (int, error) {
+func (c *Client) Write(b []byte, timeout time.Duration) (int, error) {
+	err := c.conn.SetWriteDeadline(time.Now().Add(timeout))
+
+	if err != nil {
+		return 0, err
+	}
+
 	return c.conn.Write(b)
 }
 
@@ -46,9 +68,14 @@ func (c *Client) waitForCommand() bool {
 		return false
 	}
 
-	c.cmdChan <- &Command{
+	select {
+	case <-c.ctx.Done():
+		return false
+	case c.cmdChan <- &Command{
 		Client: c,
 		Cmd:    cmdString[:len(cmdString)-1],
+	}:
+		break
 	}
 
 	return true

pkg/lock/lock_table.go

@@ -0,0 +1,72 @@
+package lock
+
+import (
+	"context"
+	"github.com/nitwhiz/omnilock/pkg/util"
+	"sync"
+	"time"
+)
+
+type Table struct {
+	locks *util.Map[string, uint64]
+}
+
+func NewTable() *Table {
+	return &Table{
+		locks: util.NewMap[string, uint64](),
+	}
+}
+
+func (t *Table) TryLock(name string, clientId uint64) bool {
+	return t.locks.TryPut(name, clientId)
+}
+
+func (t *Table) Lock(ctx context.Context, name string, clientId uint64) bool {
+	if t.TryLock(name, clientId) {
+		return true
+	}
+
+	result := false
+
+	wg := &sync.WaitGroup{}
+	isRunning := false
+
+	go func() {
+		wg.Add(1)
+		defer wg.Done()
+
+		isRunning = true
+
+		for {
+			select {
+			case <-ctx.Done():
+				return
+			case <-time.After(time.Microsecond * 250):
+				if t.TryLock(name, clientId) {
+					result = true
+					return
+				}
+
+				break
+			}
+		}
+	}()
+
+	for !isRunning {
+		<-time.After(time.Millisecond)
+	}
+
+	wg.Wait()
+
+	return result
+}
+
+func (t *Table) Unlock(name string, clientId uint64) bool {
+	return t.locks.RemoveIf(name, func(v uint64) bool {
+		return clientId == v
+	})
+}
+
+func (t *Table) UnlockAll(clientId uint64) {
+	t.locks.RemoveByValue(clientId)
+}