br: optimize communication with tikv stores. pingcap#31722

Signed-off-by: Gaoming Zhang <zhanggaoming028@gmail.com>

br/pkg/restore/client.go

@@ -882,7 +882,7 @@ func (rc *Client) setSpeedLimit(ctx context.Context, rateLimit uint64) error {
 					func() error {
 						err = rc.fileImporter.setDownloadSpeedLimit(ectx, finalStore.GetId(), rateLimit)
 						if err != nil {
-							return err
+							return errors.Trace(err)
 						}
 						return nil
 					})
@@ -1118,17 +1118,17 @@ CommunicateWithAllStores:
 					)
 					cancel()
 					if err != nil {
-						return err
+						return errors.Trace(err)
 					}
 					client := import_sstpb.NewImportSSTClient(connection)
 					_, err = client.SwitchMode(ctx, &import_sstpb.SwitchModeRequest{
 						Mode: mode,
 					})
 					if err != nil {
-						return err
+						return errors.Trace(err)
 					}
 					err = connection.Close()
-					if err != nil { // do not return. why?
+					if err != nil {
 						log.Error("close grpc connection failed in switch mode", zap.Error(err))
 					}
 					return nil

br/pkg/task/restore.go

@@ -426,15 +426,7 @@ func RunRestore(c context.Context, g glue.Glue, cmdName string, cfg *RestoreConf
 	}
 	// Init DB connection sessions
 	err = client.Init(g, mgr.GetStorage())
-	defer func() {
-		// Reset speed-limit. Do we really need to reset the speed-limit?
-		resetErr := client.ResetSpeedLimit(ctx)
-		// How should this error be handled? Even if reset speed-limit fails, restore may still succeed. TODO
-		if resetErr != nil {
-			log.Error("reset speed limit failed", zap.Error(resetErr))
-		}
-		client.Close()
-	}()
+	defer client.Close()
 
 	if err != nil {
 		return errors.Trace(err)
@@ -622,6 +614,15 @@ func RunRestore(c context.Context, g glue.Glue, cmdName string, cfg *RestoreConf
 		finish = dropToBlackhole(ctx, afterRestoreStream, errCh, updateCh)
 	}
 
+	defer func() {
+		// Reset speed-limit. ResetSpeedLimit must be called after client.InitBackupMeta has been called.
+		resetErr := client.ResetSpeedLimit(ctx)
+		// How should this error be handled? Even if reset speed-limit fails, restore may still succeed. TODO
+		if resetErr != nil {
+			log.Error("reset speed limit failed", zap.Error(resetErr))
+		}
+	}()
+
 	select {
 	case err = <-errCh:
 		err = multierr.Append(err, multierr.Combine(restore.Exhaust(errCh)...))