|
23 | 23 | package ants_test |
24 | 24 |
|
25 | 25 | import ( |
| 26 | + "context" |
26 | 27 | "log" |
27 | 28 | "os" |
28 | 29 | "runtime" |
@@ -1351,6 +1352,34 @@ func TestDefaultPoolReleaseTimeout(t *testing.T) { |
1351 | 1352 | require.NoError(t, err) |
1352 | 1353 | } |
1353 | 1354 |
|
| 1355 | +func TestDefaultPoolReleaseContext(t *testing.T) { |
| 1356 | + ants.Reboot() |
| 1357 | + for i := 0; i < 5; i++ { |
| 1358 | + _ = ants.Submit(func() { |
| 1359 | + time.Sleep(time.Second) |
| 1360 | + }) |
| 1361 | + } |
| 1362 | + require.NotZero(t, ants.Running()) |
| 1363 | + err := ants.ReleaseContext(context.Background()) |
| 1364 | + require.NoError(t, err) |
| 1365 | +} |
| 1366 | + |
| 1367 | +func TestReleaseContextWithNil(t *testing.T) { |
| 1368 | + p, err := ants.NewPool(10) |
| 1369 | + require.NoError(t, err) |
| 1370 | + for i := 0; i < 5; i++ { |
| 1371 | + _ = p.Submit(func() { |
| 1372 | + time.Sleep(time.Second) |
| 1373 | + }) |
| 1374 | + } |
| 1375 | + require.NotZero(t, p.Running()) |
| 1376 | + |
| 1377 | + // Passing nil context should release immediately without waiting for workers to exit. |
| 1378 | + err = p.ReleaseContext(nil) //nolint:staticcheck |
| 1379 | + require.NoError(t, err) |
| 1380 | + require.True(t, p.IsClosed()) |
| 1381 | +} |
| 1382 | + |
1354 | 1383 | func TestMultiPool(t *testing.T) { |
1355 | 1384 | _, err := ants.NewMultiPool(-1, 10, 8) |
1356 | 1385 | require.ErrorIs(t, err, ants.ErrInvalidMultiPoolSize) |
@@ -1536,6 +1565,189 @@ func TestMultiPoolWithFuncGeneric(t *testing.T) { |
1536 | 1565 | mp.Tune(10) |
1537 | 1566 | } |
1538 | 1567 |
|
| 1568 | +func TestMultiPoolReleaseContext(t *testing.T) { |
| 1569 | + mp, err := ants.NewMultiPool(10, 5, ants.RoundRobin) |
| 1570 | + require.NoError(t, err) |
| 1571 | + |
| 1572 | + for i := 0; i < 50; i++ { |
| 1573 | + err = mp.Submit(longRunningFunc) |
| 1574 | + require.NoError(t, err) |
| 1575 | + } |
| 1576 | + require.EqualValues(t, 50, mp.Running()) |
| 1577 | + |
| 1578 | + // Signal workers to stop, then release with a background context. |
| 1579 | + atomic.StoreInt32(&stopLongRunningFunc, 1) |
| 1580 | + err = mp.ReleaseContext(context.Background()) |
| 1581 | + require.NoError(t, err) |
| 1582 | + require.Zero(t, mp.Running()) |
| 1583 | + require.True(t, mp.IsClosed()) |
| 1584 | + atomic.StoreInt32(&stopLongRunningFunc, 0) |
| 1585 | + |
| 1586 | + // Calling ReleaseContext on a closed pool should return ErrPoolClosed. |
| 1587 | + require.ErrorIs(t, mp.ReleaseContext(context.Background()), ants.ErrPoolClosed) |
| 1588 | + |
| 1589 | + // Test with LeastTasks strategy. |
| 1590 | + mp, err = ants.NewMultiPool(10, 5, ants.LeastTasks) |
| 1591 | + require.NoError(t, err) |
| 1592 | + for i := 0; i < 50; i++ { |
| 1593 | + err = mp.Submit(longRunningFunc) |
| 1594 | + require.NoError(t, err) |
| 1595 | + } |
| 1596 | + require.EqualValues(t, 50, mp.Running()) |
| 1597 | + |
| 1598 | + atomic.StoreInt32(&stopLongRunningFunc, 1) |
| 1599 | + err = mp.ReleaseContext(context.Background()) |
| 1600 | + require.NoError(t, err) |
| 1601 | + require.Zero(t, mp.Running()) |
| 1602 | + require.True(t, mp.IsClosed()) |
| 1603 | + atomic.StoreInt32(&stopLongRunningFunc, 0) |
| 1604 | + |
| 1605 | + // Test that a cancelled context returns an error. |
| 1606 | + mp, err = ants.NewMultiPool(10, 5, ants.RoundRobin) |
| 1607 | + require.NoError(t, err) |
| 1608 | + for i := 0; i < 50; i++ { |
| 1609 | + err = mp.Submit(longRunningFunc) |
| 1610 | + require.NoError(t, err) |
| 1611 | + } |
| 1612 | + ctx, cancel := context.WithCancel(context.Background()) |
| 1613 | + cancel() // cancel immediately |
| 1614 | + err = mp.ReleaseContext(ctx) |
| 1615 | + require.Error(t, err) |
| 1616 | + atomic.StoreInt32(&stopLongRunningFunc, 1) |
| 1617 | + require.Eventually(t, func() bool { |
| 1618 | + return mp.Running() == 0 |
| 1619 | + }, 3*time.Second, 100*time.Millisecond) |
| 1620 | + atomic.StoreInt32(&stopLongRunningFunc, 0) |
| 1621 | + |
| 1622 | + // Test reboot after ReleaseContext. |
| 1623 | + mp, err = ants.NewMultiPool(10, 5, ants.RoundRobin) |
| 1624 | + require.NoError(t, err) |
| 1625 | + for i := 0; i < 50; i++ { |
| 1626 | + err = mp.Submit(longRunningFunc) |
| 1627 | + require.NoError(t, err) |
| 1628 | + } |
| 1629 | + atomic.StoreInt32(&stopLongRunningFunc, 1) |
| 1630 | + err = mp.ReleaseContext(context.Background()) |
| 1631 | + require.NoError(t, err) |
| 1632 | + atomic.StoreInt32(&stopLongRunningFunc, 0) |
| 1633 | + |
| 1634 | + mp.Reboot() |
| 1635 | + require.False(t, mp.IsClosed()) |
| 1636 | + for i := 0; i < 50; i++ { |
| 1637 | + err = mp.Submit(longRunningFunc) |
| 1638 | + require.NoError(t, err) |
| 1639 | + } |
| 1640 | + require.EqualValues(t, 50, mp.Running()) |
| 1641 | + atomic.StoreInt32(&stopLongRunningFunc, 1) |
| 1642 | + err = mp.ReleaseContext(context.Background()) |
| 1643 | + require.NoError(t, err) |
| 1644 | + atomic.StoreInt32(&stopLongRunningFunc, 0) |
| 1645 | +} |
| 1646 | + |
| 1647 | +func TestMultiPoolWithFuncReleaseContext(t *testing.T) { |
| 1648 | + ch := make(chan struct{}) |
| 1649 | + mp, err := ants.NewMultiPoolWithFunc(10, 5, longRunningPoolFunc, ants.RoundRobin) |
| 1650 | + require.NoError(t, err) |
| 1651 | + |
| 1652 | + for i := 0; i < 50; i++ { |
| 1653 | + err = mp.Invoke(ch) |
| 1654 | + require.NoError(t, err) |
| 1655 | + } |
| 1656 | + require.EqualValues(t, 50, mp.Running()) |
| 1657 | + |
| 1658 | + close(ch) |
| 1659 | + err = mp.ReleaseContext(context.Background()) |
| 1660 | + require.NoError(t, err) |
| 1661 | + require.Zero(t, mp.Running()) |
| 1662 | + require.True(t, mp.IsClosed()) |
| 1663 | + |
| 1664 | + // Calling ReleaseContext on a closed pool should return ErrPoolClosed. |
| 1665 | + require.ErrorIs(t, mp.ReleaseContext(context.Background()), ants.ErrPoolClosed) |
| 1666 | + |
| 1667 | + // Test with LeastTasks strategy. |
| 1668 | + ch = make(chan struct{}) |
| 1669 | + mp, err = ants.NewMultiPoolWithFunc(10, 5, longRunningPoolFunc, ants.LeastTasks) |
| 1670 | + require.NoError(t, err) |
| 1671 | + for i := 0; i < 50; i++ { |
| 1672 | + err = mp.Invoke(ch) |
| 1673 | + require.NoError(t, err) |
| 1674 | + } |
| 1675 | + close(ch) |
| 1676 | + err = mp.ReleaseContext(context.Background()) |
| 1677 | + require.NoError(t, err) |
| 1678 | + require.Zero(t, mp.Running()) |
| 1679 | + require.True(t, mp.IsClosed()) |
| 1680 | + |
| 1681 | + // Test that a cancelled context returns an error. |
| 1682 | + ch = make(chan struct{}) |
| 1683 | + mp, err = ants.NewMultiPoolWithFunc(10, 5, longRunningPoolFunc, ants.RoundRobin) |
| 1684 | + require.NoError(t, err) |
| 1685 | + for i := 0; i < 50; i++ { |
| 1686 | + err = mp.Invoke(ch) |
| 1687 | + require.NoError(t, err) |
| 1688 | + } |
| 1689 | + ctx, cancel := context.WithCancel(context.Background()) |
| 1690 | + cancel() |
| 1691 | + err = mp.ReleaseContext(ctx) |
| 1692 | + require.Error(t, err) |
| 1693 | + close(ch) |
| 1694 | + require.Eventually(t, func() bool { |
| 1695 | + return mp.Running() == 0 |
| 1696 | + }, 3*time.Second, 100*time.Millisecond) |
| 1697 | +} |
| 1698 | + |
| 1699 | +func TestMultiPoolWithFuncGenericReleaseContext(t *testing.T) { |
| 1700 | + ch := make(chan struct{}) |
| 1701 | + mp, err := ants.NewMultiPoolWithFuncGeneric(10, 5, longRunningPoolFuncCh, ants.RoundRobin) |
| 1702 | + require.NoError(t, err) |
| 1703 | + |
| 1704 | + for i := 0; i < 50; i++ { |
| 1705 | + err = mp.Invoke(ch) |
| 1706 | + require.NoError(t, err) |
| 1707 | + } |
| 1708 | + require.EqualValues(t, 50, mp.Running()) |
| 1709 | + |
| 1710 | + close(ch) |
| 1711 | + err = mp.ReleaseContext(context.Background()) |
| 1712 | + require.NoError(t, err) |
| 1713 | + require.Zero(t, mp.Running()) |
| 1714 | + require.True(t, mp.IsClosed()) |
| 1715 | + |
| 1716 | + // Calling ReleaseContext on a closed pool should return ErrPoolClosed. |
| 1717 | + require.ErrorIs(t, mp.ReleaseContext(context.Background()), ants.ErrPoolClosed) |
| 1718 | + |
| 1719 | + // Test with LeastTasks strategy. |
| 1720 | + ch = make(chan struct{}) |
| 1721 | + mp, err = ants.NewMultiPoolWithFuncGeneric(10, 5, longRunningPoolFuncCh, ants.LeastTasks) |
| 1722 | + require.NoError(t, err) |
| 1723 | + for i := 0; i < 50; i++ { |
| 1724 | + err = mp.Invoke(ch) |
| 1725 | + require.NoError(t, err) |
| 1726 | + } |
| 1727 | + close(ch) |
| 1728 | + err = mp.ReleaseContext(context.Background()) |
| 1729 | + require.NoError(t, err) |
| 1730 | + require.Zero(t, mp.Running()) |
| 1731 | + require.True(t, mp.IsClosed()) |
| 1732 | + |
| 1733 | + // Test that a cancelled context returns an error. |
| 1734 | + ch = make(chan struct{}) |
| 1735 | + mp, err = ants.NewMultiPoolWithFuncGeneric(10, 5, longRunningPoolFuncCh, ants.RoundRobin) |
| 1736 | + require.NoError(t, err) |
| 1737 | + for i := 0; i < 50; i++ { |
| 1738 | + err = mp.Invoke(ch) |
| 1739 | + require.NoError(t, err) |
| 1740 | + } |
| 1741 | + ctx, cancel := context.WithCancel(context.Background()) |
| 1742 | + cancel() |
| 1743 | + err = mp.ReleaseContext(ctx) |
| 1744 | + require.Error(t, err) |
| 1745 | + close(ch) |
| 1746 | + require.Eventually(t, func() bool { |
| 1747 | + return mp.Running() == 0 |
| 1748 | + }, 3*time.Second, 100*time.Millisecond) |
| 1749 | +} |
| 1750 | + |
1539 | 1751 | func TestRebootNewPoolCalc(t *testing.T) { |
1540 | 1752 | atomic.StoreInt32(&sum, 0) |
1541 | 1753 | runTimes := 1000 |
|
0 commit comments