go-transaction-manager

1
// Package gorm is an implementation of trm.Transaction interface by Transaction for *gorm.DB.
2
package gorm
3

4
import (
5
	"context"
6
	"database/sql"
7
	"errors"
8
	"sync"
9

10
	"gorm.io/gorm"
11

12
	"github.com/avito-tech/go-transaction-manager/trm/v2"
13
	"github.com/avito-tech/go-transaction-manager/trm/v2/drivers"
14
)
15

16
// Transaction is trm.Transaction for sqlx.Tx.
17
type Transaction struct {
18
	tx              *gorm.DB
19
	txMutex         sync.Mutex
20
	isClosed        *drivers.IsClosed
21
	isClosedClosure *drivers.IsClosed
22
}
23

24
// NewTransaction creates trm.Transaction for sqlx.Tx.
25
func NewTransaction(
26
	ctx context.Context,
27
	opts *sql.TxOptions,
28
	db *gorm.DB,
29
) (context.Context, *Transaction, error) {
30
	t := &Transaction{
31
		tx:              nil,
32
		txMutex:         sync.Mutex{},
33
		isClosed:        drivers.NewIsClosed(),
34
		isClosedClosure: drivers.NewIsClosed(),
35
	}
36

37
	var err error
38

39
	wg := sync.WaitGroup{}
40
	wg.Add(1)
41

42
	go func() {
43
		db = db.WithContext(ctx)
44
		// Used closure to avoid implementing nested transactions.
45
		err = db.Transaction(func(tx *gorm.DB) error {
46
			t.tx = tx
47

48
			wg.Done()
49

50
			<-t.isClosedClosure.Closed()
51

52
			return t.isClosedClosure.Err()
53
		}, opts)
54

55
		t.txMutex.Lock()
56
		defer t.txMutex.Unlock()
57
		tx := t.tx
58

59
		if tx != nil {
60
			// Return error from transaction rollback
61
			// Error from commit returns from db.Transaction closure
62
			if errors.Is(err, drivers.ErrRollbackTr) &&
63
				tx.Error != nil {
64
				err = t.tx.Error
65
			}
66

67
			t.isClosed.CloseWithCause(err)
68
		} else {
69
			wg.Done()
70
		}
71
	}()
72

73
	wg.Wait()
74

75
	if err != nil {
76
		return ctx, nil, err
77
	}
78

79
	go t.awaitDone(ctx)
80

81
	return ctx, t, nil
82
}
83

84
func (t *Transaction) awaitDone(ctx context.Context) {
85
	if ctx.Done() == nil {
86
		return
87
	}
88

89
	select {
90
	case <-ctx.Done():
91
		// Rollback will be called by context.Err()
92
		t.isClosedClosure.Close()
93
	case <-t.isClosed.Closed():
94
	}
95
}
96

97
// Transaction returns the real transaction sqlx.Tx.
98
// trm.NestedTrFactory returns IsActive as true while trm.Transaction is opened.
99
func (t *Transaction) Transaction() interface{} {
100
	return t.tx
101
}
102

103
// Begin nested transaction by save point.
104
func (t *Transaction) Begin(ctx context.Context, s trm.Settings) (context.Context, trm.Transaction, error) {
105
	t.txMutex.Lock()
106
	defer t.txMutex.Unlock()
107

108
	return NewDefaultFactory(t.tx)(ctx, s)
109
}
110

111
// Commit closes the trm.Transaction.
112
func (t *Transaction) Commit(_ context.Context) error {
113
	select {
114
	case <-t.isClosed.Closed():
115
		t.txMutex.Lock()
116
		defer t.txMutex.Unlock()
117

118
		return t.tx.Commit().Error
119
	default:
120
		t.isClosedClosure.Close()
121

122
		<-t.isClosed.Closed()
123

124
		return t.isClosed.Err()
125
	}
126
}
127

128
// Rollback the trm.Transaction.
129
func (t *Transaction) Rollback(_ context.Context) error {
130
	select {
131
	case <-t.isClosed.Closed():
132
		t.txMutex.Lock()
133
		defer t.txMutex.Unlock()
134

135
		return t.tx.Rollback().Error
136
	default:
137
		t.isClosedClosure.CloseWithCause(drivers.ErrRollbackTr)
138

139
		<-t.isClosed.Closed()
140

141
		err := t.isClosed.Err()
142
		if errors.Is(err, drivers.ErrRollbackTr) {
143
			return nil
144
		}
145

146
		return err
147
	}
148
}
149

150
// IsActive returns true if the transaction started but not committed or rolled back.
151
func (t *Transaction) IsActive() bool {
152
	return t.isClosed.IsActive()
153
}
154

155
// Closed returns a channel that's closed when transaction committed or rolled back.
156
func (t *Transaction) Closed() <-chan struct{} {
157
	return t.isClosed.Closed()
158
}
159