mapreduce_test.go 8.6 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418
  1. package mr
  2. import (
  3. "errors"
  4. "io/ioutil"
  5. "log"
  6. "runtime"
  7. "sync/atomic"
  8. "testing"
  9. "time"
  10. "git.i2edu.net/i2/go-zero/core/stringx"
  11. "git.i2edu.net/i2/go-zero/core/syncx"
  12. "github.com/stretchr/testify/assert"
  13. )
  14. var errDummy = errors.New("dummy")
  15. func init() {
  16. log.SetOutput(ioutil.Discard)
  17. }
  18. func TestFinish(t *testing.T) {
  19. var total uint32
  20. err := Finish(func() error {
  21. atomic.AddUint32(&total, 2)
  22. return nil
  23. }, func() error {
  24. atomic.AddUint32(&total, 3)
  25. return nil
  26. }, func() error {
  27. atomic.AddUint32(&total, 5)
  28. return nil
  29. })
  30. assert.Equal(t, uint32(10), atomic.LoadUint32(&total))
  31. assert.Nil(t, err)
  32. }
  33. func TestFinishNone(t *testing.T) {
  34. assert.Nil(t, Finish())
  35. }
  36. func TestFinishVoidNone(t *testing.T) {
  37. FinishVoid()
  38. }
  39. func TestFinishErr(t *testing.T) {
  40. var total uint32
  41. err := Finish(func() error {
  42. atomic.AddUint32(&total, 2)
  43. return nil
  44. }, func() error {
  45. atomic.AddUint32(&total, 3)
  46. return errDummy
  47. }, func() error {
  48. atomic.AddUint32(&total, 5)
  49. return nil
  50. })
  51. assert.Equal(t, errDummy, err)
  52. }
  53. func TestFinishVoid(t *testing.T) {
  54. var total uint32
  55. FinishVoid(func() {
  56. atomic.AddUint32(&total, 2)
  57. }, func() {
  58. atomic.AddUint32(&total, 3)
  59. }, func() {
  60. atomic.AddUint32(&total, 5)
  61. })
  62. assert.Equal(t, uint32(10), atomic.LoadUint32(&total))
  63. }
  64. func TestMap(t *testing.T) {
  65. tests := []struct {
  66. mapper MapFunc
  67. expect int
  68. }{
  69. {
  70. mapper: func(item interface{}, writer Writer) {
  71. v := item.(int)
  72. writer.Write(v * v)
  73. },
  74. expect: 30,
  75. },
  76. {
  77. mapper: func(item interface{}, writer Writer) {
  78. v := item.(int)
  79. if v%2 == 0 {
  80. return
  81. }
  82. writer.Write(v * v)
  83. },
  84. expect: 10,
  85. },
  86. {
  87. mapper: func(item interface{}, writer Writer) {
  88. v := item.(int)
  89. if v%2 == 0 {
  90. panic(v)
  91. }
  92. writer.Write(v * v)
  93. },
  94. expect: 10,
  95. },
  96. }
  97. for _, test := range tests {
  98. t.Run(stringx.Rand(), func(t *testing.T) {
  99. channel := Map(func(source chan<- interface{}) {
  100. for i := 1; i < 5; i++ {
  101. source <- i
  102. }
  103. }, test.mapper, WithWorkers(-1))
  104. var result int
  105. for v := range channel {
  106. result += v.(int)
  107. }
  108. assert.Equal(t, test.expect, result)
  109. })
  110. }
  111. }
  112. func TestMapReduce(t *testing.T) {
  113. tests := []struct {
  114. mapper MapperFunc
  115. reducer ReducerFunc
  116. expectErr error
  117. expectValue interface{}
  118. }{
  119. {
  120. expectErr: nil,
  121. expectValue: 30,
  122. },
  123. {
  124. mapper: func(item interface{}, writer Writer, cancel func(error)) {
  125. v := item.(int)
  126. if v%3 == 0 {
  127. cancel(errDummy)
  128. }
  129. writer.Write(v * v)
  130. },
  131. expectErr: errDummy,
  132. },
  133. {
  134. mapper: func(item interface{}, writer Writer, cancel func(error)) {
  135. v := item.(int)
  136. if v%3 == 0 {
  137. cancel(nil)
  138. }
  139. writer.Write(v * v)
  140. },
  141. expectErr: ErrCancelWithNil,
  142. expectValue: nil,
  143. },
  144. {
  145. reducer: func(pipe <-chan interface{}, writer Writer, cancel func(error)) {
  146. var result int
  147. for item := range pipe {
  148. result += item.(int)
  149. if result > 10 {
  150. cancel(errDummy)
  151. }
  152. }
  153. writer.Write(result)
  154. },
  155. expectErr: errDummy,
  156. },
  157. }
  158. for _, test := range tests {
  159. t.Run(stringx.Rand(), func(t *testing.T) {
  160. if test.mapper == nil {
  161. test.mapper = func(item interface{}, writer Writer, cancel func(error)) {
  162. v := item.(int)
  163. writer.Write(v * v)
  164. }
  165. }
  166. if test.reducer == nil {
  167. test.reducer = func(pipe <-chan interface{}, writer Writer, cancel func(error)) {
  168. var result int
  169. for item := range pipe {
  170. result += item.(int)
  171. }
  172. writer.Write(result)
  173. }
  174. }
  175. value, err := MapReduce(func(source chan<- interface{}) {
  176. for i := 1; i < 5; i++ {
  177. source <- i
  178. }
  179. }, test.mapper, test.reducer, WithWorkers(runtime.NumCPU()))
  180. assert.Equal(t, test.expectErr, err)
  181. assert.Equal(t, test.expectValue, value)
  182. })
  183. }
  184. }
  185. func TestMapReduceVoid(t *testing.T) {
  186. var value uint32
  187. tests := []struct {
  188. mapper MapperFunc
  189. reducer VoidReducerFunc
  190. expectValue uint32
  191. expectErr error
  192. }{
  193. {
  194. expectValue: 30,
  195. expectErr: nil,
  196. },
  197. {
  198. mapper: func(item interface{}, writer Writer, cancel func(error)) {
  199. v := item.(int)
  200. if v%3 == 0 {
  201. cancel(errDummy)
  202. }
  203. writer.Write(v * v)
  204. },
  205. expectErr: errDummy,
  206. },
  207. {
  208. mapper: func(item interface{}, writer Writer, cancel func(error)) {
  209. v := item.(int)
  210. if v%3 == 0 {
  211. cancel(nil)
  212. }
  213. writer.Write(v * v)
  214. },
  215. expectErr: ErrCancelWithNil,
  216. },
  217. {
  218. reducer: func(pipe <-chan interface{}, cancel func(error)) {
  219. for item := range pipe {
  220. result := atomic.AddUint32(&value, uint32(item.(int)))
  221. if result > 10 {
  222. cancel(errDummy)
  223. }
  224. }
  225. },
  226. expectErr: errDummy,
  227. },
  228. }
  229. for _, test := range tests {
  230. t.Run(stringx.Rand(), func(t *testing.T) {
  231. atomic.StoreUint32(&value, 0)
  232. if test.mapper == nil {
  233. test.mapper = func(item interface{}, writer Writer, cancel func(error)) {
  234. v := item.(int)
  235. writer.Write(v * v)
  236. }
  237. }
  238. if test.reducer == nil {
  239. test.reducer = func(pipe <-chan interface{}, cancel func(error)) {
  240. for item := range pipe {
  241. atomic.AddUint32(&value, uint32(item.(int)))
  242. }
  243. }
  244. }
  245. err := MapReduceVoid(func(source chan<- interface{}) {
  246. for i := 1; i < 5; i++ {
  247. source <- i
  248. }
  249. }, test.mapper, test.reducer)
  250. assert.Equal(t, test.expectErr, err)
  251. if err == nil {
  252. assert.Equal(t, test.expectValue, atomic.LoadUint32(&value))
  253. }
  254. })
  255. }
  256. }
  257. func TestMapReduceVoidWithDelay(t *testing.T) {
  258. var result []int
  259. err := MapReduceVoid(func(source chan<- interface{}) {
  260. source <- 0
  261. source <- 1
  262. }, func(item interface{}, writer Writer, cancel func(error)) {
  263. i := item.(int)
  264. if i == 0 {
  265. time.Sleep(time.Millisecond * 50)
  266. }
  267. writer.Write(i)
  268. }, func(pipe <-chan interface{}, cancel func(error)) {
  269. for item := range pipe {
  270. i := item.(int)
  271. result = append(result, i)
  272. }
  273. })
  274. assert.Nil(t, err)
  275. assert.Equal(t, 2, len(result))
  276. assert.Equal(t, 1, result[0])
  277. assert.Equal(t, 0, result[1])
  278. }
  279. func TestMapVoid(t *testing.T) {
  280. const tasks = 1000
  281. var count uint32
  282. MapVoid(func(source chan<- interface{}) {
  283. for i := 0; i < tasks; i++ {
  284. source <- i
  285. }
  286. }, func(item interface{}) {
  287. atomic.AddUint32(&count, 1)
  288. })
  289. assert.Equal(t, tasks, int(count))
  290. }
  291. func TestMapReducePanic(t *testing.T) {
  292. v, err := MapReduce(func(source chan<- interface{}) {
  293. source <- 0
  294. source <- 1
  295. }, func(item interface{}, writer Writer, cancel func(error)) {
  296. i := item.(int)
  297. writer.Write(i)
  298. }, func(pipe <-chan interface{}, writer Writer, cancel func(error)) {
  299. for range pipe {
  300. panic("panic")
  301. }
  302. })
  303. assert.Nil(t, v)
  304. assert.NotNil(t, err)
  305. assert.Equal(t, "panic", err.Error())
  306. }
  307. func TestMapReduceVoidCancel(t *testing.T) {
  308. var result []int
  309. err := MapReduceVoid(func(source chan<- interface{}) {
  310. source <- 0
  311. source <- 1
  312. }, func(item interface{}, writer Writer, cancel func(error)) {
  313. i := item.(int)
  314. if i == 1 {
  315. cancel(errors.New("anything"))
  316. }
  317. writer.Write(i)
  318. }, func(pipe <-chan interface{}, cancel func(error)) {
  319. for item := range pipe {
  320. i := item.(int)
  321. result = append(result, i)
  322. }
  323. })
  324. assert.NotNil(t, err)
  325. assert.Equal(t, "anything", err.Error())
  326. }
  327. func TestMapReduceVoidCancelWithRemains(t *testing.T) {
  328. var done syncx.AtomicBool
  329. var result []int
  330. err := MapReduceVoid(func(source chan<- interface{}) {
  331. for i := 0; i < defaultWorkers*2; i++ {
  332. source <- i
  333. }
  334. done.Set(true)
  335. }, func(item interface{}, writer Writer, cancel func(error)) {
  336. i := item.(int)
  337. if i == defaultWorkers/2 {
  338. cancel(errors.New("anything"))
  339. }
  340. writer.Write(i)
  341. }, func(pipe <-chan interface{}, cancel func(error)) {
  342. for item := range pipe {
  343. i := item.(int)
  344. result = append(result, i)
  345. }
  346. })
  347. assert.NotNil(t, err)
  348. assert.Equal(t, "anything", err.Error())
  349. assert.True(t, done.True())
  350. }
  351. func TestMapReduceWithoutReducerWrite(t *testing.T) {
  352. uids := []int{1, 2, 3}
  353. res, err := MapReduce(func(source chan<- interface{}) {
  354. for _, uid := range uids {
  355. source <- uid
  356. }
  357. }, func(item interface{}, writer Writer, cancel func(error)) {
  358. writer.Write(item)
  359. }, func(pipe <-chan interface{}, writer Writer, cancel func(error)) {
  360. drain(pipe)
  361. // not calling writer.Write(...), should not panic
  362. })
  363. assert.Equal(t, ErrReduceNoOutput, err)
  364. assert.Nil(t, res)
  365. }
  366. func BenchmarkMapReduce(b *testing.B) {
  367. b.ReportAllocs()
  368. mapper := func(v interface{}, writer Writer, cancel func(error)) {
  369. writer.Write(v.(int64) * v.(int64))
  370. }
  371. reducer := func(input <-chan interface{}, writer Writer, cancel func(error)) {
  372. var result int64
  373. for v := range input {
  374. result += v.(int64)
  375. }
  376. writer.Write(result)
  377. }
  378. for i := 0; i < b.N; i++ {
  379. MapReduce(func(input chan<- interface{}) {
  380. for j := 0; j < 2; j++ {
  381. input <- int64(j)
  382. }
  383. }, mapper, reducer)
  384. }
  385. }