[2.0] Improve extension registration logic (#2505)

[2.0] Improve extension registration logic (#2505) #1291