Skip to content

Commit

Permalink
database/mongodb: revert to old retry behavior (#8863)
Browse files Browse the repository at this point in the history
* database/mongodb: revert to old retry behavior

* add a default case for non-EOF errors
  • Loading branch information
calvn committed Apr 28, 2020
1 parent ee6c4dd commit a8e4e0e
Showing 1 changed file with 34 additions and 33 deletions.
67 changes: 34 additions & 33 deletions plugins/database/mongodb/mongodb.go
Expand Up @@ -5,6 +5,8 @@ import (
"encoding/json"
"errors"
"fmt"
"io"
"strings"
"time"

"github.com/hashicorp/vault/api"
Expand Down Expand Up @@ -95,11 +97,6 @@ func (m *MongoDB) CreateUser(ctx context.Context, statements dbplugin.Statements
return "", "", dbutil.ErrEmptyCreationStatement
}

client, err := m.getConnection(ctx)
if err != nil {
return "", "", err
}

username, err = m.GenerateUsername(usernameConfig)
if err != nil {
return "", "", err
Expand Down Expand Up @@ -132,7 +129,7 @@ func (m *MongoDB) CreateUser(ctx context.Context, statements dbplugin.Statements
Roles: mongoCS.Roles.toStandardRolesArray(),
}

if err := runCommandWithRetry(ctx, client, mongoCS.DB, createUserCmd); err != nil {
if err := m.runCommandWithRetry(ctx, mongoCS.DB, createUserCmd); err != nil {
return "", "", err
}

Expand All @@ -150,11 +147,6 @@ func (m *MongoDB) SetCredentials(ctx context.Context, statements dbplugin.Statem
m.Lock()
defer m.Unlock()

client, err := m.getConnection(ctx)
if err != nil {
return "", "", err
}

username = staticUser.Username
password = staticUser.Password

Expand All @@ -167,7 +159,7 @@ func (m *MongoDB) SetCredentials(ctx context.Context, statements dbplugin.Statem
if err != nil {
return "", "", err
}
if err := runCommandWithRetry(ctx, client, cs.Database, changeUserCmd); err != nil {
if err := m.runCommandWithRetry(ctx, cs.Database, changeUserCmd); err != nil {
return "", "", err
}

Expand All @@ -188,11 +180,6 @@ func (m *MongoDB) RevokeUser(ctx context.Context, statements dbplugin.Statements

statements = dbutil.StatementCompatibilityHelper(statements)

client, err := m.getConnection(ctx)
if err != nil {
return err
}

// If no revocation statements provided, pass in empty JSON
var revocationStatement string
switch len(statements.Revocation) {
Expand All @@ -206,7 +193,7 @@ func (m *MongoDB) RevokeUser(ctx context.Context, statements dbplugin.Statements

// Unmarshal revocation statements into mongodbRoles
var mongoCS mongoDBStatement
err = json.Unmarshal([]byte(revocationStatement), &mongoCS)
err := json.Unmarshal([]byte(revocationStatement), &mongoCS)
if err != nil {
return err
}
Expand All @@ -222,30 +209,44 @@ func (m *MongoDB) RevokeUser(ctx context.Context, statements dbplugin.Statements
WriteConcern: writeconcern.New(writeconcern.WMajority()),
}

return runCommandWithRetry(ctx, client, db, dropUserCmd)
return m.runCommandWithRetry(ctx, db, dropUserCmd)
}

// RotateRootCredentials is not currently supported on MongoDB
func (m *MongoDB) RotateRootCredentials(ctx context.Context, statements []string) (map[string]interface{}, error) {
return nil, errors.New("root credential rotation is not currently implemented in this database secrets engine")
}

// runCommandWithRetry runs a command with retry.
func runCommandWithRetry(ctx context.Context, client *mongo.Client, db string, cmd interface{}) error {
timeout := time.Now().Add(1 * time.Minute)
backoffTime := 3
for {
// Run command
result := client.Database(db).RunCommand(ctx, cmd, nil)
if result.Err() == nil {
break
}
// runCommandWithRetry runs a command and retries once more if there's a failure
// on the first attempt. This should be called with the lock held
func (m *MongoDB) runCommandWithRetry(ctx context.Context, db string, cmd interface{}) error {
// Get the client
client, err := m.getConnection(ctx)
if err != nil {
return err
}

// Run command
result := client.Database(db).RunCommand(ctx, cmd, nil)

if time.Now().After(timeout) {
return result.Err()
// Error check on the first attempt
err = result.Err()
switch {
case err == nil:
return nil
case err == io.EOF, strings.Contains(err.Error(), "EOF"):
// Call getConnection to reset and retry query if we get an EOF error on first attempt.
client, err = m.getConnection(ctx)
if err != nil {
return err
}
result = client.Database(db).RunCommand(ctx, cmd, nil)
if err := result.Err(); err != nil {
return err
}
time.Sleep(time.Duration(backoffTime) * time.Second)
backoffTime += backoffTime
default:
return err
}

return nil
}

0 comments on commit a8e4e0e

Please sign in to comment.