snapshot of https://github.com/ent/ent.git (master, 2026-09-15; blobs > 131072B omitted)

dialect/gremlin/internal/ws/conn.go

@@ -0,0 +1,315 @@
+// Copyright 2019-present Facebook Inc. All rights reserved.
+// This source code is licensed under the Apache 2.0 license found
+// in the LICENSE file in the root directory of this source tree.
+
+package ws
+
+import (
+	"bytes"
+	"context"
+	"errors"
+	"fmt"
+	"io"
+	"net/http"
+	"sync"
+	"time"
+
+	"entgo.io/ent/dialect/gremlin"
+	"entgo.io/ent/dialect/gremlin/encoding"
+	"entgo.io/ent/dialect/gremlin/encoding/graphson"
+
+	"github.com/gorilla/websocket"
+	"golang.org/x/sync/errgroup"
+)
+
+const (
+	// Time allowed to write a message to the peer.
+	writeWait = 5 * time.Second
+
+	// Time allowed to read the next pong message from the peer.
+	pongWait = 10 * time.Second
+
+	// Send pings to peer with this period. Must be less than pongWait.
+	pingPeriod = (pongWait * 9) / 10
+)
+
+type (
+	// A Dialer contains options for connecting to Gremlin server.
+	Dialer struct {
+		// Underlying websocket dialer.
+		websocket.Dialer
+
+		// Gremlin server basic auth credentials.
+		user, pass string
+	}
+
+	// Conn performs operations on a gremlin server.
+	Conn struct {
+		// Underlying websocket connection.
+		conn *websocket.Conn
+
+		// Credentials for basic authentication.
+		user, pass string
+
+		// Goroutine tracking.
+		ctx context.Context
+		grp *errgroup.Group
+
+		// Channel of outbound requests.
+		send chan io.Reader
+
+		// Map of in flight requests.
+		inflight sync.Map
+	}
+
+	// inflight tracks request state.
+	inflight struct {
+		// partially received data
+		frags []graphson.RawMessage
+
+		// response channel
+		result chan<- result
+	}
+
+	// represents an execution result.
+	result struct {
+		rsp *gremlin.Response
+		err error
+	}
+)
+
+var (
+	// DefaultDialer is a dialer with all fields set to the default values.
+	DefaultDialer = &Dialer{
+		Dialer: websocket.Dialer{
+			Proxy:            http.ProxyFromEnvironment,
+			HandshakeTimeout: 5 * time.Second,
+			WriteBufferSize:  8192,
+			ReadBufferSize:   8192,
+		},
+	}
+
+	// ErrConnClosed is returned by the Conn's Execute method when
+	// the underlying gremlin server connection is closed.
+	ErrConnClosed = errors.New("gremlin: server connection closed")
+
+	// ErrDuplicateRequest is returned by the Conns Execute method on
+	// request identifier key collision.
+	ErrDuplicateRequest = errors.New("gremlin: duplicate request")
+)
+
+// Dial creates a new connection by calling DialContext with a background context.
+func (d *Dialer) Dial(uri string) (*Conn, error) {
+	return d.DialContext(context.Background(), uri)
+}
+
+// DialContext creates a new Gremlin connection.
+func (d *Dialer) DialContext(ctx context.Context, uri string) (*Conn, error) {
+	c, rsp, err := d.Dialer.DialContext(ctx, uri, nil)
+	if err != nil {
+		return nil, fmt.Errorf("gremlin: dialing uri %s: %w", uri, err)
+	}
+	defer rsp.Body.Close()
+
+	conn := &Conn{
+		conn: c,
+		user: d.user,
+		pass: d.pass,
+		send: make(chan io.Reader),
+	}
+	conn.grp, conn.ctx = errgroup.WithContext(context.Background())
+
+	conn.grp.Go(conn.sender)
+	conn.grp.Go(conn.receiver)
+
+	return conn, nil
+}
+
+// Execute executes a request against a Gremlin server.
+func (c *Conn) Execute(ctx context.Context, req *gremlin.Request) (*gremlin.Response, error) {
+	// buffered result channel prevents receiver block on context cancellation
+	result := make(chan result, 1)
+
+	// request id must be unique across inflight request
+	if _, loaded := c.inflight.LoadOrStore(req.RequestID, &inflight{result: result}); loaded {
+		return nil, ErrDuplicateRequest
+	}
+
+	pr, pw := io.Pipe()
+	defer pr.Close()
+
+	// stream graphson encoding into request
+	c.grp.Go(func() error {
+		err := graphson.NewEncoder(pw).Encode(req)
+		if err != nil {
+			err = fmt.Errorf("encoding request: %w", err)
+		}
+		pw.CloseWithError(err)
+		return err
+	})
+
+	// local copy for single write
+	send := c.send
+
+	for {
+		select {
+		case <-c.ctx.Done():
+			c.inflight.Delete(req.RequestID)
+			return nil, ErrConnClosed
+		case <-ctx.Done():
+			c.inflight.Delete(req.RequestID)
+			return nil, ctx.Err()
+		case send <- pr:
+			send = nil
+		case result := <-result:
+			return result.rsp, result.err
+		}
+	}
+}
+
+// Close connection with a Gremlin server.
+func (c *Conn) Close() error {
+	c.grp.Go(func() error { return ErrConnClosed })
+	_ = c.grp.Wait()
+	return nil
+}
+
+func (c *Conn) sender() error {
+	pinger := time.NewTicker(pingPeriod)
+	defer pinger.Stop()
+
+	// closing connection terminates receiver
+	defer c.conn.Close()
+
+	for {
+		select {
+		case r := <-c.send:
+			// ensure write completes within a window
+			c.conn.SetWriteDeadline(time.Now().Add(writeWait))
+
+			// fetch next message writer
+			w, err := c.conn.NextWriter(websocket.BinaryMessage)
+			if err != nil {
+				return fmt.Errorf("getting message writer: %w", err)
+			}
+
+			// write mime header
+			if _, err := w.Write(encoding.GraphSON3Mime); err != nil {
+				return fmt.Errorf("writing mime header: %w", err)
+			}
+
+			// write request body
+			if _, err := io.Copy(w, r); err != nil {
+				return fmt.Errorf("writing request: %w", err)
+			}
+
+			// finish message write
+			if err := w.Close(); err != nil {
+				return fmt.Errorf("closing message writer: %w", err)
+			}
+		case <-c.ctx.Done():
+			// connection closing
+			return c.conn.WriteControl(
+				websocket.CloseMessage,
+				websocket.FormatCloseMessage(websocket.CloseNormalClosure, ""),
+				time.Time{},
+			)
+		case <-pinger.C:
+			// periodic connection keepalive
+			if err := c.conn.WriteControl(websocket.PingMessage, nil, time.Now().Add(writeWait)); err != nil {
+				return fmt.Errorf("writing ping message: %w", err)
+			}
+		}
+	}
+}
+
+func (c *Conn) receiver() error {
+	// handle keepalive responses
+	c.conn.SetReadDeadline(time.Now().Add(pongWait))
+	c.conn.SetPongHandler(func(string) error {
+		return c.conn.SetReadDeadline(time.Now().Add(pongWait))
+	})
+
+	// complete all in flight requests on termination
+	defer c.inflight.Range(func(id, ifr any) bool {
+		ifr.(*inflight).result <- result{err: ErrConnClosed}
+		c.inflight.Delete(id)
+		return true
+	})
+
+	for {
+		// rely on sender connection close during termination
+		_, r, err := c.conn.NextReader()
+		if err != nil {
+			return fmt.Errorf("writing ping message: %w", err)
+		}
+
+		// decode received response
+		var rsp gremlin.Response
+		if err := graphson.NewDecoder(r).Decode(&rsp); err != nil {
+			return fmt.Errorf("reading response: %w", err)
+		}
+
+		ifr, ok := c.inflight.Load(rsp.RequestID)
+		if !ok {
+			// context cancellation aborts inflight requests
+			continue
+		}
+
+		// handle incoming response
+		if done := c.receive(ifr.(*inflight), &rsp); done {
+			// stop tracking finished requests
+			c.inflight.Delete(rsp.RequestID)
+		}
+	}
+}
+
+func (c *Conn) receive(ifr *inflight, rsp *gremlin.Response) bool {
+	result := result{rsp: rsp}
+	switch rsp.Status.Code {
+	case gremlin.StatusSuccess:
+		// quickly handle non fragmented responses
+		if ifr.frags == nil {
+			break
+		}
+		// handle fragment
+		fallthrough
+	case gremlin.StatusPartialContent:
+		// append received fragment
+		var frag []graphson.RawMessage
+		if err := graphson.Unmarshal(rsp.Result.Data, &frag); err != nil {
+			result.err = fmt.Errorf("decoding response fragment: %w", err)
+			break
+		}
+		ifr.frags = append(ifr.frags, frag...)
+
+		// partial response requires additional fragments
+		if rsp.Status.Code == gremlin.StatusPartialContent {
+			return false
+		}
+
+		// reassemble fragmented response
+		if rsp.Result.Data, result.err = graphson.Marshal(ifr.frags); result.err != nil {
+			result.err = fmt.Errorf("assembling fragmented response: %w", result.err)
+		}
+	case gremlin.StatusAuthenticate:
+		// receiver should never block
+		c.grp.Go(func() error {
+			var buf bytes.Buffer
+			if err := graphson.NewEncoder(&buf).Encode(
+				gremlin.NewAuthRequest(rsp.RequestID, c.user, c.pass),
+			); err != nil {
+				return fmt.Errorf("encoding auth request: %w", err)
+			}
+			select {
+			case c.send <- &buf:
+			case <-c.ctx.Done():
+			}
+			return c.ctx.Err()
+		})
+		return false
+	}
+
+	ifr.result <- result
+	return true
+}