Congestion control

client.go

@@ -101,7 +101,7 @@ type client struct {
 	newRegionClientFn func(string, region.ClientType, int, time.Duration,
 		string, time.Duration, compression.Codec,
 		func(ctx context.Context, network, addr string) (net.Conn, error),
-		*slog.Logger) hrpc.RegionClient
+		*slog.Logger, int, int) hrpc.RegionClient
 
 	compressionCodec compression.Codec
 
@@ -112,6 +112,9 @@ type client struct {
 	regionDialer func(ctx context.Context, network, addr string) (net.Conn, error)
 	// logger that could be defined by user
 	logger *slog.Logger
+
+	minWindowSize int
+	maxWindowSize int
 }
 
 // NewClient creates a new HBase client.
@@ -313,6 +316,26 @@ func Logger(logger *slog.Logger) Option {
 	}
 }
 
+// CongestionControl enables per-regionserver congestion control,
+// limiting the number of outstanding Get/Scan requests. The window
+// size will adjust larged when seeing successful results and adjust
+// smaller when seeing errors within the bounds of the minWindowSize
+// and maxWindowSize.
+func CongestionControl(minWindowSize, maxWindowSize int) Option {
+	return func(c *client) {
+		if minWindowSize > maxWindowSize {
+			panic(fmt.Errorf("minWindowSize is greater than maxMindowSize: %d and %d",
+				minWindowSize, maxWindowSize))
+		}
+		if minWindowSize < 0 || maxWindowSize < 0 {
+			panic(fmt.Errorf("minWindowSize or maxWindowSize are negative: %d and %d",
+				minWindowSize, maxWindowSize))
+		}
+		c.minWindowSize = minWindowSize
+		c.maxWindowSize = maxWindowSize
+	}
+}
+
 // Close closes connections to hbase master and regionservers
 func (c *client) Close() {
 	c.closeOnce.Do(func() {

debug_state_test.go

@@ -29,6 +29,8 @@ func TestDebugStateSanity(t *testing.T) {
 		client.compressionCodec,
 		nil,
 		slog.Default(),
+		0,
+		0,
 	)
 	newClientFn := func() hrpc.RegionClient {
 		return regClient

integration_test.go

@@ -415,79 +415,96 @@ func TestPutMultipleCells(t *testing.T) {
 }
 
 func TestMultiplePutsGetsSequentially(t *testing.T) {
-	const num_ops = 100
-	keyPrefix := "row3"
-	headers := map[string][]string{"cf": nil}
-	c := gohbase.NewClient(*host, gohbase.FlushInterval(time.Millisecond))
-	defer c.Close()
-	err := performNPuts(keyPrefix, num_ops)
-	if err != nil {
-		t.Errorf("Put returned an error: %v", err)
-	}
-	for i := num_ops - 1; i >= 0; i-- {
-		key := keyPrefix + fmt.Sprintf("%d", i)
-		get, err := hrpc.NewGetStr(context.Background(), table, key, hrpc.Families(headers))
-		rsp, err := c.Get(get)
-		if err != nil {
-			t.Errorf("Get returned an error: %v", err)
-		}
-		if len(rsp.Cells) != 1 {
-			t.Errorf("Incorrect number of cells returned by Get: %d", len(rsp.Cells))
-		}
-		rsp_value := rsp.Cells[0].Value
-		if !bytes.Equal(rsp_value, []byte(fmt.Sprintf("%d", i))) {
-			t.Errorf("Get returned an incorrect result. Expected: %v, Got: %v",
-				[]byte(fmt.Sprintf("%d", i)), rsp_value)
-		}
+	for name, opts := range map[string][]gohbase.Option{
+		"no_congestion_control": nil,
+		"window20":              {gohbase.CongestionControl(1, 20)},
+		"window100":             {gohbase.CongestionControl(1, 100)},
+	} {
+		t.Run(name, func(t *testing.T) {
+			const num_ops = 100
+			keyPrefix := "row3"
+			headers := map[string][]string{"cf": nil}
+			c := gohbase.NewClient(*host, append(opts, gohbase.FlushInterval(time.Millisecond))...)
+			defer c.Close()
+			err := performNPuts(keyPrefix, num_ops)
+			if err != nil {
+				t.Errorf("Put returned an error: %v", err)
+			}
+			for i := num_ops - 1; i >= 0; i-- {
+				key := keyPrefix + fmt.Sprintf("%d", i)
+				get, err := hrpc.NewGetStr(context.Background(), table, key, hrpc.Families(headers))
+				rsp, err := c.Get(get)
+				if err != nil {
+					t.Errorf("Get returned an error: %v", err)
+				}
+				if len(rsp.Cells) != 1 {
+					t.Errorf("Incorrect number of cells returned by Get: %d", len(rsp.Cells))
+				}
+				rsp_value := rsp.Cells[0].Value
+				if !bytes.Equal(rsp_value, []byte(fmt.Sprintf("%d", i))) {
+					t.Errorf("Get returned an incorrect result. Expected: %v, Got: %v",
+						[]byte(fmt.Sprintf("%d", i)), rsp_value)
+				}
+			}
+		})
 	}
 }
 
 func TestMultiplePutsGetsParallel(t *testing.T) {
-	c := gohbase.NewClient(*host)
-	defer c.Close()
-
-	const n = 1000
-	var wg sync.WaitGroup
-	for i := 0; i < n; i++ {
-		key := fmt.Sprintf("%s_%d", t.Name(), i)
-		wg.Add(1)
-		go func() {
-			if err := insertKeyValue(c, key, "cf", []byte(key)); err != nil {
-				t.Error(key, err)
-			}
-			wg.Done()
-		}()
-	}
-	wg.Wait()
+	for name, opts := range map[string][]gohbase.Option{
+		"no_congestion_control": nil,
+		"window20":              {gohbase.CongestionControl(1, 20)},
+		"window1000":            {gohbase.CongestionControl(1, 1000)},
+	} {
+		t.Run(name, func(t *testing.T) {
+			c := gohbase.NewClient(*host, opts...)
+			defer c.Close()
 
-	// All puts are complete. Now do the same for gets.
-	headers := map[string][]string{"cf": []string{"a"}}
-	for i := n - 1; i >= 0; i-- {
-		key := fmt.Sprintf("%s_%d", t.Name(), i)
-		wg.Add(1)
-		go func() {
-			defer wg.Done()
-			get, err := hrpc.NewGetStr(context.Background(), table, key, hrpc.Families(headers))
-			if err != nil {
-				t.Error(key, err)
-				return
+			const n = 1000
+			var wg sync.WaitGroup
+			for i := 0; i < n; i++ {
+				key := fmt.Sprintf("%s_%d", t.Name(), i)
+				wg.Add(1)
+				go func() {
+					if err := insertKeyValue(c, key, "cf", []byte(key)); err != nil {
+						t.Error(key, err)
+					}
+					wg.Done()
+				}()
 			}
-			rsp, err := c.Get(get)
-			if err != nil {
-				t.Error(key, err)
-				return
+			wg.Wait()
+
+			// All puts are complete. Now do the same for gets.
+			headers := map[string][]string{"cf": []string{"a"}}
+			for i := n - 1; i >= 0; i-- {
+				key := fmt.Sprintf("%s_%d", t.Name(), i)
+				wg.Add(1)
+				go func() {
+					defer wg.Done()
+					get, err := hrpc.NewGetStr(context.Background(), table, key,
+						hrpc.Families(headers))
+					if err != nil {
+						t.Error(key, err)
+						return
+					}
+					rsp, err := c.Get(get)
+					if err != nil {
+						t.Error(key, err)
+						return
+					}
+					if len(rsp.Cells) == 0 {
+						t.Error(key, " got zero cells")
+						return
+					}
+					rsp_value := rsp.Cells[0].Value
+					if !bytes.Equal(rsp_value, []byte(key)) {
+						t.Errorf("expected %q, got %q", key, rsp_value)
+					}
+				}()
 			}
-			if len(rsp.Cells) == 0 {
-				t.Error(key, " got zero cells")
-				return
-			}
-			rsp_value := rsp.Cells[0].Value
-			if !bytes.Equal(rsp_value, []byte(key)) {
-				t.Errorf("expected %q, got %q", key, rsp_value)
-			}
-		}()
+			wg.Wait()
+		})
 	}
-	wg.Wait()
 }
 
 func TestTimestampIncreasing(t *testing.T) {

mockrc_test.go

@@ -181,7 +181,7 @@ func newMockRegionClient(addr string, ctype region.ClientType, queueSize int,
 	flushInterval time.Duration, effectiveUser string,
 	readTimeout time.Duration, codec compression.Codec,
 	dialer func(ctx context.Context, network, addr string) (net.Conn, error),
-	log *slog.Logger) hrpc.RegionClient {
+	log *slog.Logger, minWindowSize, maxWindowSize int) hrpc.RegionClient {
 	m.Lock()
 	clients[addr]++
 	m.Unlock()